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

A300 bracketed the missing factor in the A299 base
  B[n][k] = <psi_n | c_k pi^k x^k | psi_n> / lambda_n,
whose diagonal normalizes to (1, 1.12, 7.50) against the target (1, 1.5, 8).
A300 ruled out the node count as a uniform mechanism: a linear node gain
(1+nu_n) overshoots to (1, 2.24, 22.50), and no single power (1+nu_n)^p
reconciles mu (wants p~0.42) and tau (wants p~0.06). A300's stated next
step: the localizing factor may be DEGREE-dependent -- a fixed function of
the monomial degree k, not of the family/node index n.

This probe pre-registers and tests a small fixed set of degree-dependent
candidates applied as a per-column weight w(k) on the A299 base. No fitting.

PRE-REGISTERED CANDIDATES (declared before running, principled / corpus-named):
  (D4) CONTROL: w(k) = 1            -- recovers A299 base exactly.
  (D1) w(k) = 1/k!                  -- Taylor/factorial normalization; natural
                                       if the coupling reads Taylor coefficients
                                       (rho's degree-k term carries 1/k!).
  (D2) w(k) = 1/k                   -- inverse-degree damping.
  (D3) w(k) = pi^{-k} * kappa * Omega -- degree-resolved strength: the
                                       kappa*Omega/lambda_n strength carrying an
                                       explicit pi^{-k} degree power.
                                       kappa = alpha^{5/4} = Omega^{-5/4},
                                       so kappa*Omega = Omega^{-1/4}.
  (D5) w(k) = kappa * Omega * k     -- degree-resolved strength with a linear
                                       k factor (the complementary sign of D3).

For each candidate: column-weight the A299 base, renormalize the diagonal to
the (1,*) entry, and report:
  - the normalized diagonal ratio (r1, r2, r3) with r1 == 1 by construction;
  - PEAK test: column-standardize each column to z-scores; require
    argmax_n Z[n][k] == k for all k (diagonal peak) AND the sine-box control
    psi_box[n] = sqrt(2) sin((n+1) pi x) FAILS the same test;
  - L2 distance of (r1,r2,r3) to the target (1, 1.5, 8).

PRE-REGISTERED VERDICT:
  HIT iff some declared degree-dependent candidate (D1,D2,D3,D5; D4 is the
  A299 control) yields BOTH a diagonal peak AND all three normalized entries
  within 20% of (1, 1.5, 8).
  Else NULL = an EIGHTH excluded mechanism class (priors: A292 universal
  first-order, A295 dressing, A296 bare diagonal, A297 position-moment,
  A298 nodal-direction, A299 inverse-eigenvalue, A300 node-count). Report the
  closest candidate by L2 and whether its mu entry moved toward 1.5.
Multiple-comparison caveat: five candidates declared (one is the control).
"""
import os
import json
import math

import numpy as np

PI = math.pi
OMEGA = 4 * PI**3 + PI**2 + PI            # alpha^{-1}
KAPPA = OMEGA ** (-1.25)                  # alpha^{5/4}
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_degree.json")


def eigfuns(potential):
    dx = 1.0 / N
    x = np.linspace(dx, 1 - dx, N - 1)
    off = -1.0 / dx**2 * np.ones(N - 2)
    H = (np.diag(2.0 / dx**2 + potential(x))
         + np.diag(off, 1) + np.diag(off, -1))
    w, v = np.linalg.eigh(H)
    return x, w[:3], v[:, :3] / math.sqrt(dx)


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):
    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)])


def diag_peaks(M):
    """Column-standardize; require argmax over rows 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 main():
    rho_p = lambda x: 48 * PI**3 * x**2 + 6 * PI**2 * x + 2 * PI
    Vpot = lambda x: rho_p(x)**2 / (2 * OMEGA**2) + E_SELF * x**2 * (1 - x)**2
    x, lam, psi = eigfuns(Vpot)
    dx = x[1] - x[0]
    nu = [node_count(psi[:, n]) for n in range(3)]

    # A299 base
    Dm = density_matrix(psi, x, dx)
    B = Dm / np.array(lam)[:, None]

    # sine-box control eigenfunctions on the same grid
    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)])
    Dm_box = density_matrix(psi_box, x, dx)
    B_box = Dm_box / lam_box[:, None]

    # declared degree weights w(k) for k=1,2,3
    candidates = {
        "D4_control_one":      [1.0, 1.0, 1.0],
        "D1_factorial":        [1.0 / math.factorial(k) for k in (1, 2, 3)],
        "D2_inverse_degree":   [1.0 / k for k in (1, 2, 3)],
        "D3_pi_power_strength":[PI**(-k) * KAPPA * OMEGA for k in (1, 2, 3)],
        "D5_linear_strength":  [KAPPA * OMEGA * k for k in (1, 2, 3)],
    }

    results = {}
    for name, w in candidates.items():
        wv = np.array(w)
        M = B * wv[None, :]
        M_box = B_box * wv[None, :]
        ratio = norm_diag(M)
        peak = bool(diag_peaks(M))
        box_peak = bool(diag_peaks(M_box))
        valid_peak = peak and not box_peak     # control must fail
        dist = l2(ratio)
        within20 = all(abs(ratio[i] - TARGET[i]) <= 0.20 * TARGET[i]
                       for i in range(3))
        results[name] = {
            "weights_k": [float(v) for v in w],
            "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),
        }

    # closest non-control candidate by L2
    non_control = {k: v for k, v in results.items() if k != "D4_control_one"}
    closest = min(non_control, key=lambda k: non_control[k]["L2_to_target"])
    closest_ratio = results[closest]["diagonal_ratio"]

    # A299 control mu (D4)
    a299_mu = results["D4_control_one"]["diagonal_ratio"][1]
    closest_mu = closest_ratio[1]
    mu_moved_toward = abs(closest_mu - 1.5) < abs(a299_mu - 1.5)

    hit = any(v["valid_peak"] and v["within_20pct"]
              for k, v in results.items() if k != "D4_control_one")

    verdict = "HIT" if hit else "NULL_EIGHTH_CLASS"

    summary = {
        "addendum": "A303",
        "addresses": "OI-287-1",
        "eigenvalues": [round(float(v), 4) for v in lam],
        "node_counts": nu,
        "target_normalized": TARGET,
        "a299_base_diagonal": results["D4_control_one"]["diagonal_ratio"],
        "constants": {"OMEGA_alpha_inv": OMEGA, "KAPPA_alpha_5_4": KAPPA,
                      "kappa_times_Omega": KAPPA * OMEGA},
        "candidates": results,
        "n_candidates_declared": len(candidates),
        "n_noncontrol_candidates": len(non_control),
        "closest_noncontrol": closest,
        "closest_diagonal": closest_ratio,
        "closest_L2": results[closest]["L2_to_target"],
        "mu_moved_toward_1.5": bool(mu_moved_toward),
        "any_valid_peak": any(v["valid_peak"] for v in results.values()),
        "verdict": verdict,
    }

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

    print("=" * 64)
    print("A303 PROBE -- OI-287-1 degree-dependent localizing factor")
    print("=" * 64)
    print("eigenvalues:", summary["eigenvalues"], " node counts:", nu)
    print("target (normalized):", TARGET)
    print("A299 base diagonal :", results["D4_control_one"]["diagonal_ratio"])
    print("-" * 64)
    for name, v in results.items():
        print(f"{name:22s} diag={v['diagonal_ratio']}  "
              f"peak={v['peaks_diagonal']} boxpeak={v['box_control_peaks']} "
              f"valid={v['valid_peak']} L2={v['L2_to_target']} "
              f"<=20%={v['within_20pct']}")
    print("-" * 64)
    print(f"candidates declared: {len(candidates)} "
          f"(non-control: {len(non_control)})")
    print(f"closest non-control: {closest} diag={closest_ratio} "
          f"L2={results[closest]['L2_to_target']}")
    print(f"mu moved toward 1.5: {mu_moved_toward} "
          f"(A299 mu={a299_mu}, closest mu={closest_mu})")
    print(f"any valid diagonal peak: {summary['any_valid_peak']}")
    print(f"\nVERDICT: {verdict}")
    if verdict == "NULL_EIGHTH_CLASS":
        print("  -> eighth excluded mechanism class: degree-dependent weight")
        print("     (no declared candidate both peaks and lands within 20%).")
    print(f"wrote {OUT}")


if __name__ == "__main__":
    main()
