#!/usr/bin/env python3
"""verify_P311.py -- Verifier for Addendum 311 (geometric measure forced?).

Recomputes the eigenstructure and every declared candidate from scratch and
asserts the pre-registered conditions of A311. A311 builds on A307/A309: A307
localized the missing A299 selection-rule factor to the inner-product MEASURE
(W=rho lifts mu 1.12 -> 1.34); A309 showed the exact weight is not pinned by fit
(x^1 dx closest at L2=0.3252, rho dx un-forced by fit). A311 asks whether a
particular measure is FORCED by the B^4/S^3/Hopf radial geometry rather than
searched: the 4D radial volume element r^3 dr of the arena (P18 S2.1 metric
ds^2 = dr^2 + r^2 dOmega_3^2; P18 S2.2 radial operator -d^2/dr^2 - (3/r) d/dr).

  S1  Eigenstructure + A299 flat control            - checks 1-3
  S2  Forced arena measures lift mu; analytic r^3~rho - checks 4-6
  S3  Verdict + comparison + peak                   - 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 FORCED 4D arena volume W=r^3 lifts mu into the partial band [1.3,1.7]
      (mu ~ 1.34), the corpus-faithful 4D-radial operator in its r^3 measure
      moves mu toward 1.5 (mu ~ 1.22) and gives the closest fit, and the
      analytic fact rho(r) ~ 16 pi^3 r^3 holds (the r^3 term carries the bulk of
      rho on [0,1], so r^3 and rho share the leading behaviour);
 (S3) the comparison density measure W=rho reproduces A307's mu ~ 1.34, no
      candidate restores a valid diagonal peak (no HIT), and the pre-registered
      STRENGTHENED-PARTIAL verdict holds (a forced arena measure lifts mu into
      the band, geometrically forced not fitted).
"""
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 eig_radial4d(pot):
    """4D-radial operator -d^2/dr^2 - (3/r) d/dr + V (P18 S2.2), self-adjoint in
    r^3 dr. Weak form: A_ij = sum r^3 R' S' + r^3 V R S ; W = diag(r^3)."""
    x, dx = grid()
    w3 = x**3
    V = pot(x)
    n = N - 1
    r_mid = x + 0.5 * dx
    w3_mid = r_mid**3
    w3_mid_left0 = max(x[0] - 0.5 * dx, 0.0)**3
    A = np.zeros((n, n))
    for i in range(n):
        wl = w3_mid[i - 1] if i > 0 else w3_mid_left0
        wr = w3_mid[i]
        A[i, i] = (wl + wr) / dx**2
        if i + 1 < n:
            A[i, i + 1] = -wr / dx**2
            A[i + 1, i] = -wr / dx**2
    A += np.diag(w3 * V)
    A = 0.5 * (A + A.T)
    w, v = gen_eigh(A, np.diag(w3))
    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 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
rho_x = rho(x)
hopf = x**3 * np.sqrt(1.0 - x**2)

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

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

_, _, lam_r3w, psi_r3w = eig_radial4d(Vpot)
res["G_r3w"] = evaluate(psi_r3w, lam_r3w, x, dx, psi_box, lam_box, r3)

_, _, lam_rho, psi_rho = eig_weighted(Vpot, rho)
res["G_rho"] = evaluate(psi_rho, lam_rho, x, dx, psi_box, lam_box, rho_x)

_, _, lam_h, psi_h = eig_weighted(Vpot, lambda xx: xx**3 * np.sqrt(1.0 - xx**2))
res["G_Hopf"] = evaluate(psi_h, lam_h, x, dx, psi_box, lam_box, hopf)

FORCED = {"G_r3", "G_r3w"}
non_control = {k: v for k, v in res.items() if k != "G_flat"}
closest = min(non_control, key=lambda k: non_control[k]["L2"])
any_valid = any(v["valid_peak"] for v in res.values())
hit_forced = any(res[k]["valid_peak"] and res[k]["within20"] for k in FORCED)
nan_found = any(v != v for k in res for v in res[k]["ratio"])

forced_in_band = [k for k in FORCED if 1.3 <= res[k]["mu"] <= 1.7]
forced_moved = any(
    abs(res[k]["mu"] - 1.5) < abs(base_mu - 1.5) for k in FORCED)

# analytic: rho(r) leading term 16 pi^3 r^3; fraction of rho(1) it carries
r3_fraction_at_1 = (16 * PI**3) / rho(1.0)

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

print("S2  Forced arena measures lift mu; analytic r^3 ~ rho")
check(4, "FORCED 4D arena volume W=r^3 lifts mu=%.4f into band [1.3,1.7]"
      % res["G_r3"]["mu"], 1.3 <= res["G_r3"]["mu"] <= 1.7)
check(5, "FORCED 4D-radial operator (r^3 measure) mu=%.4f moved toward 1.5 vs "
      "base %.4f, and is closest by L2=%.4f"
      % (res["G_r3w"]["mu"], base_mu, res["G_r3w"]["L2"]),
      abs(res["G_r3w"]["mu"] - 1.5) < abs(base_mu - 1.5)
      and closest == "G_r3w")
check(6, "analytic: rho leading term 16 pi^3 r^3; r^3 carries %.3f of rho(1) "
      "(rho and r^3 share leading r^3 behaviour)" % r3_fraction_at_1,
      r3_fraction_at_1 > 0.9
      and abs((16 * PI**3) - 496.1) < 1.0)

print("S3  Verdict + comparison + peak")
check(7, "comparison W=rho reproduces A307 mu=%.4f ~ 1.34, and G_r3 mu=%.4f "
      "matches it within 0.01 (the leading-r^3 coincidence)"
      % (res["G_rho"]["mu"], res["G_r3"]["mu"]),
      abs(res["G_rho"]["mu"] - 1.34) < 0.02
      and abs(res["G_r3"]["mu"] - res["G_rho"]["mu"]) < 0.01)
check(8, "STRENGTHENED-PARTIAL: a FORCED arena measure lifts mu into band %s, "
      "no valid peak (no HIT), no NaN"
      % forced_in_band,
      len(forced_in_band) >= 1 and (not hit_forced)
      and (not any_valid) and (not nan_found) and forced_moved)

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