#!/usr/bin/env python3
#
#    JKS - Measurement database system
#    Copyright (C) 2013-2026  Christoph Lehner (christoph.lehner@ur.de, https://github.com/lehner/jks)
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
import jks, sys, os, numpy as np
if len(sys.argv) != 7:
    print("%s database.jks tag_in list_of_weights_in list_of_weights_out omega_grid tag_out" % sys.argv[0])
    print("")
    print("- list_of_weights_X accepts a list of")
    print("  a) in integer t for which exp(-t*omega) will be used as the spectral weight")
    print("  b) a string tag to a weight")
    print("")
    print("- Note that an error on omega_grid is not propagated, use multiple omega_grid if needed")
    print("")
    print("- Control meta parameters via JKS_META=chi2_0,chi2_1,seed,n_facets,n_samples,warmup,fiber_steps")

    sys.exit(0)

db, tag_in, list_w_in, list_w_out, omega_grid, tag_out = sys.argv[1:]

res = jks.resamples(db)

omega_grid = res.get(omega_grid).mean()

print(f"Performing sampling using energy grid [{min(omega_grid)},..,{max(omega_grid)}] ({len(omega_grid)} resolution)")

def process_weight_element(e, i):
    if isinstance(e, int):
        return np.exp(-e*omega_grid), e
    elif isinstance(e, str):
        ww = res.get(e).mean()
        assert ww.shape == omega_grid.shape, f"{e} not found"
        return ww, i
    else:
        assert False, "weight not found"

def process_weight(w):
    assert isinstance(w, list)
    t = [process_weight_element(e, i) for i, e in enumerate(w)]
    return [x[0] for x in t], [x[1] for x in t]

list_w_in, sel_w_in = process_weight(eval(list_w_in))
list_w_out, sel_w_out = process_weight(eval(list_w_out))

list_w_in = np.asarray(list_w_in, float)
list_w_out = np.asarray(list_w_out, float)

# create a temporary version of c_in projected
def get_c_in(r):
    raw_in = r[tag_in]
    return [ raw_in[s] for s in sel_w_in ]
    
c_in = res.apply(get_c_in)
c_in_mn = c_in.mean()
c_in_cv = c_in.cov()

if "JKS_META" in os.environ:
    meta = os.environ["JKS_META"]
else:
    meta = "1,2,0,1500,20000,4000,40000"

meta = meta.split(",")
assert len(meta) == 7
chi2s0, chi2s1, seed, n_facets, n_samples, warmup, fiber_steps = [int(x) for x in meta]
chi2s = [chi2s0, chi2s1]

print(f"""Meta parameters:
 chi2s       = {chi2s}
 seed        = {seed}
 n_facets    = {n_facets}
 n_samples   = {n_samples}
 warmup      = {warmup}
 fiber_steps = {fiber_steps}
""")
    
mn = {}
cv = {}
for chi2 in chi2s:
    print(f"Sampling for chi^2 = {chi2}")
    mn[chi2], cv[chi2], _ = jks.positive_laplace.sample_admissible(
        list_w_in, c_in_mn, c_in_cv,
        list_w_out, omega_grid,
        chi=chi2**0.5, floor=None,
        n_facets=n_facets,
        n_samples=n_samples,
        warmup=warmup,
        fiber_steps=fiber_steps,
        seed=0, joint=True
    )

# print(np.abs(mn[2] / mn[1] - 1), "noise")

mn = mn[chi2s[0]]

ii = len(list_w_in)
io = len(list_w_out)
assert len(mn) == ii + io
mn_in = mn[0:ii]
mn_out = mn[ii:]

# cv = chi^2 cv_stat + cv_sys
# cv[a] = a cv_stat + cv_sys
# cv[b] = b cv_stat + cv_sys
# (cv[b] - cv[a]) / (b - a) = cv_stat
# (cv[b] a - cv[a] b) / (a - b) = cv_sys

def get_c_out(r, cor):
    d_c_in = np.array(get_c_in(r)) - c_in_mn
    d_c_out = d_c_in @ cor
    return mn_out + d_c_out

cv_stat = (cv[chi2s[0]] - cv[chi2s[1]]) / (chi2s[0] - chi2s[1])
cv_sys = (cv[chi2s[0]] * chi2s[1] - cv[chi2s[1]] * chi2s[0]) / (chi2s[0] - chi2s[1])

def get_cor(cv):
    cv_in_out = cv[0:ii, ii:(ii+io)]
    cv_in = cv[0:ii, 0:ii]
    cv_out = cv[ii:(ii+io), ii:(ii+io)]
    # should rescale with ratio of plain jks error on input and (max - min) / 2 sample error for chi2=1
    err_in = np.array(c_in_cv).diagonal() ** 0.5
    err_in_ref = cv_in.diagonal() ** 0.5 * (3 ** 0.5)  # 3 comes from uniform versus gaussian variance
    err_scale = np.diag(err_in_ref / err_in)
    return err_scale @ np.linalg.inv(cv_in) @ cv_in_out

cor_in_out = get_cor(cv_stat)    
jk = res.apply(lambda r: get_c_out(r, cor_in_out))
jk = jk.addvar(tag_out, [mn_out[i] + cv_sys[ii + i][ii + i]**0.5 for i in range(len(mn_out))], "Systematic error from %s" % str(sys.argv))
res.add(tag_out, jk)

res.save(db)

print("Done")
