#!/usr/bin/env python3
"""verify_P309.py -- Verifier for Addendum 309 (density-measure family sweep).

Recomputes the eigenstructure and every declared measure from scratch and
asserts the pre-registered conditions of A309. A309 builds on A307: A307
localized the missing A299 selection-rule factor to the inner-product MEASURE
(W=rho lifts mu from 1.12 to 1.34) but left the exact weight unpinned. A309
sweeps a declared principled density-measure family and tests whether any
member restores a valid diagonal peak, and whether plain rho dx sits at a
distinguished / forced point.

  S1  Eigenstructure + A299 flat control          - checks 1-3
  S2  Measure family: W=rho lift, no peak, no HIT - checks 4-6
  S3  NULL-ISH verdict conditions                 - checks 7-8

The defining facts checked:
 (S1) the flat-dx control recovers the A299 base diagonal (1, 1.12, 7.50),
      eigenvalues ordered, node counts (0,1,2);
 (S2) the density measure W=rho lifts mu to ~1.34 reproducing A307; NO declared
      measure yields a valid diagonal peak; the HIT criterion is met by none;
 (S3) the NULL-ISH verdict holds: the closest non-control member is the
      NON-canonical power W=x^1 at L2~0.3252 (identical to A307, not improved),
      so rho dx is NOT forced by L2; and the rho-power sub-family mu is NOT
      monotone in p (it peaks at p=1, the local mu-maximum), so the
      monotone-crossing distinguished-point condition also fails.
"""
import sys
import math

import numpy as np


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):
    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):
    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)])
ones = np.ones_like(x)
rho_x = rho(x)

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

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

for p in (0.5, 1.0, 2.0):
    _, _, lp, pp = eig_weighted(Vpot, (lambda q: (lambda xx: rho(xx)**q))(p))
    key = "W_rho" if p == 1.0 else f"W_rho_pow{p}"
    res[key] = evaluate(pp, lp, x, dx, psi_box, lam_box, rho_x**p)

for m in (1, 2):
    _, _, lm, pm = eig_weighted(Vpot, (lambda q: (lambda xx: xx**q * rho(xx)))(m))
    key = "W_xrho_m1" if m == 1 else "W_x2rho_m2"
    res[key] = evaluate(pm, lm, x, dx, psi_box, lam_box, x**m * rho_x)

CANON = {"W_rho", "W_rho_pow0.5", "W_rho_pow2.0", "W_xrho_m1", "W_x2rho_m2"}
non_control = {k: v for k, v in res.items() if k != "W_flat"}
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())
nan_found = any(v != v for k in res for v in res[k]["ratio"])

rho_pow_mus = [res["W_rho_pow0.5"]["mu"], res["W_rho"]["mu"],
               res["W_rho_pow2.0"]["mu"]]
rho_pow_monotone = (rho_pow_mus[0] <= rho_pow_mus[1] <= rho_pow_mus[2]) \
    or (rho_pow_mus[0] >= rho_pow_mus[1] >= rho_pow_mus[2])
a307_closest_l2 = 0.3252

# ---- 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" % nu, nu == [0, 1, 2])
check(3, "W_flat control diagonal (1, %.2f, %.2f) within 0.02 of (1,1.12,7.50)"
      % (res["W_flat"]["ratio"][1], res["W_flat"]["ratio"][2]),
      abs(res["W_flat"]["ratio"][1] - 1.12) < 0.02
      and abs(res["W_flat"]["ratio"][2] - 7.50) < 0.02)

print("S2  Measure family: W=rho lift, no peak, no HIT")
check(4, "W=rho measure mu=%.2f ~ 1.34 (reproduces A307 lead)"
      % res["W_rho"]["mu"], abs(res["W_rho"]["mu"] - 1.34) < 0.02)
check(5, "no declared measure yields a valid diagonal peak; no NaN",
      (not any_valid) and (not nan_found))
check(6, "HIT criterion (valid peak AND within 20%) met by NONE",
      not hit)

print("S3  NULL-ISH verdict conditions")
check(7, "closest is %s (NON-canonical) at L2=%.4f, NOT improving A307 0.3252; "
      "rho dx not forced by L2"
      % (closest, res[closest]["L2"]),
      closest == "W_pow_x1" and (closest not in CANON)
      and abs(res[closest]["L2"] - a307_closest_l2) < 0.01
      and res[closest]["L2"] >= a307_closest_l2 - 1e-6)
check(8, "rho-power mu NOT monotone in p (%.2f,%.2f,%.2f): peaks at p=1, so no "
      "monotone-crossing distinguished point"
      % tuple(rho_pow_mus),
      (not rho_pow_monotone)
      and rho_pow_mus[1] > rho_pow_mus[0] and rho_pow_mus[1] > rho_pow_mus[2])

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