#!/usr/bin/env python3
"""probe_selection_nodal.py — selection-rule nodal/derivative probe.

A297 excluded the position-moment reading and showed the families
differ by node count (Dirichlet linear onset). This tests the
surviving door: does a NODAL or DERIVATIVE pairing couple level n to
the degree-n monomial of rho? Four fixed functionals, no free knobs.

Pairings (each builds a 3x3 matrix M[n][k], n,k in {1,2,3}):
  P1 derivative overlap   <psi_n' | x^k | psi_n'>   (psi_n' has n lobes)
  P2 kinetic coupling     <psi_n | (-d2/dx2)(x^k psi_n)>  diag part
  P3 density-Taylor overlap  <psi_n | c_k pi^k x^k | psi_n> normalized
  P4 ladder transition    <psi_n | x^k | psi_1>  (off-diagonal to ground)

Criteria (from A297, pre-registered in the plan):
  T1 diagonal preference: z-score columns; argmax_k(row n)==n for all
     n, AND the box control sin(n pi x) does NOT satisfy it.
  T2 coefficient reproduction: normalized diagonal ratio within 20%
     of target (1, 1.5, 8) = (2,3,16)/2.
  Direction check: is the diagonal monotone INCREASING (the (2,3,16)
     direction the position moments got backwards)?
VERDICT per pairing; global NULL = fifth excluded mechanism class.

Copyright Léon Fernando Vlegels -- CC BY 4.0
"""
import json
import math
import os

import numpy as np

PI = math.pi
OM = 4 * PI**3 + PI**2 + PI
ES = 13.177
C = (2.0, 3.0, 16.0)
N = 1500
HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "..", "..", "..", "..", "Outputs",
                   "selection_rule_nodal.json")
TGT = [c / C[0] for c in C]   # (1, 1.5, 8)


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 zcols(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 Z


def diag_pref(M):
    Z = zcols(M)
    am = [int(np.argmax(Z[n])) for n in range(3)]
    return all(am[n] == n for n in range(3)), am


def evaluate(name, M, Mbox, res):
    diag, am = diag_pref(M)
    diagb, amb = diag_pref(Mbox)
    d = [float(M[n, n]) for n in range(3)]
    ratio = [v / d[0] if d[0] != 0 else float("nan") for v in d]
    within20 = all(abs(ratio[i] - TGT[i]) <= 0.2 * TGT[i]
                   for i in range(3)) if not any(
                       math.isnan(r) for r in ratio) else False
    increasing = (not any(math.isnan(r) for r in ratio)
                  and ratio[0] < ratio[1] < ratio[2])
    hit = bool(diag and not diagb)
    res[name] = {
        "M": M.tolist(), "argmax_per_level": am,
        "box_diagonal_preference": bool(diagb),
        "diagonal_preference": bool(diag),
        "diagonal_ratio": ratio, "target": TGT,
        "within_20pct": bool(within20),
        "diagonal_increasing": bool(increasing),
        "verdict": "HIT" if (hit and within20) else
                   ("partial: diagonal preference, ratio off"
                    if hit else "NULL"),
    }
    print("== %s" % name)
    print("   argmax %s (box %s) -> diag pref %s/%s; ratio "
          "(1, %.2f, %.2f) increasing=%s within20=%s -> %s"
          % (am, amb, diag, diagb, ratio[1], ratio[2],
             increasing, within20, res[name]["verdict"]))


def main():
    res = {"target_ratio": TGT, "pairings": {}}
    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)
    dx = x[1] - x[0]
    res["eigenvalues"] = [float(e) for e in evals]
    xb = x
    psib = np.column_stack([math.sqrt(2) * np.sin((n + 1) * PI * xb)
                            for n in range(3)])

    def deriv(p):
        d = np.gradient(p, dx, axis=0)
        return d

    P = res["pairings"]

    # P1 derivative overlap <psi_n'|x^k|psi_n'>
    dpsi = deriv(psi)
    dpsib = deriv(psib)
    M1 = np.array([[float(np.sum(dpsi[:, n]**2 * x**k) * dx)
                    for k in (1, 2, 3)] for n in range(3)])
    M1b = np.array([[float(np.sum(dpsib[:, n]**2 * xb**k) * dx)
                     for k in (1, 2, 3)] for n in range(3)])
    evaluate("P1_derivative_overlap", M1, M1b, P)

    # P2 kinetic coupling: diag of <psi_n| -d2/dx2 (x^k psi_n)>
    def kinetic(ps, n, k):
        f = x**k * ps[:, n]
        d2 = np.gradient(np.gradient(f, dx), dx)
        return float(np.sum(ps[:, n] * (-d2)) * dx)
    M2 = np.array([[kinetic(psi, n, k) for k in (1, 2, 3)]
                   for n in range(3)])
    M2b = np.array([[kinetic(psib, n, k) for k in (1, 2, 3)]
                    for n in range(3)])
    evaluate("P2_kinetic_coupling", M2, M2b, P)

    # P3 density-Taylor overlap: <psi_n| c_k pi^k x^k |psi_n>
    ck_pik = [C[k - 1] * PI**k for k in (1, 2, 3)]
    M3 = np.array([[float(np.sum(psi[:, n]**2 * ck_pik[ki] * x**(ki + 1))
                          * dx) for ki in range(3)] for n in range(3)])
    M3b = np.array([[float(np.sum(psib[:, n]**2 * ck_pik[ki]
                                  * xb**(ki + 1)) * dx)
                     for ki in range(3)] for n in range(3)])
    evaluate("P3_density_taylor_overlap", M3, M3b, P)

    # P4 ladder transition <psi_n|x^k|psi_1>
    M4 = np.array([[float(np.sum(psi[:, n] * x**k * psi[:, 0]) * dx)
                    for k in (1, 2, 3)] for n in range(3)])
    M4b = np.array([[float(np.sum(psib[:, n] * xb**k * psib[:, 0]) * dx)
                     for k in (1, 2, 3)] for n in range(3)])
    evaluate("P4_ladder_transition", M4, M4b, P)

    hits = [k for k, v in P.items() if v["verdict"] == "HIT"]
    partials = [k for k, v in P.items()
                if v["verdict"].startswith("partial")]
    increasing = [k for k, v in P.items() if v["diagonal_increasing"]]
    res["VERDICT"] = (
        "HIT: " + ", ".join(hits) if hits else
        "NULL: no pairing reproduces index-matching (fifth excluded "
        "mechanism class). partial=%s; right-direction(increasing "
        "diagonal)=%s" % (partials or "none", increasing or "none"))

    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    json.dump(res, open(OUT, "w"), indent=1)
    print("\nVERDICT:", res["VERDICT"])
    print("written:", os.path.normpath(OUT))


if __name__ == "__main__":
    main()
