#!/usr/bin/env python3
"""verify_P313.py -- Verifier for Addendum 313 (full B^4 arena measure).

Recomputes the eigenstructure and every declared candidate from scratch and
asserts the pre-registered conditions of A313. A313 builds on A311: A311 showed
the inner-product measure that carries the A307 lift is the GEOMETRICALLY FORCED
radial volume r^3 dr of the B^4 arena (P18 S2.1 metric ds^2 = dr^2 + r^2
dOmega_3^2), the flat control recovers A299 (1,1.12,7.50), the r^3 measure gives
mu=1.337, and the corpus-faithful 4D-radial operator -d^2/dr^2 - (3/r) d/dr + V
in its r^3 measure gives mu=1.222 / L2=0.399 (closest). A311 left two limits: mu
short of 1.5 and no diagonal peak, with the hypothesis that the residual sits in
the angular (l) and Hopf-fiber (S^1) sectors the radial reduction drops.

A313 tests that hypothesis with the full-arena object of P18 S4.3: O_{l,k} =
-d^2/dr^2 - (3/r) d/dr + l(l+2)/r^2 + V_eff + omega^k gamma, family n carrying
angular momentum l = n and fiber index k = n. Each family is the ground radial
mode of its own l-sector operator, self-adjoint in r^3 dr (angular candidate) or
in the Hopf-lapse measure r^3 |sin(pi r)| (fiber candidate, with the
omega^k gamma family shift gamma cos(2 pi n / 3)).

  S1  Eigenstructure + A299 flat control + A311 pure-radial - checks 1-4
  S2  Full-arena angular/fiber sectors and the motion they supply - checks 5-7
  S3  Verdict (STRENGTHENED-PARTIAL): mu overshoots, no peak, residual not
      closed by naive l=n - check 8

The defining facts checked:
 (S1) flat eigenvalues ordered, node counts (0,1,2); the flat-dx control
      recovers A299 (1, 1.12, 7.50); the FORCED r^3 measure gives mu ~ 1.337
      (A311); the FORCED 4D-radial operator gives mu ~ 1.222 / L2 ~ 0.399 and is
      the closest candidate (A311);
 (S2) the full-arena angular candidate (l = n) moves mu strongly past 1.40
      (it overshoots to ~2.01); the fiber candidate moves mu further (~2.27);
      the ANGULAR sector supplies the larger share of the motion;
 (S3) the pre-registered STRENGTHENED-PARTIAL verdict holds (a full-arena
      candidate exceeds mu = 1.40, the angular sector supplies the motion), no
      candidate restores a valid diagonal peak (no HIT), and no NaN; the residual
      to 1.5 is NOT cleanly closed by the naive l = n assignment, it overshoots.
"""
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
GAMMA = 0.75
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 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 radial_matrices(x, dx, V_on_grid, weight):
    """Symmetric 4D-radial A and mass W for the positive weight (self-adjoint in
    weight dr): A_ij = sum W R' S' + W V R S ; W = diag(weight)."""
    n = N - 1
    w_mid = 0.5 * (weight[:-1] + weight[1:])
    w_mid_left0 = max(0.5 * (0.0 + weight[0]), 0.0)
    A = np.zeros((n, n))
    for i in range(n):
        wl = w_mid[i - 1] if i > 0 else w_mid_left0
        wr = w_mid[i] if i + 1 < n else weight[i]
        A[i, i] = (wl + wr) / dx**2
        if i + 1 < n:
            A[i, i + 1] = -w_mid[i] / dx**2
            A[i + 1, i] = -w_mid[i] / dx**2
    A += np.diag(weight * V_on_grid)
    A = 0.5 * (A + A.T)
    return A, np.diag(weight)


def eig_radial4d(pot):
    """4D-radial -d^2/dr^2 - (3/r) d/dr + V (P18 S2.2), self-adjoint in r^3 dr."""
    x, dx = grid()
    w3 = x**3
    A, Wm = radial_matrices(x, dx, pot(x), w3)
    w, v = gen_eigh(A, Wm)
    o = np.argsort(w)
    w = w[o]
    v = v[:, o]
    vv = v[:, :3].copy()
    for k in range(3):
        nrm = math.sqrt(float(np.sum(vv[:, k]**2 * w3) * dx))
        vv[:, k] = vv[:, k] / nrm
    return x, dx, w[:3], vv


def eig_angular(pot, weight, shift_fn):
    """Full-arena: family n is the ground radial mode of O_{l=n} with the
    centrifugal term l(l+2)/r^2, l = n, plus shift_fn(n), self-adjoint in
    weight dr."""
    x, dx = grid()
    w_arr = weight(x)
    psis = np.zeros((N - 1, 3))
    lams = np.zeros(3)
    for n in range(3):
        l = n
        V_n = pot(x) + l * (l + 2) / x**2 + shift_fn(n)
        A, Wm = radial_matrices(x, dx, V_n, w_arr)
        w, v = gen_eigh(A, Wm)
        o = np.argsort(w)
        g = v[:, o[0]]
        nrm = math.sqrt(float(np.sum(g**2 * w_arr) * dx))
        psis[:, n] = g / nrm
        lams[n] = w[o[0]]
    return x, dx, lams, psis


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], "tau": r[2]}


# ---- 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)
r3 = x**3
lapse = np.abs(np.sin(PI * x))
fiber_w = r3 * lapse

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

_, _, lam_r3, psi_r3 = eig_weighted(Vpot, lambda xx: xx**3)
res["C_r3"] = evaluate(psi_r3, lam_r3, x, dx, psi_box, lam_box, r3)

_, _, lam_4d, psi_4d = eig_radial4d(Vpot)
res["C_4Dself"] = evaluate(psi_4d, lam_4d, x, dx, psi_box, lam_box, r3)

_, _, lam_ang, psi_ang = eig_angular(Vpot, lambda xx: xx**3, lambda n: 0.0)
res["C_ang"] = evaluate(psi_ang, lam_ang, x, dx, psi_box, lam_box, r3)


def fiber_shift(n):
    return GAMMA * math.cos(2 * PI * n / 3.0)


_, _, lam_fib, psi_fib = eig_angular(
    Vpot, lambda xx: xx**3 * np.abs(np.sin(PI * xx)), fiber_shift)
res["C_fiber"] = evaluate(psi_fib, lam_fib, x, dx, psi_box, lam_box, fiber_w)

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

r3_mu = res["C_r3"]["mu"]
self_mu = res["C_4Dself"]["mu"]
ang_mu = res["C_ang"]["mu"]
fib_mu = res["C_fiber"]["mu"]
pr_closest_to_15 = min((r3_mu, self_mu), key=lambda m: abs(m - 1.5))
fa_mu_gt_140 = any(res[k]["mu"] > 1.40 for k in FULL_ARENA)
angular_motion = ang_mu - r3_mu
fiber_motion = fib_mu - ang_mu
sector = "angular" if abs(angular_motion) >= abs(fiber_motion) else "fiber"
strengthened_partial = (fa_mu_gt_140 or fa_valid) and (not hit)

# ---- assertions ----------------------------------------------------------
print("S1  Eigenstructure + A299 flat control + A311 pure-radial")
check(1, "flat eigenvalues ordered (%.2f, %.2f, %.2f), node counts (0,1,2)=%s"
      % (lam_flat[0], lam_flat[1], lam_flat[2], nu),
      lam_flat[0] < lam_flat[1] < lam_flat[2] and nu == [0, 1, 2])
check(2, "C_flat control diagonal (1, %.2f, %.2f) reproduces A299 (1,1.12,7.50)"
      % (res["C_flat"]["ratio"][1], res["C_flat"]["ratio"][2]),
      abs(res["C_flat"]["ratio"][1] - 1.12) < 0.02
      and abs(res["C_flat"]["ratio"][2] - 7.50) < 0.02)
check(3, "FORCED r^3 measure gives mu=%.4f ~ 1.337 (A311)" % r3_mu,
      abs(r3_mu - 1.337) < 0.01)
check(4, "FORCED 4D-radial operator (r^3 measure) mu=%.4f ~ 1.222, L2=%.4f "
      "~ 0.399, and is the closest candidate (A311)"
      % (self_mu, res["C_4Dself"]["L2"]),
      abs(self_mu - 1.222) < 0.01
      and abs(res["C_4Dself"]["L2"] - 0.399) < 0.01
      and closest == "C_4Dself")

print("S2  Full-arena angular/fiber sectors and the motion they supply")
check(5, "full-arena angular sector (l=n) moves mu=%.4f strongly past 1.40 "
      "(it OVERSHOOTS 1.5)" % ang_mu,
      ang_mu > 1.40 and ang_mu > 1.5)
check(6, "full-arena fiber candidate mu=%.4f moves mu further than angular "
      "(%.4f), both well past 1.40" % (fib_mu, ang_mu),
      fib_mu > ang_mu and fib_mu > 1.40)
check(7, "the ANGULAR sector supplies the larger share of the motion: "
      "angular motion %+.4f >= fiber motion %+.4f (sector=%s)"
      % (angular_motion, fiber_motion, sector),
      abs(angular_motion) >= abs(fiber_motion) and sector == "angular")

print("S3  Verdict: STRENGTHENED-PARTIAL (mu overshoots, no peak)")
check(8, "STRENGTHENED-PARTIAL: a full-arena candidate exceeds mu=1.40 (%s), no "
      "valid diagonal peak (no HIT), no NaN; residual not cleanly closed (mu "
      "overshoots 1.5, naive l=n over-applies the motion)"
      % fa_mu_gt_140,
      strengthened_partial and (not hit) and (not any_valid)
      and (not nan_found) and (ang_mu > 1.5) and (fib_mu > 1.5))

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