#!/usr/bin/env python3
"""probe_selection_eigenstructure.py — selection-rule eigenstructure probe.

OI-287-1 keystone test. Three coupling mechanisms are excluded (A292,
A295, A296), all of which perturbed the observation operator from
outside. This probes the un-tried alternative: do the eigenfunctions
psi_n of the UNPERTURBED operator already carry degree-n character, so
that index-matching (family n <-> degree-n monomial of rho, A296) is a
property of the bare spectrum rather than an inserted coupling?

Operator (matches verify_P296.py exactly):
  H = -d^2/dx^2 + V(x),   V = [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,   Dirichlet on [0,1].

Pre-registered tests (criteria fixed in the plan before running):
  T1 diagonal preference: A[n][k]=<psi_n|x^k|psi_n>, k=1..3,
     z-scored down each k-column; HIT if argmax_k(row n)==n for all n.
  T2 coefficient reproduction (strong): recover (2,3,16) ratios from
     the matrix elements with NO tuning; HIT if within 20% normalized.
  T3 control: same as T1 on the box psi_n=sqrt2 sin(n pi x); the
     potential must show a diagonal preference the box does not.
VERDICT: HIT only if T1 diagonal AND T3 control fails; else NULL =
fourth excluded mechanism class.

Copyright Léon Fernando Vlegels — MIT
"""
import json
import math
import os

import numpy as np

PI = math.pi
OM = 4 * PI**3 + PI**2 + PI
ES = 13.177
C_TARGET = (2.0, 3.0, 16.0)   # A296 degree coefficients
N = 1500
HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "..", "..", "..", "..", "Outputs",
                   "selection_rule_eigenstructure.json")


def eigfuns(potential):
    """First three Dirichlet eigenpairs of -d2/dx2 + potential on
    (0,1), returned as (eigvals[:3], psi[:, :3]) L2-normalised."""
    dx = 1.0 / N
    x = np.linspace(dx, 1 - dx, N - 1)
    V = potential(x)
    off = -1.0 / dx**2 * np.ones(N - 2)
    H = np.diag(2.0 / dx**2 + V) + np.diag(off, 1) + np.diag(off, -1)
    w, v = np.linalg.eigh(H)
    psi = v[:, :3] / math.sqrt(dx)        # L2 normalise
    return x, w[:3], psi


def matrix_elements(x, psi):
    """A[n][k] = <psi_n| x^k |psi_n>, n,k in {1,2,3}."""
    dx = x[1] - x[0]
    A = np.zeros((3, 3))
    for n in range(3):
        p2 = psi[:, n] ** 2
        for ki, k in enumerate((1, 2, 3)):
            A[n, ki] = float(np.sum(p2 * x**k) * dx)
    return A


def zscore_columns(A):
    """Standardise each k-column across the three levels."""
    Z = np.zeros_like(A)
    for k in range(3):
        col = A[:, k]
        sd = col.std()
        Z[:, k] = (col - col.mean()) / sd if sd > 0 else 0.0
    return Z


def diagonal_preference(A):
    Z = zscore_columns(A)
    argmax_per_row = [int(np.argmax(Z[n])) for n in range(3)]
    diagonal = all(argmax_per_row[n] == n for n in range(3))
    return diagonal, argmax_per_row, Z


def main():
    res = {"operator": "H=-d2/dx2 + [rho']^2/(2 Omega^2) + "
           "E_self x^2(1-x)^2, Dirichlet, N=%d" % N,
           "c_target": C_TARGET}

    # corpus operator
    rho_p = lambda x: 48 * PI**3 * x**2 + 6 * PI**2 * x + 2 * PI
    Vpot = lambda x: rho_p(x)**2 / (2 * OM**2) + ES * x**2 * (1 - x)**2
    x, evals, psi = eigfuns(Vpot)
    A = matrix_elements(x, psi)
    diag, argmax_rows, Z = diagonal_preference(A)
    res["eigenvalues"] = [float(e) for e in evals]
    res["A_matrix_xk"] = A.tolist()
    res["A_zscored"] = Z.tolist()
    res["T1_argmax_per_level"] = argmax_rows
    res["T1_diagonal_preference"] = bool(diag)

    # near-center leading polynomial degree (descriptive)
    near = x < 0.06
    leading = []
    for n in range(3):
        seg = psi[near, n]
        xs = x[near]
        # fit log|psi| vs log x slope -> leading power near 0
        m = np.abs(seg) > 0
        if m.sum() > 5:
            sl = np.polyfit(np.log(xs[m]), np.log(np.abs(seg[m])), 1)[0]
            leading.append(round(float(sl), 3))
        else:
            leading.append(None)
    res["near_center_leading_power"] = leading

    # T2: recover (2,3,16) ratios with no tuning. Two readings:
    #  (a) raw diagonal <psi_n|x^n|psi_n>
    #  (b) the diagonal element divided by the n=1 diagonal
    diag_xn = [float(A[n, n]) for n in range(3)]
    norm = [d / diag_xn[0] for d in diag_xn]
    tgt_norm = [c / C_TARGET[0] for c in C_TARGET]
    within20 = all(abs(norm[i] - tgt_norm[i]) <= 0.2 * tgt_norm[i]
                   for i in range(3))
    res["T2_diag_xn"] = diag_xn
    res["T2_normalised_ratio"] = norm
    res["T2_target_normalised"] = tgt_norm
    res["T2_within_20pct"] = bool(within20)

    # T3: box control
    xb = x
    psib = np.column_stack([math.sqrt(2) * np.sin((n + 1) * PI * xb)
                            for n in range(3)])
    Ab = matrix_elements(xb, psib)
    diagb, argmaxb, Zb = diagonal_preference(Ab)
    res["T3_box_A_matrix"] = Ab.tolist()
    res["T3_box_argmax_per_level"] = argmaxb
    res["T3_box_diagonal_preference"] = bool(diagb)

    hit = bool(diag and not diagb)
    res["VERDICT"] = ("HIT: index-matching in bare spectrum"
                      if hit else
                      "NULL: no index-matching in bare matrix elements "
                      "(fourth excluded mechanism class)")

    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    json.dump(res, open(OUT, "w"), indent=1)

    print("eigenvalues (levels 1,2,3):",
          [round(e, 4) for e in res["eigenvalues"]])
    print("\nA[n][k] = <psi_n | x^k | psi_n>  (rows n=1,2,3; cols k=1,2,3)")
    for n in range(3):
        print("  n=%d: %s" % (n + 1, ["%.4f" % v for v in A[n]]))
    print("z-scored down columns:")
    for n in range(3):
        print("  n=%d: %s  argmax=k%d"
              % (n + 1, ["%+.2f" % v for v in Z[n]], argmax_rows[n] + 1))
    print("\nT1 diagonal preference (corpus): ", res["T1_diagonal_preference"])
    print("T3 diagonal preference (box ctrl):", res["T3_box_diagonal_preference"])
    print("near-center leading power per level:", leading)
    print("T2 normalised diag ratio:", [round(v, 3) for v in norm],
          "vs target", [round(v, 3) for v in tgt_norm],
          "-> within20%:", within20)
    print("\nVERDICT:", res["VERDICT"])
    print("written:", os.path.normpath(OUT))


if __name__ == "__main__":
    main()
