#!/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, time, 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 via JKS_META=chi2 (7 legacy fields still accepted: chi2 is the first,")
    print("  the rest are ignored).  The result is deterministic (certificate-based),")
    print("  so no seed/n_samples/warmup/facets/fibre parameters are needed any more.")

    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_arr = res.get(omega_grid).mean()
grid_tag = omega_grid

print(f"Performing positivity bounds using energy grid [{min(omega_grid_arr)},..,{max(omega_grid_arr)}] ({len(omega_grid_arr)} resolution)")

def process_weight_element(e, i):
    if isinstance(e, int):
        return np.exp(-e*omega_grid_arr), e
    elif isinstance(e, str):
        ww = res.get(e).mean()
        assert ww.shape == omega_grid_arr.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()

# full input covariance (the diagonal-only version mis-states the tension by
# O(1) on correlated data, e.g. lqcd: chi_min 1.09 (diag) vs 1.80 (full))
c_in_cv = np.array(c_in.cov())

if "JKS_META" in os.environ:
    meta = os.environ["JKS_META"]
else:
    meta = "4,8,0,1500,20000,4000,500"

meta = meta.split(",")
if len(meta) == 1:
    chi2s = [int(meta[0])]
    print(f"Meta parameters: chi2 = {chi2s[0]}")
else:
    assert len(meta) == 7, "JKS_META takes 1 field (chi2) or the 7 legacy fields"
    chi2s0, chi2s1, seed, n_facets, n_samples, warmup, conditional = [int(x) for x in meta]
    chi2s = [chi2s0]
    print(f"""Meta parameters (legacy 7-field format; only chi2 is used):
 chi2        = {chi2s[0]}   (legacy chi2_1={chi2s1}, seed={seed}, n_facets={n_facets},
                             n_samples={n_samples}, warmup={warmup}, conditional={conditional} ignored)
""")

chi2 = float(chi2s[0])

# --- tension between the data and positivity ---------------------------------
# exact (NNLS), not the faceted lower bound of the old min_chi
chi_min_in = jks.positive_laplace.min_chi_exact(list_w_in, c_in_mn, c_in_cv, omega_grid_arr,
                                                floor=None)
# A nonzero chi_min is the NORMAL state, not a red flag: a correlator dominated by
# its ground state sits close to the boundary of the moment cone (a single
# exponential is an extreme ray of it), so an ordinary sub-sigma fluctuation of the
# central value lands outside.  What matters is whether chi^2_min is comparable to
# the chi^2 being used: only then is the admissible set a thin cap.
ratio_cap = chi_min_in ** 2 / chi2
print(f"Positivity tension: chi_min = {chi_min_in:.4f}  (chi^2_min = {chi_min_in**2:.4f}, "
      f"{ratio_cap:.1%} of chi2 {chi2})")
if chi_min_in == 0.0:
    print(" the data are reproduced exactly by a positive spectrum on this grid")
elif ratio_cap <= 0.1:
    print(" small next to chi2 -- consistent with ordinary noise on the central value")
else:
    print(f" the data lie {chi_min_in:.3f} sigma outside the positive cone; the")
    print(f" admissible set at chi2={chi2} is a spherical cap.  Raise chi2 (e.g. to")
    print(f" ~{(chi_min_in + 1.0)**2:.0f}) if the quoted bands look narrow, or check")
    print(f" whether the tension is structural: if chi_min jumps when an input")
    print(f" weight is added, the window cannot fit it rather than it being noise.")
    print(f"NOTE: for N~30 resamples, chi2 < ~9 makes ~10% of datasets infeasible.")
print("")

# --- exact admissible range of every output weight (deterministic) -----------
t0 = time.time()
lo, hi, M, info = jks.positive_laplace.positivity_bounds(
    list_w_in, c_in_mn, c_in_cv, list_w_out, omega_grid_arr,
    chi=chi2 ** 0.5, floor=None)
t1 = time.time()
io = len(list_w_out)

bad = [j for j in range(io) if info["status"][j] != "ok"]
if bad:
    for j in bad:
        print(f"WARNING: output {sel_w_out[j]}: dual status {info['status'][j]} (no band)")
    print("  (chi below chi_min or a numerically diverged certificate; raise chi2)")
central = 0.5 * (lo + hi)
half = 0.5 * (hi - lo)
# 1-sigma statistical error of the central value:  sqrt( (M S M^T)_jj )
stat = np.sqrt(np.clip(np.einsum("ja,ab,jb->j", M, c_in_cv, M), 0.0, None))

print(f"Admissible range computed in {t1 - t0:.2f} s (max certificate violation "
      f"{max(v for v in info['violation_rel'] if v == v):.1e} relative)")
print(" out    central        stat(1sd)    half-width   hw/stat")
for j in range(io):
    s = f" {sel_w_out[j]!s:>6}  {central[j]: .5e}  {stat[j]: .3e}  {half[j]: .3e}  "
    s += f"{half[j] / stat[j]:6.1f}" if stat[j] > 0 else "    inf"
    # data-dominated outputs have hw ~ chi * stat (the band is chi sigma wide);
    # much wider means the location is set by the admissible range, not the data
    s += "   [band-dominated]" if (not np.isnan(central[j])) and (stat[j] == 0 or half[j] > 1.5 * chi2 ** 0.5 * stat[j]) else ""
    print(s)
print("")
print(" hw/stat ~ %.1f is the data-dominated baseline (band width = chi * stat);"
      % (chi2 ** 0.5))
print(" [band-dominated] marks outputs whose location is set by the admissible")
print(" range rather than the data; their central value is the band midpoint.")
print("")

# --- store: central = band midpoint; stat blocks = affine propagation through
# the certificate slope M; the band itself as a single '!band' variation whose
# shift is the half-width, so [central - shift, central + shift] = [lo, hi].
def get_c_out(r):
    d_c_in = np.array(get_c_in(r)) - c_in_mn
    return central + M @ d_c_in

jk = res.apply(get_c_out)

desc = ("Admissible-range band half-width (range = central +- shift): positivity "
        "bounds via dual certificate, chi2=%g, omega grid tag '%s' [%g..%g] x %d nodes, "
        "full input covariance, deterministic.  Input weights: %s"
        % (chi2, grid_tag, min(omega_grid_arr), max(omega_grid_arr), len(omega_grid_arr),
           sel_w_in))
jk = jk.addvar("band", central + half, desc)
res.add(tag_out, jk)

res.save(db)

print("Done")
