#!/usr/bin/env python3
# Copyright Léon Fernando Vlegels -- CC BY 4.0
"""probe_selection_measure.py -- Probe for OI-287-1 (localizing factor, A307).

THREE prior exclusions bound this problem. A300 ruled out a factor depending
only on the FAMILY/node index n. A303 ruled out a factor depending only on the
monomial DEGREE k. A305 ruled out ANY kernel acting on the existing overlap
matrix B: every kernel that equals one on the matched pair (n=k) leaves the
diagonal entry untouched, so mu cannot move from 1.12 toward 1.5; the
localizing information is simply NOT PRESENT in the overlap matrix B to be
recovered by any post-overlap reweighting.

Those three exclusions force a single conclusion: the localizing factor must
enter EARLIER than B, either in the INNER PRODUCT / MEASURE used to form the
overlap, or in the EIGENSTATES themselves (a deformed operator). This probe
tests exactly that locus. No fitting.

A299 base recap:
  R[n][k] = <psi_n | c_k pi^k x^k | psi_n>,  B = R / lambda_n,
  diagonal normalizes to (1, 1.12, 7.50) against target (1, 1.5, 8), L2=0.63.

Operator (canonical):
  H = -d^2/dx^2 + [rho'(x)]^2/(2 Omega^2) + E_self x^2 (1-x)^2
  rho'(x) = 48 pi^3 x^2 + 6 pi^2 x + 2 pi
  Omega   = 4 pi^3 + pi^2 + pi
  E_self  = 13.177
  rho(x)  = 16 pi^3 x^3 + 3 pi^2 x^2 + 2 pi x
  Dirichlet [0,1], N=1500, psi normalized sum(psi^2) dx = 1.

PRE-REGISTERED CANDIDATES (declared before running; principled, corpus-named):
  (M5) CONTROL flat measure dx: recovers the A299 base (1, 1.12, 7.50).
  (M1) DENSITY-WEIGHTED MEASURE: form overlap and eigenstate normalization in
       the measure rho(x) dx instead of dx.
         R_rho[n][k] = <psi_n | c_k pi^k x^k | psi_n>_rho / <psi_n|psi_n>_rho,
         <f>_rho = integral f rho dx,
       and lambda_n from the rho-weighted generalized eigenproblem
         H psi = lambda rho psi.
       The natural geometric-density inner product of P18's density-coupling
       sector. Compute the diagonal.
  (M2) MEASURE x^a dx for a small fixed corpus-motivated set a in {1,2,3}
       (the layer / dimension counts). lambda_n from H psi = lambda x^a psi;
       overlap and normalization in x^a dx. Tests whether ANY simple measure
       shift lifts mu.
  (M3) DEFORMED OPERATOR: add the density coupling INTO the operator as P18
       has it, H' = H + alpha rho(x) with alpha = 1/Omega; recompute
       eigenstates psi'_n; form the SAME A299 overlap (flat dx) in psi'_n.
       Tests whether the localizing info is in the dressed eigenstates.
  (M4) GENERALIZED eigenproblem H psi = lambda W(x) psi with W = rho (a
       Sturm-Liouville weight); the eigenstates themselves change; form the
       diagonal overlap in these (flat dx for the overlap, W-normalized
       eigenstates). Distinct from M1: M1 forms the overlap in the rho
       measure too; M4 changes only the eigenstates.

For each candidate: normalized diagonal ratio (r1=1 by construction), valid
peak (column-standardized argmax_n of column k == k for all k AND the
sqrt(2) sin((n+1) pi x) box control, weighted by the SAME construction, FAILS
the same test), L2 to (1, 1.5, 8), and whether mu moved from 1.12 toward 1.5.

PRE-REGISTERED VERDICT:
  HIT iff some declared candidate yields BOTH a valid diagonal peak AND all
  three normalized entries within 20% of (1, 1.5, 8).
  PARTIAL-SIGNAL (reported distinctly, NOT a hit): some candidate moves mu
  materially toward 1.5 (mu in [1.3, 1.7]) even without a full peak. This
  would localize WHERE the factor lives (measure vs eigenstates) and is a
  real positive result short of HIT.
  NULL = a TENTH excluded mechanism class; report closest by L2 and mu
  movement.
Multiple-comparison caveat: candidates are counted. Honesty over a hit.
"""
import os
import json
import math

import numpy as np

# scipy is preferred for the generalized eigenproblem; fall back to a numpy
# symmetric (Cholesky) reduction when scipy is unavailable. Both solve the
# same generalized symmetric-definite problem H v = lambda W v with W SPD.
try:
    from scipy.linalg import eigh as _sp_eigh

    def gen_eigh(A, Bm):
        return _sp_eigh(A, Bm)
except Exception:
    def gen_eigh(A, Bm):
        # W = L L^T (Cholesky). Solve standard symmetric problem
        #   C = L^{-1} A L^{-T},  C y = lambda y,  v = L^{-T} y.
        L = np.linalg.cholesky(Bm)
        Linv = np.linalg.inv(L)
        C = Linv @ A @ Linv.T
        C = 0.5 * (C + C.T)
        w, y = np.linalg.eigh(C)
        v = Linv.T @ y
        return w, v


PI = math.pi
OMEGA = 4 * PI**3 + PI**2 + PI            # alpha^{-1}
E_SELF = 13.177
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_measure.json")


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 * 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 hamiltonian(x, dx, potential):
    off = -1.0 / dx**2 * np.ones(N - 2)
    return (np.diag(2.0 / dx**2 + potential(x))
            + np.diag(off, 1) + np.diag(off, -1))


def eigfuns_flat(potential):
    """Standard eigenproblem H psi = lambda psi, flat-dx normalized."""
    x, dx = grid()
    H = hamiltonian(x, dx, potential)
    w, v = np.linalg.eigh(H)
    return x, dx, w[:3], v[:, :3] / math.sqrt(dx)


def eigfuns_weighted(potential, weight):
    """Generalized eigenproblem H psi = lambda W psi with W = diag(weight(x)).

    Discretized: H_mat psi = lambda W_mat psi with W_mat = diag(weight*dx)
    on the same finite-difference grid. The dx factor cancels in lambda but
    keeps W_mat the discrete representation of the measure weight(x) dx.
    psi normalized so sum(psi^2 * weight) dx = 1 (the W-inner-product norm).
    """
    x, dx = grid()
    H = hamiltonian(x, dx, potential)
    Wdiag = weight(x)
    if np.any(Wdiag <= 0):
        raise ValueError("weight must be strictly positive on the interior")
    Wm = np.diag(Wdiag)
    w, v = gen_eigh(H, Wm)
    order = np.argsort(w)
    w = w[order]
    v = v[:, order]
    vv = v[:, :3].copy()
    # W-normalize: sum(psi^2 * weight) dx = 1
    for n in range(3):
        nrm = math.sqrt(float(np.sum(vv[:, n]**2 * Wdiag) * 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 density_matrix(psi, x, dx, measure=None):
    """R[n][k] = <psi_n | c_k pi^k x^k | psi_n>_measure / <psi_n|psi_n>_measure.

    measure=None -> flat dx (the A299 form, psi already sum(psi^2)dx=1).
    measure=array -> integral weight; normalization divides by the measure norm.
    """
    if measure is None:
        return np.array(
            [[float(np.sum(psi[:, n]**2 * C[k - 1] * PI**k * x**k) * dx)
              for k in (1, 2, 3)] for n in range(3)])
    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):
            num = float(np.sum(psi[:, n]**2 * C[k - 1] * PI**k * x**k
                               * measure) * dx)
            M[n, k - 1] = num / 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(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 evaluate(name, psi, lam, x, dx, psi_box, lam_box, measure=None):
    """Build B = R / lambda for both the operator states and the box control,
    in the given measure, and report the diagonal / peak / L2 / mu-movement."""
    Dm = density_matrix(psi, x, dx, measure)
    B = Dm / np.array(lam)[:, None]
    Dm_box = density_matrix(psi_box, x, dx, measure)
    B_box = Dm_box / lam_box[:, None]
    ratio = norm_diag(B)
    peak = bool(diag_peaks(B))
    box_peak = bool(diag_peaks(B_box))
    valid_peak = peak and not box_peak
    dist = l2(ratio)
    within20 = all(abs(ratio[i] - TARGET[i]) <= 0.20 * TARGET[i]
                   for i in range(3))
    mu = ratio[1]
    return {
        "diagonal_ratio": [round(v, 4) for v in ratio],
        "peaks_diagonal": peak,
        "box_control_peaks": box_peak,
        "valid_peak": valid_peak,
        "L2_to_target": round(dist, 4),
        "within_20pct": bool(within20),
        "mu": round(mu, 4),
        "mu_in_partial_band": bool(1.3 <= mu <= 1.7),
        "eigenvalues": [round(float(v), 4) for v in lam],
    }


def main():
    # ---- flat-measure baseline states (the A299 operator) ----------------
    x, dx, lam_flat, psi_flat = eigfuns_flat(Vpot)
    nu = [node_count(psi_flat[:, n]) for n in range(3)]
    psi_box, lam_box = box_states(x)

    results = {}

    # M5 control: flat dx -> recovers A299 (1, 1.12, 7.50)
    results["M5_control_flat"] = evaluate(
        "M5", psi_flat, lam_flat, x, dx, psi_box, lam_box, measure=None)

    # M1 density-weighted measure rho(x) dx, with rho-weighted eigenproblem
    rho_x = rho(x)
    _, _, lam_rho, psi_rho = eigfuns_weighted(Vpot, rho)
    nu_rho = [node_count(psi_rho[:, n]) for n in range(3)]
    results["M1_measure_rho"] = evaluate(
        "M1", psi_rho, lam_rho, x, dx, psi_box, lam_box, measure=rho_x)

    # M2 measure x^a dx, a in {1,2,3}; eigenproblem H psi = lambda x^a psi
    for a in (1, 2, 3):
        wfun = (lambda aa: (lambda xx: xx**aa))(a)
        wx = x**a
        _, _, lam_a, psi_a = eigfuns_weighted(Vpot, wfun)
        results[f"M2_measure_x{a}"] = evaluate(
            f"M2_x{a}", psi_a, lam_a, x, dx, psi_box, lam_box, measure=wx)

    # M3 deformed operator H' = H + alpha rho(x), alpha = 1/Omega; flat overlap
    alpha = 1.0 / OMEGA

    def Vdef(xx):
        return Vpot(xx) + alpha * rho(xx)

    _, _, lam_def, psi_def = eigfuns_flat(Vdef)
    results["M3_deformed_op"] = evaluate(
        "M3", psi_def, lam_def, x, dx, psi_box, lam_box, measure=None)

    # M4 generalized SL eigenproblem H psi = lambda rho psi; flat overlap,
    # rho-W-normalized eigenstates (states change; overlap stays flat dx)
    results["M4_sturm_liouville_rho"] = evaluate(
        "M4", psi_rho, lam_rho, x, dx, psi_box, lam_box, measure=None)

    # ---- aggregate -------------------------------------------------------
    a299_mu = results["M5_control_flat"]["mu"]
    non_control = {k: v for k, v in results.items() if k != "M5_control_flat"}
    closest = min(non_control, key=lambda k: non_control[k]["L2_to_target"])
    closest_ratio = results[closest]["diagonal_ratio"]
    closest_mu = results[closest]["mu"]

    # mu movement toward 1.5 for each candidate, relative to A299 base 1.12
    for k, v in results.items():
        v["mu_moved_toward_1.5"] = bool(
            abs(v["mu"] - 1.5) < abs(a299_mu - 1.5))

    any_valid = any(v["valid_peak"] for v in results.values())
    hit = any(v["valid_peak"] and v["within_20pct"]
              for k, v in non_control.items())

    # partial-signal: any non-control candidate with mu in [1.3, 1.7]
    partial_candidates = [k for k, v in non_control.items()
                          if v["mu_in_partial_band"]]
    partial = len(partial_candidates) > 0
    any_mu_moved = any(v["mu_moved_toward_1.5"] for v in non_control.values())

    if hit:
        verdict = "HIT"
    elif partial:
        verdict = "PARTIAL_SIGNAL"
    else:
        verdict = "NULL_TENTH_CLASS"

    # NaN guard
    nan_found = False
    for v in results.values():
        for val in v["diagonal_ratio"] + v["eigenvalues"]:
            if val != val:
                nan_found = True

    summary = {
        "addendum": "A307",
        "addresses": "OI-287-1",
        "target_normalized": TARGET,
        "a299_base_diagonal": results["M5_control_flat"]["diagonal_ratio"],
        "a299_base_mu": a299_mu,
        "node_counts_flat": nu,
        "node_counts_rho": nu_rho,
        "constants": {"OMEGA_alpha_inv": OMEGA, "E_self": E_SELF,
                      "alpha_deform": alpha},
        "candidates": results,
        "n_candidates_declared": len(results),
        "n_noncontrol_candidates": len(non_control),
        "closest_noncontrol": closest,
        "closest_diagonal": closest_ratio,
        "closest_L2": results[closest]["L2_to_target"],
        "closest_mu": closest_mu,
        "any_valid_peak": any_valid,
        "any_mu_moved_toward_1.5": bool(any_mu_moved),
        "partial_signal_candidates": partial_candidates,
        "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("=" * 66)
    print("A307 PROBE -- OI-287-1 measure / deformed-eigenstate localizing factor")
    print("=" * 66)
    print("target (normalized):", TARGET)
    print("A299 base diagonal :", results["M5_control_flat"]["diagonal_ratio"],
          " (control, flat dx)")
    print("node counts flat:", nu, " rho:", nu_rho)
    print("-" * 66)
    for name, v in results.items():
        print(f"{name:24s} diag={v['diagonal_ratio']}  "
              f"peak={v['peaks_diagonal']} box={v['box_control_peaks']} "
              f"valid={v['valid_peak']} L2={v['L2_to_target']} "
              f"mu={v['mu']} muMoved={v['mu_moved_toward_1.5']} "
              f"<=20%={v['within_20pct']}")
    print("-" * 66)
    print(f"candidates declared: {len(results)} "
          f"(non-control: {len(non_control)})")
    print(f"closest non-control: {closest} diag={closest_ratio} "
          f"L2={results[closest]['L2_to_target']} mu={closest_mu}")
    print(f"any valid diagonal peak: {any_valid}")
    print(f"any mu moved toward 1.5: {any_mu_moved}")
    print(f"partial-signal candidates (mu in [1.3,1.7]): {partial_candidates}")
    print(f"NaN found: {nan_found}")
    print(f"\nVERDICT: {verdict}")
    if verdict == "NULL_TENTH_CLASS":
        print("  -> tenth excluded mechanism class: measure shift / deformed")
        print("     eigenstates (no declared candidate peaks within 20%, and")
        print("     no candidate moved mu into the partial band [1.3,1.7]).")
    elif verdict == "PARTIAL_SIGNAL":
        print("  -> mu moved materially toward 1.5 without a full peak:")
        print(f"     localizes the factor to {partial_candidates}.")
    print(f"wrote {OUT}")


if __name__ == "__main__":
    main()
