#!/usr/bin/env python3
"""verify_P307.py -- Verifier for Addendum 307 (measure / deformed eigenstates).

Recomputes the eigenstructure and every declared candidate from scratch and
asserts the pre-registered conditions of A307.

  S1  Eigenstructure + A299 flat control      - checks 1-3
  S2  Measure shifts move mu; eigenstate-only does not  - checks 4-6
  S3  PARTIAL-SIGNAL verdict conditions hold   - checks 7-8

The defining facts checked:
 (S1) the flat-dx control recovers the A299 base diagonal (1, 1.12, 7.50),
      eigenvalues are ordered, node counts are (0,1,2);
 (S2) the density measure M1 (rho dx) and the x^2/x^3 measures move mu into
      the partial band [1.3,1.7], while the pure eigenstate change M4
      (Sturm-Liouville with rho weight, flat overlap) does NOT move mu toward
      1.5 -- the lift lives in the MEASURE, not the bare eigenstates;
 (S3) the PARTIAL-SIGNAL verdict holds: no candidate yields a valid diagonal
      peak (so no HIT), but at least one non-control candidate sits with mu in
      [1.3,1.7]; the closest non-control by L2 is M2_x1 at L2 ~ 0.33, well
      below the A299 control distance 0.63.
"""
import sys
import math

import numpy as np

try:
    from scipy.linalg import eigh as _sp_eigh

    def gen_eigh(A, Bm):
        return _sp_eigh(A, Bm)
except Exception:
    def gen_eigh(A, Bm):
        L = np.linalg.cholesky(Bm)
        Linv = np.linalg.inv(L)
        Cmat = Linv @ A @ Linv.T
        Cmat = 0.5 * (Cmat + Cmat.T)
        w, y = np.linalg.eigh(Cmat)
        return w, Linv.T @ y


PASS = FAIL = 0


def check(n, desc, cond):
    global PASS, FAIL
    ok = bool(cond)
    PASS += ok
    FAIL += (not ok)
    print(f"  [{'PASS' if ok else 'FAIL'}] {n:>2}. {desc}")


PI = math.pi
OM = 4 * PI**3 + PI**2 + PI
ES = 13.177
C = (2.0, 3.0, 16.0)
TGT = [c / C[0] for c in C]
N = 1500


def rho_p(x):
    return 48 * PI**3 * x**2 + 6 * PI**2 * x + 2 * PI


def rho(x):
    return 16 * PI**3 * x**3 + 3 * PI**2 * x**2 + 2 * PI * x


def Vpot(x):
    return rho_p(x)**2 / (2 * OM**2) + ES * x**2 * (1 - x)**2


def grid():
    dx = 1.0 / N
    return np.linspace(dx, 1 - dx, N - 1), dx


def hmat(x, dx, pot):
    off = -1.0 / dx**2 * np.ones(N - 2)
    return (np.diag(2.0 / dx**2 + pot(x))
            + np.diag(off, 1) + np.diag(off, -1))


def eig_flat(pot):
    x, dx = grid()
    w, v = np.linalg.eigh(hmat(x, dx, pot))
    return x, dx, w[:3], v[:, :3] / math.sqrt(dx)


def eig_weighted(pot, weight):
    x, dx = grid()
    H = hmat(x, dx, pot)
    Wd = weight(x)
    w, v = gen_eigh(H, np.diag(Wd))
    o = np.argsort(w)
    w = w[o]
    v = v[:, o]
    vv = v[:, :3].copy()
    for n in range(3):
        nrm = math.sqrt(float(np.sum(vv[:, n]**2 * Wd) * dx))
        vv[:, n] = vv[:, n] / nrm
    return x, dx, w[:3], vv


def node_count(col):
    s = np.sign(col[np.abs(col) > 1e-9])
    return int(np.sum(s[1:] != s[:-1]))


def densM(psi, x, dx, measure=None):
    if measure is None:
        return np.array(
            [[float(np.sum(psi[:, n]**2 * C[k - 1] * PI**k * x**k) * dx)
              for k in (1, 2, 3)] for n in range(3)])
    M = np.zeros((3, 3))
    for n in range(3):
        norm = float(np.sum(psi[:, n]**2 * measure) * dx)
        for k in (1, 2, 3):
            M[n, k - 1] = float(np.sum(
                psi[:, n]**2 * C[k - 1] * PI**k * x**k * measure) * dx) / norm
    return M


def diag_peaks(M):
    Z = np.zeros_like(M)
    for k in range(3):
        c = M[:, k]
        Z[:, k] = (c - c.mean()) / c.std() if c.std() > 0 else 0.0
    return all(int(np.argmax(Z[:, k])) == k for k in range(3))


def norm_diag(M):
    d = [M[n, n] for n in range(3)]
    return [v / d[0] for v in d]


def l2(r):
    return math.sqrt(sum((r[i] - TGT[i])**2 for i in range(3)))


def evaluate(psi, lam, x, dx, psi_box, lam_box, measure=None):
    B = densM(psi, x, dx, measure) / np.array(lam)[:, None]
    Bb = densM(psi_box, x, dx, measure) / lam_box[:, None]
    r = norm_diag(B)
    valid = diag_peaks(B) and not diag_peaks(Bb)
    within = all(abs(r[i] - TGT[i]) <= 0.20 * TGT[i] for i in range(3))
    return {"ratio": r, "valid_peak": valid, "L2": l2(r),
            "within20": within, "mu": r[1]}


# ---- recompute -----------------------------------------------------------
x, dx, lam_flat, psi_flat = eig_flat(Vpot)
nu = [node_count(psi_flat[:, n]) for n in range(3)]
psi_box = np.column_stack([math.sqrt(2.0) * np.sin((n + 1) * PI * x)
                           for n in range(3)])
lam_box = np.array([(n + 1)**2 * PI**2 for n in range(3)])

res = {}
res["M5"] = evaluate(psi_flat, lam_flat, x, dx, psi_box, lam_box, None)
base_mu = res["M5"]["mu"]

rho_x = rho(x)
_, _, lam_rho, psi_rho = eig_weighted(Vpot, rho)
nu_rho = [node_count(psi_rho[:, n]) for n in range(3)]
res["M1"] = evaluate(psi_rho, lam_rho, x, dx, psi_box, lam_box, rho_x)

for a in (1, 2, 3):
    _, _, la, pa = eig_weighted(Vpot, (lambda aa: (lambda xx: xx**aa))(a))
    res[f"M2_x{a}"] = evaluate(pa, la, x, dx, psi_box, lam_box, x**a)

alpha = 1.0 / OM
_, _, lam_def, psi_def = eig_flat(lambda xx: Vpot(xx) + alpha * rho(xx))
res["M3"] = evaluate(psi_def, lam_def, x, dx, psi_box, lam_box, None)

res["M4"] = evaluate(psi_rho, lam_rho, x, dx, psi_box, lam_box, None)

non_control = {k: v for k, v in res.items() if k != "M5"}
closest = min(non_control, key=lambda k: non_control[k]["L2"])
any_valid = any(v["valid_peak"] for v in res.values())
hit = any(v["valid_peak"] and v["within20"] for v in non_control.values())

partial_band = [k for k, v in non_control.items() if 1.3 <= v["mu"] <= 1.7]
m4_moved = abs(res["M4"]["mu"] - 1.5) < abs(base_mu - 1.5)
nan_found = any(v != v for k in res for v in res[k]["ratio"])

# ---- assertions ----------------------------------------------------------
print("S1  Eigenstructure + A299 flat control")
check(1, "flat eigenvalues ordered (%.2f, %.2f, %.2f)" % tuple(lam_flat),
      lam_flat[0] < lam_flat[1] < lam_flat[2])
check(2, "node counts (0,1,2) flat=%s rho=%s" % (nu, nu_rho),
      nu == [0, 1, 2] and nu_rho == [0, 1, 2])
check(3, "M5 flat control diagonal (1, %.2f, %.2f) within 0.02 of (1,1.12,7.50)"
      % (res["M5"]["ratio"][1], res["M5"]["ratio"][2]),
      abs(res["M5"]["ratio"][1] - 1.12) < 0.02
      and abs(res["M5"]["ratio"][2] - 7.50) < 0.02)

print("S2  Measure shifts move mu; eigenstate-only does not")
check(4, "M1 rho-measure mu=%.2f lands in partial band [1.3,1.7]"
      % res["M1"]["mu"], 1.3 <= res["M1"]["mu"] <= 1.7)
check(5, "M2_x2 mu=%.2f and M2_x3 mu=%.2f both in partial band [1.3,1.7]"
      % (res["M2_x2"]["mu"], res["M2_x3"]["mu"]),
      1.3 <= res["M2_x2"]["mu"] <= 1.7 and 1.3 <= res["M2_x3"]["mu"] <= 1.7)
check(6, "M4 SL-eigenstate-only mu=%.2f did NOT move toward 1.5 (vs base %.2f)"
      % (res["M4"]["mu"], base_mu), not m4_moved)

print("S3  PARTIAL-SIGNAL verdict conditions hold")
check(7, "no candidate yields a valid diagonal peak (no HIT); no NaN",
      (not any_valid) and (not hit) and (not nan_found))
check(8, "closest is %s at L2=%.2f (below A299 control 0.63); >=1 candidate "
      "in partial band %s" % (closest, res[closest]["L2"], partial_band),
      res[closest]["L2"] < 0.63 and len(partial_band) >= 1)

print(f"\n{'='*60}\nRESULT: {PASS} PASS / {FAIL} FAIL")
sys.exit(0 if FAIL == 0 else 1)