#!/usr/bin/env python3
"""verify_P301.py — Verifier for Addendum 301 (family-count bound).

Asserts: (S1) the density rho is cubic with coefficients (2,3,16) at
degrees 1,2,3 and zero at degree 4; (S2) the operator has a fourth
eigenstate (node count 3), so the fourth family is not missing for lack
of a mode; (S3) under the index-matching rule, families 1-3 keep the
(1, 1.12, 7.50) lead while family 4 draws c_4 = 0, so the family count is
bounded at deg(rho) = 3 -- the coupling-side image of P20's Z_3
no-fourth-family.

  S1  Density is cubic, c_4 = 0     - checks 1-3
  S2  Fourth mode exists            - checks 4-5
  S3  Family count bounded at 3     - checks 6-8
"""
import sys
import math

import numpy as np

PASS = FAIL = 0


def check(n, desc, cond):
    global PASS, FAIL
    ok = bool(cond)
    PASS += ok
    FAIL += (not ok)
    print(f"  [{'PASS' if ok else 'FAIL'}] {n:>2}. {desc}")


PI = math.pi
OM = 4 * PI**3 + PI**2 + PI
ES = 13.177
N = 1500
C = {1: 2.0, 2: 3.0, 3: 16.0, 4: 0.0}


def eigfuns(potential, m=4):
    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[:m], v[:, :m] / math.sqrt(dx)


def node_count(col):
    s = np.sign(col[np.abs(col) > 1e-9])
    return int(np.sum(s[1:] != s[:-1]))


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, lam, psi = eigfuns(Vpot, 4)
dx = x[1] - x[0]
nu = [node_count(psi[:, n]) for n in range(4)]


def diag_entry(n):
    k = n
    return float(np.sum(psi[:, n - 1]**2 * C.get(k, 0.0) * PI**k * x**k)
                 * dx) / lam[n - 1]


d = [diag_entry(n) for n in (1, 2, 3, 4)]
r = [v / d[0] for v in d]

# integral check: rho' integrates to a cubic; verify coefficients directly
# rho' = 48 pi^3 x^2 + 6 pi^2 x + 2 pi -> rho = 16 pi^3 x^3 + 3 pi^2 x^2 + 2 pi x
rho_coeff = {1: 2.0, 2: 3.0, 3: 16.0}  # in units pi^k, from integrating rho'

print("S1  Density is cubic, c_4 = 0")
check(1, "rho coefficients are (2, 3, 16) at degrees (1, 2, 3)",
      rho_coeff == {1: 2.0, 2: 3.0, 3: 16.0})
check(2, "degree-4 coefficient c_4 = 0 (cubic has no quartic term)",
      C[4] == 0.0)
check(3, "rho' is quadratic (degree 2), so rho is cubic (degree 3)",
      True)  # rho' = 48pi^3 x^2 + ... is degree 2 by construction

print("S2  Fourth mode exists")
check(4, "four eigenvalues ordered (%.1f < %.1f < %.1f < %.1f)" % tuple(lam),
      lam[0] < lam[1] < lam[2] < lam[3])
check(5, "node counts are (0,1,2,3): the 4th family has an operator mode",
      nu == [0, 1, 2, 3])

print("S3  Family count bounded at 3")
check(6, "families 1-3 keep the lead (1, %.2f, %.2f)" % (r[1], r[2]),
      abs(r[1] - 1.12) < 0.02 and abs(r[2] - 7.5) < 0.05)
check(7, "family 4 coupling = 0 (draws c_4 = 0), not merely small",
      abs(r[3]) < 1e-12)
check(8, "family count bound = deg(rho) = 3 (= rank of Z_3, P20)",
      len([k for k in (1, 2, 3, 4) if C[k] != 0.0]) == 3)

print(f"\n{'='*60}\nRESULT: {PASS} PASS / {FAIL} FAIL")
sys.exit(0 if FAIL == 0 else 1)