#!/usr/bin/env python3
# Copyright Léon Fernando Vlegels -- CC BY 4.0
"""probe_selection_modemap.py -- Probe for OI-287-1 (A316).

FIT-VS-FORCED probe. A315 settled that the magnitude residual of the charged-
lepton selection rule is NOT a measure choice (A311: the inner-product measure is
the geometrically forced B^4 radial volume r^3 dr) and NOT an angular
coupling-strength (A315: no scalar centrifugal coupling g lands mu = 1.5; the
target lies in the gap between two disjoint mode-selection ranges). What A315 left
open is the FAMILY-TO-MODE SELECTION RULE: which eigenstate of the forced
4D-radial operator each of the three families is mapped to. A316 enumerates a
SHORT, PRE-DECLARED list of principled family-to-mode maps and asks, for each,
whether it reproduces the target diagonal (1, 1.5, 8) with a valid column-
standardized peak, and whether that map is FORCED by the corpus or a FIT.

THE FORCED GEOMETRY (held fixed across every map, verbatim from A311/A313):
  4D-radial operator  O_l = -d^2/dr^2 - (3/r) d/dr + l(l+2)/r^2 + V_eff(r),
  self-adjoint in r^3 dr; eigenstates AND the degree-k density overlap both formed
  in r^3 dr (the A311/A313 C_4Dself / C_ang object). V_eff(r) =
  rho'(r)^2/(2 Omega^2) + E_self r^2 (1-r)^2, rho'(r) = 48 pi^3 r^2 + 6 pi^2 r +
  2 pi, Omega = 4 pi^3 + pi^2 + pi, E_self = 13.177, rho(r) = 16 pi^3 r^3 +
  3 pi^2 r^2 + 2 pi r, C = (2,3,16) -> target normalized (1, 1.5, 8). The
  eigensolver (gen_eigh, numpy Cholesky reduction; scipy unavailable), the
  4D-radial assembly (radial4d_matrices), and the box control sqrt(2)sin((n+1)pi r)
  in the same r^3 measure are all reused verbatim from A313/A315.

A "family-to-mode map" assigns to each family n in {1,2,3} a specific eigenstate
of the operator (possibly in a specific l-sector), and the diagonal entry n is the
degree-n density overlap of that eigenstate, divided by its eigenvalue (the
A299/A311 normalization). The PRE-DECLARED maps (declared here before any number is
computed; NO post-hoc additions):

  (S-radial)  family n -> the n-th radial mode (modes 0,1,2) of the l=0 operator.
              NODE-COUNT MAP: node count = n-1. The A297-A301 thread used this; it
              is the corpus-canonical "node count distinguishes families" reading.
              This is family A of A315 at g=0 -> exactly A313 C_4Dself, mu=1.222.
              [CORPUS-CANONICAL family definition #1]

  (S-lsector) family n -> the ground mode of the l = n-1 sector (one family per
              angular sector). The A313 C_ang construction at l = 0,1,2.
              [principled but non-canonical: ties family to angular sector]

  (S-lground) family n -> the ground mode of the l = n sector (l = 1,2,3).
              A shifted-sector variant; principled, non-canonical.

  (S-Z3)      family n -> the eigenstate in the Z_3 character-k eigenspace,
              k = n-1, of the layer-cycle operator T_cycle. P20 S5.1: the three
              fermion families ARE the Z_3 eigenspaces of T_cycle, T_cycle psi^(k)
              = omega^k psi^(k), omega = exp(2 pi i / 3), k = 0,1,2 (P20 Prop.
              "Family Origin"). [CORPUS-CANONICAL family definition #2]
              CONCRETE REALIZATION IN THE RADIAL REDUCTION: T_cycle is not present
              as a distinct operator in the pure radial reduction. The corpus
              footprint of T_cycle on the radial operator (P18 S4.3 O_{l,k}) is the
              real diagonal layer-cycle term Re(omega^k) gamma = gamma cos(2 pi k/3)
              added to the SAME l=0 base operator, with gamma = 3/4 the P18 S4.2
              layer-cycle coefficient (this is exactly the A313 C_fiber family
              shift). So S-Z3 here is: family n = ground mode of
              O_{l=0} + gamma cos(2 pi (n-1)/3) on r^3 dr. Because the three shifts
              for k=0,1,2 are gamma*(1, -1/2, -1/2), the k=1 and k=2 operators are
              IDENTICAL, so their ground modes coincide; the Z_3 phase is a complex
              label the real radial reduction cannot separate. We state this
              honestly: in this reduction S-Z3 has only two distinct radial
              operators (k=0 and k=1,2 degenerate) and therefore CANNOT by itself
              produce three distinct diagonal entries. We test it as constructed
              and report this structural degeneracy as the finding for S-Z3. (An
              alternative honest reading -- S-Z3 reduces to S-radial because the
              Z_3 character has no radial footprint -- is reported alongside.)

  (S-diag)    family n -> the eigenstate (among the lowest several modes of the
              l=0 operator) whose degree-n density overlap is MAXIMAL.
              SELF-CONSISTENT diagonal selection. [principled but self-referential:
              it selects to make the diagonal peak by construction, so a peak here
              is NOT evidence of a forced rule -- explicitly a fit]

PRE-REGISTERED VERDICT (declared before running):
  HIT (would CLOSE the magnitude piece of OI-287-1) iff the CORPUS-CANONICAL map
    -- S-radial (node-count) OR S-Z3 (character map), the two principled family
    DEFINITIONS in the corpus -- reproduces (1, 1.5, 8) to within 20% on all three
    entries AND restores a valid column-standardized diagonal peak (peak holds for
    the map AND the box control fails the same test in the same measure). ONLY a
    canonical map counts for a HIT.
  FIT-FLAGGED iff some NON-canonical map (S-lsector, S-lground, S-diag) hits the
    target with a peak but the canonical maps do NOT: reported as a FITTED
    selection, explicitly not a derivation, with which canonical map is closest and
    how far.
  NULL iff no declared map reproduces the target / restores a peak: report the
    closest map + its diagonal, and the structural feature that still blocks the
    peak.

Honesty discipline: the canonical maps (node-count, Z_3-character) are privileged
BY THE CORPUS, declared before computing. If one hits, it is a genuine derivation.
If only an ad-hoc map hits, it is a fit and is named as such. tau and peak are
reported, not just mu. Five maps declared (multiple-comparison count = 5). No NaN;
r^3 > 0 on (0,1); every operator symmetric-definite.
"""
import os
import json
import math

import numpy as np


def gen_eigh(A, Bm):
    """Generalized symmetric-definite eig A v = lambda W v, W SPD, via numpy
    Cholesky reduction (scipy unavailable). Verbatim from A313/A315.
        W = L L^T ; C = L^{-1} A L^{-T} ; C y = lambda y ; v = L^{-T} y.
    """
    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)
    v = Linv.T @ y
    return w, v


PI = math.pi
OMEGA = 4 * PI**3 + PI**2 + PI            # alpha^{-1}
E_SELF = 13.177
GAMMA = 0.75                             # P18 S4.2 layer-cycle coefficient 3/4
C = (2.0, 3.0, 16.0)                     # raw target coefficients
TARGET = [c / C[0] for c in C]           # normalized: (1, 1.5, 8)
N = 1500

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "..", "..", "..", "..", "Outputs",
                   "selection_rule_modemap.json")

# CORPUS-CANONICAL family definitions (the two privileged maps, declared here).
CANONICAL_MAPS = {"S_radial", "S_Z3"}
ALL_MAPS = ["S_radial", "S_lsector", "S_lground", "S_Z3", "S_diag"]


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 * OMEGA**2) + E_SELF * x**2 * (1 - x)**2


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


def radial4d_matrices(x, dx, V_on_grid, weight):
    """Symmetric 4D-radial operator A and mass W, self-adjoint in weight dr.
    Verbatim from A313/A315 (probe_selection_fullarena.radial4d_matrices)."""
    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 0.5 * (weight[i] + 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):
    """Solve O_l = -d^2/dr^2 - (3/r)d/dr + l(l+2)/r^2 + V_eff + shift in r^3 dr.
    Returns x, dx, eigenvalues, r^3-normalized eigenvectors (lowest n_modes)."""
    x, dx = grid()
    w_arr = x**3
    cent = l * (l + 2) / x**2 if l > 0 else np.zeros_like(x)
    A, Wm = radial4d_matrices(x, dx, Vpot(x) + cent + shift, w_arr)
    w, v = gen_eigh(A, Wm)
    order = np.argsort(w)
    w = w[order]
    v = v[:, order]
    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):
    """Degree-k normalized density overlap row for one mode, divided by lam
    (A299/A311 normalization). measure = r^3 (w_arr)."""
    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 node_count(col):
    s = np.sign(col[np.abs(col) > 1e-9])
    return int(np.sum(s[1:] != s[:-1]))


def diag_peaks(M):
    """Column-standardized argmax of column k == k for all k."""
    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(ratio):
    return math.sqrt(sum((ratio[i] - TARGET[i])**2 for i in range(3)))


def box_states(x):
    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)])
    return psi_box, lam_box


def build_matrix(states, lams, x, dx, w_arr):
    """Assemble the 3x3 B-matrix from three selected (state, eigenvalue) pairs:
    B[n][k] = overlap_k(state_n) / lam_n, all in r^3 dr."""
    M = np.zeros((3, 3))
    for n in range(3):
        M[n, :] = overlap_row(states[:, n], x, dx, w_arr, lams[n])
    return M


def box_matrix(psi_box, lam_box, x, dx, w_arr):
    M = np.zeros((3, 3))
    for n in range(3):
        M[n, :] = overlap_row(psi_box[:, n], x, dx, w_arr, lam_box[n])
    return M


def select_states(mapname):
    """Build the three selected eigenstates (states[:, n], lams[n], nodes[n]) for
    one family-to-mode map, all r^3-normalized. Returns also any structural note."""
    x, dx = grid()
    w_arr = x**3
    states = np.zeros((N - 1, 3))
    lams = np.zeros(3)
    nodes = [0, 0, 0]
    note = ""

    if mapname == "S_radial":
        # family n -> n-th radial mode (modes 0,1,2) of the l=0 operator
        _, _, w0, v0 = solve_lsector(0, 0.0, n_modes=4)
        for n in range(3):
            states[:, n] = v0[:, n]
            lams[n] = w0[n]
            nodes[n] = node_count(v0[:, n])
        note = "node-count map: family n = mode n-1 of l=0 operator"

    elif mapname == "S_lsector":
        # family n -> ground mode of l = n-1 sector
        for n in range(3):
            _, _, wl, vl = solve_lsector(n, 0.0, n_modes=2)
            states[:, n] = vl[:, 0]
            lams[n] = wl[0]
            nodes[n] = node_count(vl[:, 0])
        note = "one family per angular sector: family n = ground of l=n-1"

    elif mapname == "S_lground":
        # family n -> ground mode of l = n sector (l = 1,2,3)
        for n in range(3):
            _, _, wl, vl = solve_lsector(n + 1, 0.0, n_modes=2)
            states[:, n] = vl[:, 0]
            lams[n] = wl[0]
            nodes[n] = node_count(vl[:, 0])
        note = "shifted sector: family n = ground of l=n"

    elif mapname == "S_Z3":
        # family n -> ground mode of O_{l=0} + gamma cos(2 pi (n-1)/3)
        # (the Re(omega^k) gamma layer-cycle footprint, P18 S4.3 / A313 C_fiber).
        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], n_modes=2)
            states[:, n] = vl[:, 0]
            lams[n] = wl[0]
            nodes[n] = node_count(vl[:, 0])
        # honest structural note: k=1 and k=2 shifts are equal (-gamma/2), so the
        # two operators are identical and their ground modes coincide.
        note = ("Z_3 character map via the real Re(omega^k)gamma layer-cycle "
                "shift (P18 S4.3 O_{l,k}); shifts gamma*(1,-1/2,-1/2): the k=1 "
                "and k=2 operators are IDENTICAL, so families 2 and 3 collapse "
                "to the same radial state -- the real radial reduction cannot "
                "separate the Z_3 phase. S-Z3 has only two distinct operators.")

    elif mapname == "S_diag":
        # family n -> the lowest-several l=0 mode whose degree-n overlap is maximal
        _, _, w0, v0 = solve_lsector(0, 0.0, n_modes=6)
        n_avail = v0.shape[1]
        rows = np.zeros((n_avail, 3))
        for m in range(n_avail):
            rows[m, :] = overlap_row(v0[:, m], x, dx, w_arr, w0[m])
        used = set()
        for n in range(3):
            # pick the mode (not already used) maximizing degree-n overlap
            cand = sorted(range(n_avail),
                          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]
            nodes[n] = node_count(v0[:, pick])
        note = ("self-consistent diagonal selection: family n = the mode "
                "maximizing degree-n overlap (selects FOR a peak; a peak here "
                "is by construction, not evidence of a forced rule)")

    return x, dx, w_arr, states, lams, nodes, note


def evaluate_map(mapname, psi_box, lam_box):
    x, dx, w_arr, states, lams, nodes, note = select_states(mapname)
    B = build_matrix(states, lams, x, dx, w_arr)
    B_box = box_matrix(psi_box, lam_box, x, dx, w_arr)
    ratio = norm_diag(B)
    peak = bool(diag_peaks(B))
    box_peak = bool(diag_peaks(B_box))
    within20 = all(abs(ratio[i] - TARGET[i]) <= 0.20 * TARGET[i]
                   for i in range(3))
    # degeneracy guard: two families collapse to the SAME radial state. The
    # signature is identical eigenvalues AND identical state shape (the full
    # density-overlap rows agree, not the single diagonal entry, which differs
    # by degree). For S-Z3 the k=1,2 operators are identical, so families 2,3
    # share an eigenvalue and an entire overlap row.
    degenerate = False
    for a in range(3):
        for b in range(a + 1, 3):
            same_eig = abs(lams[a] - lams[b]) < 1e-6
            same_row = all(abs(B[a, k] - B[b, k]) < 1e-6 for k in range(3))
            if same_eig and same_row:
                degenerate = True
    return {
        "map": mapname,
        "canonical": mapname in CANONICAL_MAPS,
        "note": note,
        "diagonal_ratio": [round(float(v), 4) for v in ratio],
        "mu": round(float(ratio[1]), 4),
        "tau": round(float(ratio[2]), 4),
        "peaks_diagonal": peak,
        "box_control_peaks": box_peak,
        "valid_peak": bool(peak and not box_peak),
        "within_20pct": bool(within20),
        "hits_target": bool(within20 and peak and not box_peak),
        "L2_to_target": round(float(l2(ratio)), 4),
        "eigenvalues": [round(float(v), 4) for v in lams],
        "node_counts": nodes,
        "degenerate_selection": degenerate,
    }


def main():
    x, dx = grid()
    psi_box, lam_box = box_states(x)

    maps = {}
    for name in ALL_MAPS:
        maps[name] = evaluate_map(name, psi_box, lam_box)

    # canonical reproduction anchor: S_radial at modes 0,1,2 must be the A313
    # C_4Dself diagonal (1, 1.222, 8.286), mu=1.222.
    repro_c4dself = abs(maps["S_radial"]["mu"] - 1.222) < 0.01

    # which canonical maps hit?
    canonical_hits = [m for m in CANONICAL_MAPS if maps[m]["hits_target"]]
    # which non-canonical maps hit?
    noncanon_hits = [m for m in ALL_MAPS
                     if m not in CANONICAL_MAPS and maps[m]["hits_target"]]
    any_hit = any(maps[m]["hits_target"] for m in ALL_MAPS)

    # closest map overall (by L2), and closest canonical map
    closest = min(ALL_MAPS, key=lambda m: maps[m]["L2_to_target"])
    closest_canon = min(CANONICAL_MAPS, key=lambda m: maps[m]["L2_to_target"])

    # NaN guard
    nan_found = False
    for m in maps:
        for val in maps[m]["diagonal_ratio"] + maps[m]["eigenvalues"]:
            if val != val:
                nan_found = True

    # ---- VERDICT ------------------------------------------------------
    if canonical_hits:
        verdict = "FORCED"
    elif noncanon_hits:
        verdict = "FIT_FLAGGED"
    else:
        verdict = "NULL"

    summary = {
        "addendum": "A316",
        "addresses": "OI-287-1",
        "builds_on": ["A311", "A313", "A315"],
        "thesis": ("the magnitude residual is a FAMILY-TO-MODE SELECTION rule "
                   "(A315): which eigenstate of the forced 4D-radial operator each "
                   "family maps to. Enumerate a short pre-declared list of "
                   "principled maps; which reproduces (1,1.5,8) with a peak; is the "
                   "corpus-CANONICAL map (node-count or Z_3-character) the one?"),
        "target_normalized": TARGET,
        "constants": {"OMEGA_alpha_inv": OMEGA, "E_self": E_SELF,
                      "gamma_layer_cycle": GAMMA},
        "canonical_maps": sorted(CANONICAL_MAPS),
        "all_maps_declared": ALL_MAPS,
        "n_maps_declared": len(ALL_MAPS),
        "S_radial_reproduces_A313_C4Dself": bool(repro_c4dself),
        "maps": maps,
        "canonical_maps_that_hit": sorted(canonical_hits),
        "noncanonical_maps_that_hit": sorted(noncanon_hits),
        "any_map_hits": bool(any_hit),
        "closest_map_overall": closest,
        "closest_map_diagonal": maps[closest]["diagonal_ratio"],
        "closest_map_mu": maps[closest]["mu"],
        "closest_map_tau": maps[closest]["tau"],
        "closest_map_L2": maps[closest]["L2_to_target"],
        "closest_canonical_map": closest_canon,
        "closest_canonical_diagonal": maps[closest_canon]["diagonal_ratio"],
        "closest_canonical_L2": maps[closest_canon]["L2_to_target"],
        "S_Z3_structural_note": maps["S_Z3"]["note"],
        "S_Z3_degenerate": maps["S_Z3"]["degenerate_selection"],
        "blocking_feature": (
            "no declared map both lands within 20%% of (1,1.5,8) and restores a "
            "valid column-standardized peak; the canonical S_radial gives "
            "mu=%.4f (short of 1.5, no peak) and the canonical S_Z3 is "
            "rank-degenerate in the real radial reduction (k=1,2 collapse)"
            % maps["S_radial"]["mu"]) if verdict == "NULL" else None,
        "nan_found": nan_found,
        "verdict": verdict,
    }

    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, "w") as f:
        json.dump(summary, f, indent=2)

    print("=" * 72)
    print("A316 PROBE -- OI-287-1: the FAMILY-TO-MODE selection rule. Which")
    print("              principled map gives (1,1.5,8) with a peak; is it FORCED?")
    print("=" * 72)
    print("target (normalized):", TARGET)
    print(f"S_radial reproduces A313 C_4Dself (mu~1.222): {repro_c4dself}")
    print(f"canonical maps (privileged by corpus): {sorted(CANONICAL_MAPS)}")
    print("-" * 72)
    for name in ALL_MAPS:
        m = maps[name]
        tag = "CANON" if m["canonical"] else "     "
        print(f"{name:10s} [{tag}] diag={m['diagonal_ratio']} mu={m['mu']} "
              f"tau={m['tau']} peak={m['valid_peak']} "
              f"<=20%={m['within_20pct']} HIT={m['hits_target']} "
              f"L2={m['L2_to_target']}")
        print(f"             nodes={m['node_counts']} "
              f"eig={m['eigenvalues']} degen={m['degenerate_selection']}")
    print("-" * 72)
    print(f"maps declared: {len(ALL_MAPS)} (multiple-comparison count)")
    print(f"canonical maps that hit: {sorted(canonical_hits)}")
    print(f"non-canonical maps that hit: {sorted(noncanon_hits)}")
    print(f"closest map overall: {closest} diag={maps[closest]['diagonal_ratio']} "
          f"L2={maps[closest]['L2_to_target']}")
    print(f"closest canonical map: {closest_canon} "
          f"diag={maps[closest_canon]['diagonal_ratio']} "
          f"L2={maps[closest_canon]['L2_to_target']}")
    print(f"S_Z3 note: {maps['S_Z3']['note']}")
    print(f"NaN found: {nan_found}")
    print(f"\nVERDICT: {verdict}")
    if verdict == "FORCED":
        print(f"  -> a CORPUS-CANONICAL map ({sorted(canonical_hits)}) reproduces")
        print("     (1,1.5,8) within 20%% with a valid peak. This is a genuine")
        print("     DERIVATION; the magnitude piece of OI-287-1 CLOSES.")
    elif verdict == "FIT_FLAGGED":
        print(f"  -> a NON-canonical map ({sorted(noncanon_hits)}) hits the target,")
        print("     but the canonical maps do NOT. This is a FITTED selection,")
        print("     explicitly NOT a derivation. Closest canonical map: "
              f"{closest_canon} (L2={maps[closest_canon]['L2_to_target']}).")
    else:
        print("  -> NO declared map reproduces (1,1.5,8) with a valid peak. The")
        print(f"     closest is {closest} (diag={maps[closest]['diagonal_ratio']}, "
              f"L2={maps[closest]['L2_to_target']}).")
        print("     The canonical S_radial gives mu short of 1.5 with no peak;")
        print("     the canonical S_Z3 is rank-degenerate in the real radial")
        print("     reduction (k=1,2 collapse). The family-to-mode rule that")
        print("     lands the target with a peak is not among the principled maps.")
    print(f"wrote {OUT}")


if __name__ == "__main__":
    main()
