#!/usr/bin/env python3
"""verify_P316.py -- Verifier for Addendum 316 (family-to-mode selection, OI-287-1).

Recomputes the pre-declared family-to-mode maps from scratch and asserts the
pre-registered conditions of A316. A315 settled that the magnitude residual to
mu = 1.5 is neither a measure choice (A311: the forced r^3 dr arena volume) nor an
angular coupling-strength (A315: no scalar g crosses 1.5), but a FAMILY-TO-MODE
SELECTION rule: which eigenstate of the forced 4D-radial operator each family maps
to. A316 enumerates five principled maps and asks which reproduces (1,1.5,8) with a
valid peak and whether the corpus-CANONICAL map (node-count S_radial or
Z_3-character S_Z3) is the one.

Maps (eigenstates and overlap both in r^3 dr; gen_eigh numpy Cholesky reduction):
  S_radial  : family n = n-th mode of l=0 operator (node-count; CANONICAL).
  S_lsector : family n = ground mode of l=n-1 sector (A313 C_ang).
  S_lground : family n = ground mode of l=n sector.
  S_Z3      : family n = ground mode of O_{l=0}+gamma cos(2 pi(n-1)/3) (CANONICAL).
  S_diag    : family n = mode maximizing degree-n overlap (self-consistent fit).

Honest finding: NULL. The canonical S_radial gives (1,1.2216,8.2859), mu short of
1.5 with no peak (reproduces A313 C_4Dself, L2=0.399, the closest map). The
canonical S_Z3 is rank-degenerate (the real Re(omega^k)gamma shifts gamma*(1,-1/2,
-1/2) make the k=1,2 operators identical, collapsing families 2 and 3). No map both
hits within 20% and restores a peak. The verifier encodes this honest verdict: it
asserts each map's diagonal/mu/tau/peak/L2, that neither canonical map hits, that
no map hits at all, which map is closest, the S_Z3 degeneracy, and that the NULL
verdict's defining condition (no map both within 20% and peaking) holds.

  S1  canonical S_radial (node-count) + closest-map facts   - checks 1-3
  S2  non-canonical maps + canonical S_Z3 degeneracy        - checks 4-6
  S3  no map hits + NULL verdict honest condition + no NaN   - checks 7-8
"""
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
CANONICAL = {"S_radial", "S_Z3"}


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 radial_matrices(x, dx, V_on_grid, 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 solve_lsector(l, shift, n_modes=4):
    x, dx = grid()
    w_arr = x**3
    cent = l * (l + 2) / x**2 if l > 0 else np.zeros_like(x)
    A, Wm = radial_matrices(x, dx, Vpot(x) + cent + shift, w_arr)
    w, v = gen_eigh(A, Wm)
    o = np.argsort(w)
    w = w[o]
    v = v[:, o]
    vv = v[:, :n_modes].copy()
    for k in range(n_modes):
        nrm = math.sqrt(float(np.sum(vv[:, k]**2 * w_arr) * dx))
        vv[:, k] = vv[:, k] / nrm
    return x, dx, w[:n_modes], vv


def overlap_row(col, x, dx, w_arr, lam):
    norm = float(np.sum(col**2 * w_arr) * dx)
    return [float(np.sum(col**2 * C[k - 1] * PI**k * x**k * w_arr) * dx)
            / norm / lam for k in (1, 2, 3)]


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 l2(ratio):
    return math.sqrt(sum((ratio[i] - TGT[i])**2 for i in range(3)))


def box_states(x):
    pb = np.column_stack([math.sqrt(2.0) * np.sin((n + 1) * PI * x)
                          for n in range(3)])
    lb = np.array([(n + 1)**2 * PI**2 for n in range(3)])
    return pb, lb


def select_states(mapname):
    x, dx = grid()
    w_arr = x**3
    states = np.zeros((N - 1, 3))
    lams = np.zeros(3)
    if mapname == "S_radial":
        _, _, w0, v0 = solve_lsector(0, 0.0, 4)
        for n in range(3):
            states[:, n] = v0[:, n]
            lams[n] = w0[n]
    elif mapname == "S_lsector":
        for n in range(3):
            _, _, wl, vl = solve_lsector(n, 0.0, 2)
            states[:, n] = vl[:, 0]
            lams[n] = wl[0]
    elif mapname == "S_lground":
        for n in range(3):
            _, _, wl, vl = solve_lsector(n + 1, 0.0, 2)
            states[:, n] = vl[:, 0]
            lams[n] = wl[0]
    elif mapname == "S_Z3":
        shifts = [GAMMA * math.cos(2 * PI * k / 3.0) for k in range(3)]
        for n in range(3):
            _, _, wl, vl = solve_lsector(0, shifts[n], 2)
            states[:, n] = vl[:, 0]
            lams[n] = wl[0]
    elif mapname == "S_diag":
        _, _, w0, v0 = solve_lsector(0, 0.0, 6)
        na = v0.shape[1]
        rows = np.array([overlap_row(v0[:, m], x, dx, w_arr, w0[m])
                         for m in range(na)])
        used = set()
        for n in range(3):
            cand = sorted(range(na), key=lambda m: -rows[m, n])
            pick = next(m for m in cand if m not in used)
            used.add(pick)
            states[:, n] = v0[:, pick]
            lams[n] = w0[pick]
    return x, dx, w_arr, states, lams


def eval_map(mapname, pb, lb):
    x, dx, w_arr, states, lams = select_states(mapname)
    M = np.array([overlap_row(states[:, n], x, dx, w_arr, lams[n])
                  for n in range(3)])
    Mbox = np.array([overlap_row(pb[:, n], x, dx, w_arr, lb[n])
                     for n in range(3)])
    d = [M[n, n] for n in range(3)]
    ratio = [v / d[0] for v in d]
    peak = diag_peaks(M)
    box_peak = diag_peaks(Mbox)
    valid_peak = peak and not box_peak
    within20 = all(abs(ratio[i] - TGT[i]) <= 0.20 * TGT[i] for i in range(3))
    # degeneracy: two families collapse to the same radial state -- identical
    # eigenvalue AND identical full overlap row (the single diagonal entry
    # differs by degree, so compare the whole row).
    degen = any(abs(lams[a] - lams[b]) < 1e-6
                and all(abs(M[a, k] - M[b, k]) < 1e-6 for k in range(3))
                for a in range(3) for b in range(a + 1, 3))
    return {
        "mu": ratio[1], "tau": ratio[2], "ratio": ratio,
        "valid_peak": valid_peak, "within20": within20,
        "hit": within20 and valid_peak, "L2": l2(ratio),
        "lams": [float(v) for v in lams], "degen": degen,
    }


# ---- recompute -----------------------------------------------------------
ALL = ["S_radial", "S_lsector", "S_lground", "S_Z3", "S_diag"]
pb, lb = box_states(grid()[0])
R = {m: eval_map(m, pb, lb) for m in ALL}

canon_hits = [m for m in CANONICAL if R[m]["hit"]]
noncanon_hits = [m for m in ALL if m not in CANONICAL and R[m]["hit"]]
any_hit = any(R[m]["hit"] for m in ALL)
closest = min(ALL, key=lambda m: R[m]["L2"])

nan_found = any(v != v for m in ALL
                for v in R[m]["ratio"] + R[m]["lams"] + [R[m]["mu"], R[m]["tau"]])

verdict = ("FORCED" if canon_hits
           else "FIT_FLAGGED" if noncanon_hits
           else "NULL")

# ---- assertions ----------------------------------------------------------
print("S1  canonical S_radial (node-count) + closest-map facts")
check(1, "canonical S_radial reproduces A313 C_4Dself: diagonal "
      "(1, %.4f, %.4f), mu=%.4f ~ 1.222, L2=%.4f (the closest map)"
      % (R["S_radial"]["mu"], R["S_radial"]["tau"], R["S_radial"]["mu"],
         R["S_radial"]["L2"]),
      abs(R["S_radial"]["mu"] - 1.222) < 0.01
      and abs(R["S_radial"]["tau"] - 8.286) < 0.02)
check(2, "canonical S_radial mu=%.4f is SHORT of the target 1.5 and restores NO "
      "valid peak (peak=%s): the node-count map does not hit"
      % (R["S_radial"]["mu"], R["S_radial"]["valid_peak"]),
      R["S_radial"]["mu"] < 1.5 and not R["S_radial"]["valid_peak"]
      and not R["S_radial"]["hit"])
check(3, "the closest map overall is %s with L2=%.4f and diagonal "
      "(1, %.4f, %.4f); it is a CANONICAL map but does not hit"
      % (closest, R[closest]["L2"], R[closest]["mu"], R[closest]["tau"]),
      closest == "S_radial" and abs(R[closest]["L2"] - 0.3991) < 0.01
      and closest in CANONICAL and not R[closest]["hit"])

print("S2  non-canonical maps + canonical S_Z3 degeneracy")
check(4, "non-canonical S_lsector (A313 C_ang) overshoots: mu=%.4f (>1.5), "
      "tau=%.4f, no valid peak, does not hit"
      % (R["S_lsector"]["mu"], R["S_lsector"]["tau"]),
      R["S_lsector"]["mu"] > 1.5 and not R["S_lsector"]["valid_peak"]
      and not R["S_lsector"]["hit"]
      and abs(R["S_lsector"]["mu"] - 2.009) < 0.01)
check(5, "non-canonical S_lground (mu=%.4f) and self-consistent S_diag "
      "(mu=%.4f, = S_radial) neither hit: S_diag selects modes 0,1,2 and "
      "reproduces S_radial, giving NO peak even when selecting for one"
      % (R["S_lground"]["mu"], R["S_diag"]["mu"]),
      not R["S_lground"]["hit"] and not R["S_diag"]["hit"]
      and abs(R["S_diag"]["mu"] - R["S_radial"]["mu"]) < 1e-6
      and not R["S_diag"]["valid_peak"])
check(6, "canonical S_Z3 is RANK-DEGENERATE in the real radial reduction: the "
      "Re(omega^k)gamma shifts gamma*(1,-1/2,-1/2) make the k=1,2 operators "
      "identical, collapsing families 2 and 3 (degenerate=%s); S_Z3 cannot "
      "produce three distinct entries and does not hit (mu=%.4f)"
      % (R["S_Z3"]["degen"], R["S_Z3"]["mu"]),
      R["S_Z3"]["degen"] and not R["S_Z3"]["hit"])

print("S3  no map hits + NULL verdict honest condition + no NaN")
check(7, "VERDICT NULL: no canonical map hits (%s) and no non-canonical map "
      "hits (%s), so the magnitude piece does NOT close; defining condition: no "
      "declared map both lands within 20%% of (1,1.5,8) AND restores a valid peak"
      % (canon_hits, noncanon_hits),
      verdict == "NULL" and not canon_hits and not noncanon_hits
      and not any_hit
      and not any(R[m]["within20"] and R[m]["valid_peak"] for m in ALL))
check(8, "clean: no NaN across all 5 maps; the two corpus-canonical maps "
      "(node-count S_radial, Z_3-character S_Z3) are declared before computing "
      "and NEITHER hits, so the result is a genuine NULL, not a fit "
      "(canonical_hits=%s)" % canon_hits,
      (not nan_found) and not canon_hits and verdict == "NULL"
      and abs(R["S_radial"]["mu"] - 1.222) < 0.01)

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