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

A309 builds directly on A307. A307 localized the missing factor of the A299
selection rule to the INNER-PRODUCT MEASURE: with the density measure rho(x) dx
and the rho-weighted generalized eigenproblem H psi = lambda rho psi, mu lifts
from the A299 flat-measure base 1.12 to 1.34, diagonal (1, 1.34, 8.72); and the
M1-vs-M4 contrast in A307 proved the lift is attached to the MEASURE, not the
eigenstates. A307 was a PARTIAL-SIGNAL: mu moved toward 1.5 but no declared
measure restored a valid diagonal peak, and the exact measure was not pinned.

A309 sweeps a DECLARED, PRINCIPLED one-parameter / small-discrete family of
density-type measures W(x). No arbitrary fitting: every member is either a
corpus-named object (the P18 density rho, its powers, the P18 moment-measure
x^m rho) or a natural geometric / dimension-count weight (x^a). For each W we
solve the generalized eigenproblem H psi = lambda W psi by the numpy Cholesky
reduction (scipy unavailable in the sandbox: L = chol(W), C = L^{-1} H L^{-T},
eigh(C), psi = L^{-T} y), form the A299-style diagonal overlap IN the measure W

    R[n][k] = <psi_n | c_k pi^k x^k | psi_n>_W / <psi_n|psi_n>_W ,
    B = R / lambda_n,   <f>_W = integral f W dx ,

and report: the normalized diagonal ratio (r1 = 1 by construction), valid_peak
(column-standardized argmax over rows of column k equals k for all k AND the
sqrt(2) sin((n+1) pi x) box control, weighted the SAME way, FAILS the same
test), L2 to the normalized target (1, 1.5, 8), and mu = ratio[1].

Operator (canonical, identical to A307/A299):
  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.  C = (2, 3, 16), target normalized (1, 1.5, 8).
  Box control psi_box[n] = sqrt(2) sin((n+1) pi x).

DECLARED MEASURE FAMILY (listed up front; all principled; counted for the
multiple-comparison caveat):
  (W-flat)  W = 1                       CONTROL. Must recover A299 (1, 1.12, 7.50).
  (W-rho)   W = rho(x)                  P18 geometric density measure. The A307
                                        lead; expected mu ~ 1.34.
  (W-pow)   W = x^a, a in {1, 2, 3}     layer / dimension-count weights. A307 saw
                                        x^1 closest by L2.
  (W-rhop)  W = rho(x)^p, p in {1/2, 1, 2}  powers of the density: does a power
                                        of rho sharpen the peak? (p=1 duplicates
                                        W-rho; kept once, the duplicate dropped.)
  (W-mom)   W = x^m rho(x), m in {0, 1, 2}  P18 moment-measure mu_n = int x^n rho.
                                        (m=0 duplicates W-rho; kept once.)

Unique members actually evaluated (duplicates folded):
  W=1 ; W=x ; W=x^2 ; W=x^3 ; W=rho ; W=rho^{1/2} ; W=rho^2 ;
  W=x*rho ; W=x^2*rho .
That is 9 measures (1 control + 8 non-control). All strictly positive on the
open interval (0,1): rho>0 there, x^a>0, products of positives positive, so
every W is a valid positive-definite Sturm-Liouville weight.

PRE-REGISTERED VERDICT (declared before running):
  HIT iff some declared PRINCIPLED non-control measure yields BOTH a valid
       diagonal peak AND all three normalized entries within 20% of (1,1.5,8).
  STRENGTHENED-PARTIAL iff no full peak (no HIT) but the family pins a
       DISTINGUISHED measure -- specifically either
         (a) the L2-closest non-control member is itself a corpus-canonical
             measure (rho, a power of rho, or a moment-measure x^m rho), OR
         (b) mu is MONOTONE in a declared one-parameter sub-family exponent and
             crosses a corpus-meaningful value,
       improving on or sharpening A307's partial signal. Report the exact
       closest measure, its diagonal, its L2, and its mu.
  NULL-ISH iff the family does not improve on A307: closest L2 >= A307's
       0.3252 and no canonical member is distinguished.

The honest expectation from A307 is STRENGTHENED-PARTIAL at best: the flat-x
power x^1 was closest there (L2=0.33) and is NOT corpus-canonical, while the
canonical rho sat at mu=1.34 but L2=0.74. A309 asks whether ADDING the moment
and rho-power members changes that picture -- whether a canonical density-type
measure becomes the distinguished point, or whether the closest member stays
non-canonical (which would leave rho dx UN-forced by L2 and the result an
honest strengthened-partial resting on the M1-vs-M4 measure-locus argument, not
on rho being the L2 minimum).

Multiple-comparison caveat: 8 non-control measures are declared and counted.
Honesty over a hit. No derivation/closure claim unless the HIT criterion is met
verbatim.
"""
import os
import json
import math

import numpy as np


def gen_eigh(A, Bm):
    """Generalized symmetric-definite eig H v = lambda W v, W SPD, via numpy
    Cholesky reduction (scipy unavailable in the sandbox).
        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
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_measurefamily.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, W = diag(weight(x)).

    psi W-normalized: sum(psi^2 * weight) dx = 1.
    """
    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()
    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):
    """R[n][k] = <psi_n | c_k pi^k x^k | psi_n>_W / <psi_n|psi_n>_W.

    measure is an array (the weight W(x) on the grid). The flat control passes
    W = ones, which reproduces the A299 flat overlap exactly.
    """
    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(psi, lam, x, dx, psi_box, lam_box, measure):
    """Diagonal / peak / L2 / mu for one weighted measure W = measure(array)."""
    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],
        "tau": round(ratio[2], 4),
        "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],
        "node_counts": [node_count(psi[:, n]) for n in range(3)],
    }


# Canonical-measure tag: corpus-named density-type weights (rho, its powers,
# moment-measures x^m rho). The flat x^a powers are NOT tagged canonical: they
# are natural geometric weights but not P18 density objects.
def is_canonical(name):
    return name in ("W_rho", "W_rho_pow0.5", "W_rho_pow2.0",
                    "W_xrho_m1", "W_x2rho_m2")


def main():
    x, dx = grid()
    # flat-measure baseline states (the A299 operator)
    _, _, lam_flat, psi_flat = eigfuns_flat(Vpot)
    psi_box, lam_box = box_states(x)
    ones = np.ones_like(x)
    rho_x = rho(x)

    results = {}

    # ---- (W-flat) control ----
    results["W_flat"] = evaluate(psi_flat, lam_flat, x, dx,
                                 psi_box, lam_box, ones)

    # ---- (W-pow) W = x^a, a in {1,2,3} ----
    for a in (1, 2, 3):
        wfun = (lambda aa: (lambda xx: xx**aa))(a)
        _, _, lam_a, psi_a = eigfuns_weighted(Vpot, wfun)
        results[f"W_pow_x{a}"] = evaluate(psi_a, lam_a, x, dx,
                                          psi_box, lam_box, x**a)

    # ---- (W-rho) and (W-rhop) W = rho^p, p in {0.5, 1, 2} ----
    for p in (0.5, 1.0, 2.0):
        wfun = (lambda pp: (lambda xx: rho(xx)**pp))(p)
        _, _, lam_p, psi_p = eigfuns_weighted(Vpot, wfun)
        key = "W_rho" if p == 1.0 else f"W_rho_pow{p}"
        results[key] = evaluate(psi_p, lam_p, x, dx,
                                psi_box, lam_box, rho_x**p)

    # ---- (W-mom) W = x^m rho, m in {1,2}  (m=0 == W_rho, already done) ----
    for m in (1, 2):
        wfun = (lambda mm: (lambda xx: xx**mm * rho(xx)))(m)
        _, _, lam_m, psi_m = eigfuns_weighted(Vpot, wfun)
        key = "W_xrho_m1" if m == 1 else "W_x2rho_m2"
        results[key] = evaluate(psi_m, lam_m, x, dx,
                                psi_box, lam_box, x**m * rho_x)

    # ---- aggregate -------------------------------------------------------
    flat_mu = results["W_flat"]["mu"]
    non_control = {k: v for k, v in results.items() if k != "W_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"]
    closest_l2 = results[closest]["L2_to_target"]
    closest_is_canonical = is_canonical(closest)

    for k, v in results.items():
        v["mu_moved_toward_1.5"] = bool(
            abs(v["mu"] - 1.5) < abs(flat_mu - 1.5))
        v["is_canonical_measure"] = is_canonical(k)

    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_candidates = [k for k, v in non_control.items()
                          if v["mu_in_partial_band"]]

    # mu monotone in the rho-power exponent p?  p: 0.5 -> 1 -> 2
    rho_pow_mus = [results["W_rho_pow0.5"]["mu"], results["W_rho"]["mu"],
                   results["W_rho_pow2.0"]["mu"]]
    rho_pow_monotone = (rho_pow_mus[0] <= rho_pow_mus[1] <= rho_pow_mus[2]) \
        or (rho_pow_mus[0] >= rho_pow_mus[1] >= rho_pow_mus[2])

    # A307 reference: closest non-control there was M2_x1 at L2 = 0.3252.
    a307_closest_l2 = 0.3252
    improves_l2 = closest_l2 < a307_closest_l2

    # STRENGTHENED-PARTIAL conditions (a) or (b)
    cond_a = closest_is_canonical
    # corpus-meaningful crossing value for monotone mu sub-family: 1.5 target
    cond_b = rho_pow_monotone and (min(rho_pow_mus) < 1.5 < max(rho_pow_mus))
    strengthened_partial = (not hit) and (cond_a or cond_b)
    nullish = (not hit) and (not strengthened_partial) \
        and (closest_l2 >= a307_closest_l2) and (not closest_is_canonical)

    if hit:
        verdict = "HIT"
    elif strengthened_partial:
        verdict = "STRENGTHENED_PARTIAL"
    elif nullish:
        verdict = "NULL_ISH"
    else:
        # no full peak, closest not canonical, no monotone crossing, but L2
        # may or may not beat A307: honest residual partial pointing to A307.
        verdict = "PARTIAL_NO_NEW_DISTINGUISHED"

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

    summary = {
        "addendum": "A309",
        "addresses": "OI-287-1",
        "builds_on": "A307",
        "target_normalized": TARGET,
        "flat_control_diagonal": results["W_flat"]["diagonal_ratio"],
        "flat_control_mu": flat_mu,
        "W_rho_diagonal": results["W_rho"]["diagonal_ratio"],
        "W_rho_mu": results["W_rho"]["mu"],
        "constants": {"OMEGA_alpha_inv": OMEGA, "E_self": E_SELF},
        "measures": results,
        "n_measures_declared": len(results),
        "n_noncontrol_measures": len(non_control),
        "closest_noncontrol": closest,
        "closest_diagonal": closest_ratio,
        "closest_L2": closest_l2,
        "closest_mu": closest_mu,
        "closest_is_canonical_measure": closest_is_canonical,
        "a307_closest_L2": a307_closest_l2,
        "improves_on_a307_L2": bool(improves_l2),
        "rho_power_mus_p_0.5_1_2": [round(m, 4) for m in rho_pow_mus],
        "rho_power_mu_monotone": bool(rho_pow_monotone),
        "rho_power_mu_crosses_1.5": bool(
            min(rho_pow_mus) < 1.5 < max(rho_pow_mus)),
        "any_valid_peak": any_valid,
        "partial_signal_candidates": partial_candidates,
        "strengthened_partial_cond_a_closest_canonical": bool(cond_a),
        "strengthened_partial_cond_b_monotone_crossing": bool(cond_b),
        "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("=" * 70)
    print("A309 PROBE -- OI-287-1 principled density-measure FAMILY sweep")
    print("=" * 70)
    print("target (normalized):", TARGET)
    print("flat control diagonal:", results["W_flat"]["diagonal_ratio"],
          " (W=1, recovers A299)")
    print("-" * 70)
    for name, v in results.items():
        tag = "CANON" if v["is_canonical_measure"] else "     "
        print(f"{name:16s} [{tag}] diag={v['diagonal_ratio']}  "
              f"valid={v['valid_peak']} L2={v['L2_to_target']} "
              f"mu={v['mu']} <=20%={v['within_20pct']}")
    print("-" * 70)
    print(f"measures declared: {len(results)} "
          f"(non-control: {len(non_control)})")
    print(f"closest non-control: {closest} diag={closest_ratio} "
          f"L2={closest_l2} mu={closest_mu} canonical={closest_is_canonical}")
    print(f"rho^p mus (p=0.5,1,2): {[round(m,4) for m in rho_pow_mus]} "
          f"monotone={rho_pow_monotone} crosses1.5="
          f"{min(rho_pow_mus) < 1.5 < max(rho_pow_mus)}")
    print(f"improves on A307 L2 (0.3252): {improves_l2}")
    print(f"any valid diagonal peak: {any_valid}")
    print(f"partial-band candidates (mu in [1.3,1.7]): {partial_candidates}")
    print(f"NaN found: {nan_found}")
    print(f"\nVERDICT: {verdict}")
    if verdict == "HIT":
        print("  -> a principled measure restored BOTH a valid peak and the")
        print("     diagonal within 20%. OI-287-1 closure candidate.")
    elif verdict == "STRENGTHENED_PARTIAL":
        if cond_a:
            print(f"  -> distinguished measure: closest member {closest} is")
            print("     corpus-canonical. The density-type measure is pinned.")
        if cond_b:
            print("  -> mu monotone in the rho-power exponent and crosses 1.5.")
    elif verdict == "NULL_ISH":
        print("  -> family does not improve on A307; no canonical member")
        print("     distinguished by L2.")
    else:
        print("  -> no full peak; closest member is non-canonical, so rho dx")
        print("     is NOT forced by L2. The measure locus from A307 stands on")
        print("     the M1-vs-M4 argument, not on rho being the L2 minimum.")
    print(f"wrote {OUT}")


if __name__ == "__main__":
    main()
