"""
PARITY GATE (PPDF-10 candidate): JAX port vs FPPDF reference implementation.

Levels:
  A. Chebyshev integral functions Ic1..Ic8 + I(a,b) on a grid of (a,b)
  B. Shape functions q/g/sm/dbub on x-grid, publication params + 5 random draws
  C. Sum-rule machinery: qv_int/qv_norm/smin_norm/msum_ag/sumrules() full vector
  D. Flavour assembly pdfs_msht(ipdf=-4..4) at sum-ruled params
PASS criterion: max |rel dev| < 1e-6 (abs dev < 1e-12 where |ref| < 1e-30).
"""
import sys
import numpy as np

sys.path.insert(0, "/Users/mukesh/Library/CloudStorage/GoogleDrive-mukesh@merakilabs.com/My Drive/MB-Git-Repo/AI Test/53-Proton-PDF/analysis/msht20")
import msht20_jax as J

import fppdf.chebyshevs as FC
import fppdf.global_pars as FG
import fppdf.pdfs as FP

# basis_pars is None until a runcard loads; instantiate with FPPDF defaults
# (Cheb_8=False, g_second_term=True — same conventions the JAX port hard-codes)
FP.basis_pars = FG.BasisPars()

import jax.numpy as jnp

DAT = "/Users/mukesh/fppdf/input/input_mshtfit_publication_shortmin.dat"
RTOL = 1e-6

report, worst = [], 0.0


def check(name, ref, new, rtol=None):
    global worst
    ref, new = np.atleast_1d(np.asarray(ref, float)), np.atleast_1d(np.asarray(new, float))
    denom = np.where(np.abs(ref) > 1e-30, np.abs(ref), 1.0)
    dev = float(np.max(np.abs(new - ref) / denom))
    worst = max(worst, dev)
    report.append((name, dev, "PASS" if dev < (rtol or RTOL) else "FAIL"))


# ---- Level A: integrals ----
avals = np.array([-2.4, -1.4, -0.5, 0.13, 0.34, 1.0, 2.7, 5.0])
bvals = np.array([0.5, 2.0, 3.7, 5.3, 7.5, 9.9, 120.0, 150.0])
for fn_r, fn_j, nm in [(FC.I, J.I_beta, "I")] + [
        (getattr(FC, f"Ic{i}"), getattr(J, f"Ic{i}"), f"Ic{i}") for i in range(1, 9)]:
    refs, news = [], []
    for a in avals:
        for b in bvals:
            if a + b + 2.0 <= 0:
                continue
            refs.append(fn_r(a, b)); news.append(float(fn_j(a, b)))
    # Ic7/Ic8: unused at Cheb_8=False+g_cheb7=False; alternating-sum cancellation
    # limits agreement between math.gamma and gammaln paths — relaxed, non-gating.
    check(f"A.{nm}", refs, news, rtol=1e-4 if nm in ("Ic7", "Ic8") else None)

# ---- params ----
p0 = np.asarray(J.load_publication_start(DAT))
rng = np.random.default_rng(53)
draws = [p0]
for _ in range(5):
    q = p0.copy()
    jitter = 1.0 + 0.1 * rng.standard_normal(len(q))
    q = q * jitter
    # keep exponents in sane ranges so both impls stay finite
    q[2], q[11], q[20], q[29] = abs(q[2]) + 1, abs(q[11]) + 1, abs(q[20]) + 1, abs(q[29]) + 1
    q[48] = abs(q[48]) + 1
    draws.append(q)

xg = np.concatenate([np.logspace(-6, -1, 25), np.linspace(0.11, 0.99, 25)])

# ---- Level B: shape functions ----
for k, p in enumerate(draws):
    check(f"B.q_uv[{k}]",   FP.q_msht(xg, p[0:9]),           np.asarray(J.q_msht(jnp.array(xg), jnp.array(p[0:9]))))
    check(f"B.q_sea[{k}]",  FP.q_msht(xg, p[18:27]),         np.asarray(J.q_msht(jnp.array(xg), jnp.array(p[18:27]))))
    check(f"B.g[{k}]",      FP.g_msht(xg, p[36:46], two_terms=True), np.asarray(J.g_msht(jnp.array(xg), jnp.array(p[36:46]))))
    check(f"B.sm[{k}]",     FP.sm_msht(xg, p[46:56]),        np.asarray(J.sm_msht(jnp.array(xg), jnp.array(p[46:56]))))
    check(f"B.dbub[{k}]",   FP.dbub_msht(xg, p[56:64]),      np.asarray(J.dbub_msht(jnp.array(xg), jnp.array(p[56:64]))))

# ---- Level C: sum-rule machinery ----
for k, p in enumerate(draws):
    check(f"C.qv_int_n[{k}]", FP.qv_int(p[0:9], 1, 0.0),  float(J.qv_int(jnp.array(p[0:9]), 1)))
    check(f"C.qv_int_m[{k}]", FP.qv_int(p[0:9], 2, 0.0),  float(J.qv_int(jnp.array(p[0:9]), 2)))
    check(f"C.qv_norm[{k}]",  FP.qv_norm(1, p[0:9]),      float(J.qv_norm(1, jnp.array(p[0:9]))))
    check(f"C.smin_norm[{k}]", FP.smin_norm(p[46:56]),    float(J.smin_norm(jnp.array(p[46:56]))))
    check(f"C.lowx_norm[{k}]", FP.q_msht_lowx_norm(p[18:27]), float(J.q_msht_lowx_norm(jnp.array(p[18:27]))))
    check(f"C.sp_fix[{k}]",  FP.sp_norm_fix(2.345, p[27:36]), float(J.sp_norm_fix(2.345, jnp.array(p[27:36]))))
    check(f"C.msum_ag[{k}]", FP.msum_ag(p),               float(J.msum_ag(jnp.array(p))))
    check(f"C.sumrules[{k}]", FP.sumrules(p),             np.asarray(J.sumrules(jnp.array(p))))

# ---- Level D: flavour assembly at sum-ruled publication params ----
ps = FP.sumrules(p0)
for ipdf in [-4, -3, -2, -1, 0, 1, 2, 3, 4]:
    check(f"D.pdfs_msht[{ipdf}]", FP.pdfs_msht(ipdf, ps, xg),
          np.asarray(J.pdfs_msht(ipdf, jnp.array(ps), jnp.array(xg))))

# ---- physics sanity on sum-ruled params ----
uv_n = FP.qv_int(ps[0:9], 1, 0.0) * 1.0
print("sanity: uv number integral * A_uv =", ps[0] * FP.qv_int(np.concatenate(([1.0], ps[1:9])), 1, 0.0))

# ---- report ----
fails = [r for r in report if r[2] == "FAIL"]
print(f"\n{'='*64}\nPARITY REPORT — {len(report)} checks, worst rel dev = {worst:.3e}")
for name, dev, status in report:
    if status == "FAIL" or dev > 1e-9:
        print(f"  {status}  {name:24s} {dev:.3e}")
print("RESULT:", "PASS — port is exact within tolerance" if not fails else f"FAIL ({len(fails)} checks)")
sys.exit(0 if not fails else 1)