#!/usr/bin/env python3
"""verify_constraint_end_checks.py — the CONSTRAINT-END (bottom-knot) power law.

Pre-registered verification battery for the constraint-end results (the mirror of the
high-wealth q-theorem), stated in `statement.md` §5 and proved in
`constraint_end_proof.md`. Notation: the constraint-end approach exponent is q° (written
`q_circ` here and `q_\\circ` in the math — the ring-subscript partner of the high-wealth q*):

  (C1)  worst-atom limiting MPC          kbar = 1 - wp^(1/rho)*ThornR         (BST eq:MPCmaxDef)
  (C2)  THE NEW THEOREM: approach exp.   c(m)=kbar*m^e - K*(m^e)^(1+q_circ)(1+o(1)),
        i.e. the MPC deficit  kbar - c/m^e ~ K*(m^e)^q_circ,  q° = rho             (st-thm-CE)
  (C3)  finite-horizon recursion         kbar_{T-1} = 1/(1+MPSmin),  MPSmin=wp^(1/rho)ThornR
  (C4)  grid-design corollary            bottom knot is inside the constraint asymptote to
        tolerance tol iff  m^e_0 <= (tol*kbar/K)^(1/rho)  (mirror of the q*hEx top rule)

PRE-REGISTRATION (gates frozen here BEFORE this battery first ran; never weakened):
  The PREDICTION is exact: q° = rho, kbar = 1 - wp^(1/rho)*ThornR, kbar_{T-1}=1/(1+MPSmin).
  Gates are the derivation's exact values plus a tolerance for (a) income discretization,
  (b) the finite EGM asset grid, (c) the free-exponent nonlinear-fit noise. Bands:
     GATE_KBAR    = 1e-3   # |fit intercept - kbar_formula|   (measured ~1e-8; also c/m^e monotone up)
     GATE_QC      = 0.06   # |free-fit q_circ - rho|          (measured 0.01-0.04 on the robust window)
     GATE_QC_REL  = 0.03   # |q_circ/rho - 1|                 (redundant relative guard)
     GATE_T1      = 1e-3   # |mpc_{T-1} - 1/(1+MPSmin)|        (measured ~1e-7)
     GATE_GRID    = 0.20   # relative error of the predicted vs measured inside-asymptote m^e_0
  Measurement window for the free-exponent fit: m^e in [WIN_LO, WIN_HI], below the
  bottom<->top crossover and above the near-constraint numerical floor, gated by the
  per-point Euler residual (EULER_GATE) — the same guard the top-side batteries use.

Env: pure numpy/scipy longdouble reference EGM (`review/R4_egm_lib.py`); no HARK needed.
Run:  python3 verify_constraint_end_checks.py   (writes verify_constraint_end_checks_out.txt)
"""
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "review"))
import numpy as np
import R4_egm_lib as A
from scipy.optimize import curve_fit

LD = np.longdouble

# ------------------------------------------------------------------ frozen gates
GATE_KBAR   = 1e-3
GATE_QC     = 0.06
GATE_QC_REL = 0.03
GATE_T1     = 1e-3
GATE_GRID   = 0.20
WIN_LO, WIN_HI = 1e-4, 2e-2
EULER_GATE = 1e-6

# Calibrations: zero-income NATURAL-borrowing case (true theta=0 atom => m_underbar=0,
# m^e=m, constraint-end MPC = kbar<1). rho is varied to test that q° tracks rho (the
# decisive check), across GIC-satisfying (R,G,beta,wp). Deep near-constraint grid.
CALS = [
    dict(tag="NB-rho2",     R=1.04, G=1.00, beta=0.96, rho=2.0, wp=0.05),
    dict(tag="NB-rho1.5",   R=1.04, G=1.00, beta=0.96, rho=1.5, wp=0.05),
    dict(tag="NB-rho3",     R=1.04, G=1.00, beta=0.96, rho=3.0, wp=0.05),
    dict(tag="NB-rho2-wp10",R=1.03, G=1.00, beta=0.95, rho=2.0, wp=0.10),
    dict(tag="NB-rho2.5",   R=1.05, G=1.01, beta=0.94, rho=2.5, wp=0.07),
]
SOLVE = dict(N=7, sig=0.10, aMin=1e-10, aMax=1e6, Na=9000, tol=1e-16, maxit=30000)

OUT = []; P = OUT.append; failures = []
def check(label, ok):
    P(f"  [{'PASS' if ok else 'FAIL'}] {label}")
    if not ok:
        failures.append(label)

def kbar_formula(d, wp):
    return float(1 - LD(wp) ** (1 / d["rho"]) * d["PR"])

def solve_cal(cal):
    th, p = A.with_unemp_atom(SOLVE["sig"], SOLVE["N"], unemp_prob=cal["wp"], inc_unemp=0.0)
    d = A.solve(cal["R"], cal["G"], cal["beta"], cal["rho"], th, p,
                aMin=SOLVE["aMin"], aMax=SOLVE["aMax"], Na=SOLVE["Na"],
                tol=SOLVE["tol"], maxit=SOLVE["maxit"])
    return d, th, p

def one_step_Tm1_mpc(cal, th, p, a_small=1e-9):
    """One backward EGM step from terminal c_T(m)=m; MPC c_{T-1}/m^e as a->0."""
    d = A.derived(cal["R"], cal["G"], cal["beta"], cal["rho"])
    disc = d["beta"] * d["R"] * d["G"] ** (-d["rho"])
    a = LD(a_small)
    mp = d["RG"] * a + th                 # terminal m'; worst atom th=0 -> mp=RG a
    with np.errstate(divide="ignore"):
        EvP = (np.where(mp > 0, mp, np.inf) ** (-d["rho"]) * p).sum()
    c = (disc * EvP) ** (-1 / d["rho"])
    return float(c / (a + c))

def main():
    P("=" * 78)
    P("verify_constraint_end_checks.py — constraint-end power law (q_circ = rho)")
    P("=" * 78)
    P(f"gates: |q_circ-rho|<={GATE_QC} (rel<={GATE_QC_REL}), |kbar-fit|<={GATE_KBAR:.0e}, "
      f"|T-1|<={GATE_T1:.0e}, grid<={GATE_GRID}; window m^e in [{WIN_LO:.0e},{WIN_HI:.0e}]")
    P("")
    for cal in CALS:
        d, th, p = solve_cal(cal)
        rho = float(cal["rho"]); wp = cal["wp"]
        kbar = kbar_formula(d, wp)
        MPSmin = float(LD(wp) ** (1 / d["rho"]) * d["PR"])
        m = np.asarray(d["m"], dtype=float); c = np.asarray(d["c"], dtype=float)
        me = m                                   # m_underbar = 0
        mpc = c / me
        er = np.abs(A.euler_resid_rel(d, d["m"]))
        P(f"--- {cal['tag']}: R={cal['R']} G={cal['G']} beta={cal['beta']} rho={rho} wp={wp}"
          f"  | kbar={kbar:.6f}  q*={d['qs']:.2f}  PG={float(d['PG']):.5f}")

        # (C1) & (C2): free-exponent fit  mpc = kb - K*me^q_circ  on the clean window
        sel = (me > WIN_LO) & (me < WIN_HI) & (er < EULER_GATE) & (mpc > 0) & (mpc < kbar)
        x, y = me[sel], mpc[sel]
        def f(x, kb, K, q_circ): return kb - K * np.power(x, q_circ)
        popt, _ = curve_fit(f, x, y, p0=[kbar, 1.0, rho], maxfev=40000)
        kb, K, q_circ = popt
        P(f"      n_fit={sel.sum():4d}  free-fit: q_circ={q_circ:.4f}  K={K:.4g}  kb_fit={kb:.6f}")
        check(f"C1[{cal['tag']}] kbar: |kb_fit-formula|={abs(kb-kbar):.2e} (kbar={kbar:.6f})",
              abs(kb - kbar) <= GATE_KBAR)
        check(f"C1b[{cal['tag']}] MPC rises to kbar from below (c/m^e < kbar on window)",
              bool(np.all(y < kbar + 1e-9)))
        check(f"C2[{cal['tag']}] q_circ=rho: q_circ={q_circ:.4f} rho={rho} |q_circ-rho|={abs(q_circ-rho):.4f}",
              abs(q_circ - rho) <= GATE_QC)
        check(f"C2rel[{cal['tag']}] |q_circ/rho-1|={abs(q_circ/rho-1):.4f}", abs(q_circ/rho - 1) <= GATE_QC_REL)

        # (C3) finite-horizon T-1 recursion
        mpc_T1 = one_step_Tm1_mpc(cal, th, p)
        pred_T1 = 1.0 / (1.0 + MPSmin)
        check(f"C3[{cal['tag']}] T-1: c/m^e={mpc_T1:.6f} vs 1/(1+MPSmin)={pred_T1:.6f} "
              f"|d|={abs(mpc_T1-pred_T1):.2e}", abs(mpc_T1 - pred_T1) <= GATE_T1)

        # (C4) grid-design: predict m^e_0 giving relative deviation = tol from the kbar line,
        # then verify the measured deviation there matches tol.  dev(me) = (K/kbar)*me^rho.
        tol = 1e-2
        me0_pred = (tol * kbar / K) ** (1.0 / rho)
        c0 = float(d["eval_c"](LD(me0_pred)))
        dev_meas = (kbar * me0_pred - c0) / (kbar * me0_pred)
        check(f"C4[{cal['tag']}] grid-design: at m^e0={me0_pred:.4g} predicted dev={tol:.0e}, "
              f"measured={dev_meas:.3e} (rel err {abs(dev_meas-tol)/tol:.2f})",
              abs(dev_meas - tol) / tol <= GATE_GRID)
        P("")

    P("-" * 78)
    P(f"RESULT: {'ALL PASS' if not failures else 'FAILURES: ' + str(failures)}")
    text = "\n".join(OUT)
    print(text)
    with open(__file__.replace(".py", "_out.txt"), "w") as fh:
        fh.write(text + "\n")
    return 0 if not failures else 1

if __name__ == "__main__":
    raise SystemExit(main())