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

A300 and A303 together excluded the two separable hypotheses for 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 a factor depending only on the FAMILY/node index n (a node
power overshoots; no single power reconciles mu and tau). A303 ruled out a
factor depending only on the monomial DEGREE k (a column weight cannot move a
column's argmax nor lift a diagonal entry relative to its own column). Both
results point to the same conclusion: the missing localizing factor must
couple n and k JOINTLY, suppressing entries with n != k, and must NOT be
separable into f(n)*g(k).

This probe pre-registers and tests a small fixed set of non-separable joint
kernels K[n,k] applied to the A299 base, M[n,k] = B[n,k] * K[n,k]. No fitting.

PRE-REGISTERED CANDIDATES (declared before running; principled / corpus-named):
  (J0) CONTROL: K[n,k] = 1                 -- recovers the A299 base exactly.
  (J1) Gaussian off-diagonal suppression:
       K[n,k] = exp(-(k-(n+1))^2 / (2 s^2)), s = 1 FIXED (no free param).
       Rewards the column k whose degree matches the family index (k = n+1),
       suppressing off-diagonal. Non-separable: depends on (k-(n+1)).
  (J2) Kronecker-band: K[n,k] = rho_band^{|k-(n+1)|}, rho_band = 1/2 FIXED.
       A banded suppression; 1 on the band k = n+1, geometric decay off it.
       Non-separable: depends on |k-(n+1)|.
  (J3) Self-structure kernel: K[n,k] = (R[n,k])^(p-1) normalized so the
       weighted entry is M[n,k] = B[n,k] * (overlap[n,k])^(p-1), i.e. the
       matrix suppresses its own off-diagonal via a fixed power p in {2}.
       overlap[n,k] = <psi_n| pi^k x^k |psi_n> (the unnormalized density
       moment). Non-separable: depends on both indices through psi_n. Tested
       at p=2 (J3a, overlap^1 multiplicative) and the limiting hard form.
  (J3a) p = 2: M[n,k] = B[n,k] * overlap_hat[n,k], overlap_hat row-normalized
       to its diagonal entry. Non-separable.
  (J4) Hard projection (limiting control): K[n,k] = 1 if n==k else 0.
       M = diag(B). Shows what perfect off-diagonal suppression gives.
  (J5) SEPARABILITY CONTROL: K[n,k] = f(n)*g(k) with declared
       f(n) = (1+nu_n) (node gain, the A300 mechanism) and
       g(k) = 1/k (inverse degree, the A303 mechanism). This MUST fail to
       beat the joint kernels, confirming non-separability is the operative
       feature. Separable by construction.

For each candidate: form M = B (.*) K, renormalize the diagonal to the (1,*)
entry, and report:
  - the normalized diagonal ratio (r1, r2, r3) with r1 == 1 by construction;
  - valid_peak: 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), weighted by the SAME K, FAILS the
    same test;
  - L2 distance of (r1,r2,r3) to the target (1, 1.5, 8).

PRE-REGISTERED VERDICT:
  HIT iff some declared joint kernel (J1,J2,J3a,J4,J5; J0 is the control)
  yields BOTH a valid diagonal peak AND all three normalized entries within
  20% of (1, 1.5, 8).
  Else NULL = a NINTH 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, A303
  degree). Report the closest candidate by L2, whether any peaked, and
  whether the separable control (J5) underperformed the joint kernels.
Multiple-comparison caveat: six candidates declared (one is the control).
Honesty over a hit.
"""
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_jointkernel.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 overlap_matrix(psi, x, dx):
    """Unnormalized density moments <psi_n | pi^k x^k | psi_n> (no c_k)."""
    return np.array([[float(np.sum(psi[:, n]**2 * PI**k * x**k) * dx)
                      for k in (1, 2, 3)] for n in range(3)])


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]

    # unnormalized overlap (for the self-structure kernel J3a)
    OV = overlap_matrix(psi, x, dx)
    # row-normalize the overlap to its own diagonal entry -> overlap_hat
    OV_hat = OV / np.array([OV[n, n] for n in range(3)])[:, 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]
    OV_box = overlap_matrix(psi_box, x, dx)
    OV_box_hat = OV_box / np.array([OV_box[n, n] for n in range(3)])[:, None]

    # ---- declared joint kernels K[n,k] (n,k in {0,1,2} indexing row/col) ----
    def K_gauss(s=1.0):
        return np.array([[math.exp(-((k - (n + 1))**2) / (2 * s**2))
                          for k in (1, 2, 3)] for n in range(3)])

    def K_band(rho=0.5):
        return np.array([[rho**abs(k - (n + 1))
                          for k in (1, 2, 3)] for n in range(3)])

    def K_hard():
        return np.eye(3)

    def K_separable():
        f = np.array([1 + nu[n] for n in range(3)], dtype=float)   # node gain
        g = np.array([1.0 / k for k in (1, 2, 3)])                 # inv degree
        return f[:, None] * g[None, :]

    # J3a self-structure: K = overlap_hat (so M = B .* overlap_hat).
    # For the box control, use the box's own overlap_hat.
    kernels = {
        "J0_control_one":      (np.ones((3, 3)),        np.ones((3, 3))),
        "J1_gaussian_s1":      (K_gauss(1.0),           K_gauss(1.0)),
        "J2_band_half":        (K_band(0.5),            K_band(0.5)),
        "J3a_self_structure":  (OV_hat,                 OV_box_hat),
        "J4_hard_projection":  (K_hard(),               K_hard()),
        "J5_separable_control":(K_separable(),          K_separable()),
    }

    results = {}
    for name, (K, Kb) in kernels.items():
        M = B * K
        M_box = B_box * Kb
        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 the test
        dist = l2(ratio)
        within20 = all(abs(ratio[i] - TARGET[i]) <= 0.20 * TARGET[i]
                       for i in range(3))
        results[name] = {
            "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),
            "separable": name == "J5_separable_control",
        }

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

    a299_mu = results["J0_control_one"]["diagonal_ratio"][1]
    closest_mu = closest_ratio[1]
    mu_moved_toward = abs(closest_mu - 1.5) < abs(a299_mu - 1.5)

    # joint candidates = all non-control except the separable control
    joint = {k: v for k, v in non_control.items()
             if k != "J5_separable_control"}
    best_joint_L2 = min(v["L2_to_target"] for v in joint.values())
    sep_L2 = results["J5_separable_control"]["L2_to_target"]
    # separable underperforms iff its L2 is worse than the best joint kernel
    # AND it does not produce a valid peak that the joints lack
    sep_valid = results["J5_separable_control"]["valid_peak"]
    any_joint_valid = any(v["valid_peak"] for v in joint.values())
    separable_underperforms = (sep_L2 >= best_joint_L2) and not (
        sep_valid and not any_joint_valid)

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

    verdict = "HIT" if hit else "NULL_NINTH_CLASS"

    summary = {
        "addendum": "A305",
        "addresses": "OI-287-1",
        "eigenvalues": [round(float(v), 4) for v in lam],
        "node_counts": nu,
        "target_normalized": TARGET,
        "a299_base_diagonal": results["J0_control_one"]["diagonal_ratio"],
        "constants": {"OMEGA_alpha_inv": OMEGA, "KAPPA_alpha_5_4": KAPPA},
        "candidates": results,
        "n_candidates_declared": len(kernels),
        "n_noncontrol_candidates": len(non_control),
        "n_joint_candidates": len(joint),
        "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_valid,
        "best_joint_L2": best_joint_L2,
        "separable_control_L2": sep_L2,
        "separable_underperforms_joint": bool(separable_underperforms),
        "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("A305 PROBE -- OI-287-1 non-separable joint-kernel localizing factor")
    print("=" * 64)
    print("eigenvalues:", summary["eigenvalues"], " node counts:", nu)
    print("target (normalized):", TARGET)
    print("A299 base diagonal :", results["J0_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(kernels)} "
          f"(non-control: {len(non_control)}, joint: {len(joint)})")
    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: {any_valid}")
    print(f"best joint L2={best_joint_L2}  separable(J5) L2={sep_L2}")
    print(f"separable control underperforms joints: {separable_underperforms}")
    print(f"\nVERDICT: {verdict}")
    if verdict == "NULL_NINTH_CLASS":
        print("  -> ninth excluded mechanism class: non-separable joint kernel")
        print("     (no declared kernel both peaks and lands within 20%).")
    print(f"wrote {OUT}")


if __name__ == "__main__":
    main()
