#!/usr/bin/env python
"""Numerical verification of the closed forms of Propositions 1-5 (paper/sections/theory.tex).

Monte Carlo with N = 50 traders and R = 200,000 draws of the one-period model
    s_i = v + eta_{g(i)} + u_i,   x_i = beta s_i,   Y = sum_i x_i + z,   P = lambda Y.
Every row of results/tables/theory_checks.csv compares a closed form with its simulated
counterpart and reports the relative error; the script exits with a non-zero code if any
relative error exceeds 2 %.

Run:  .venv/bin/python scripts/90_theory_checks.py
"""
from __future__ import annotations

import math
import sys
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "results" / "tables" / "theory_checks.csv"

# ----------------------------------------------------------------------------- primitives
N = 50               # traders
R = 200_000          # Monte Carlo draws
SV, SE, SU, SZ = 1.0, 1.0, 0.8, 5.0   # sigma_v, sigma_eta, sigma_u, sigma_z
BETA = 1.0
TOL = 0.02
rng = np.random.default_rng(20261027)

SS2 = SV**2 + SE**2 + SU**2           # signal variance sigma_s^2
RHO0 = SV**2 / SS2                    # different / private models
RHO1 = (SV**2 + SE**2) / SS2          # same model

rows: list[dict] = []


def add(prop: str, case: str, quantity: str, theory: float, sim: float) -> None:
    rel = abs(sim - theory) / abs(theory) if theory != 0 else abs(sim - theory)
    rows.append(dict(proposition=prop, case=case, quantity=quantity, theory=theory,
                     simulation=sim, rel_error=rel))


# ----------------------------------------------------------------------------- closed forms
def omega(H: float, Phi: float) -> float:
    return SE**2 * (H + (1.0 - Phi) / N) + SU**2 / N + SZ**2 / (BETA**2 * N**2)


def closed_forms(sizes: list[int]) -> dict:
    Phi = sum(sizes) / N
    H = sum((n / N) ** 2 for n in sizes)
    Om = omega(H, Phi)
    tau = SV**2 / (SV**2 + Om)
    lam = tau / (BETA * N)
    w = (N * H - Phi) / (N - 1)                     # share of same-model ordered pairs
    rhobar = RHO0 + w * SE**2 / SS2
    sync = 0.5 + ((1 - w) * math.asin(RHO0) + w * math.asin(RHO1)) / math.pi
    return dict(H=H, Phi=Phi, Omega=Om, tau=tau, lam=lam, w=w, rhobar=rhobar, sync=sync,
                mse=SV**2 * (1 - tau))


# ----------------------------------------------------------------------------- simulation
def simulate(sizes: list[int]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Draw (v, x, Y) for a population with shared-model cohorts of the given sizes
    (first cohort occupies traders 0..n_1-1, etc.); remaining traders use private models."""
    v = rng.normal(0.0, SV, R)
    eps = rng.normal(0.0, SU, (R, N))
    start = 0
    for n in sizes:
        eps[:, start:start + n] += rng.normal(0.0, SE, (R, 1))
        start += n
    if start < N:
        eps[:, start:] += rng.normal(0.0, SE, (R, N - start))
    x = BETA * (v[:, None] + eps)
    Y = x.sum(axis=1) + rng.normal(0.0, SZ, R)
    return v, x, Y


def flow_stats(v: np.ndarray, x: np.ndarray, Y: np.ndarray) -> dict:
    cov = np.cov(v, Y, ddof=0)
    lam_hat = cov[0, 1] / cov[1, 1]
    tau_hat = cov[0, 1] ** 2 / (cov[0, 0] * cov[1, 1])
    P = lam_hat * Y
    mse = np.var(P - v)
    C = np.corrcoef(x, rowvar=False)
    off = ~np.eye(N, dtype=bool)
    rhobar = C[off].mean()
    sg = np.sign(x)
    A = (sg.T @ sg) / R                           # E[sign_i sign_j]
    sync = ((1.0 + A[off]) / 2.0).mean()
    return dict(tau=tau_hat, lam=lam_hat, mse=mse, rhobar=rhobar, sync=sync)


def check_population(prop: str, label: str, sizes: list[int], what=("tau", "lam", "mse", "rhobar", "sync")):
    th = closed_forms(sizes)
    v, x, Y = simulate(sizes)
    sim = flow_stats(v, x, Y)
    for q in what:
        add(prop, label, q, th[q], sim[q])
    return th, sim, (v, x, Y)


# ============================================================================= Props 1-2
print(f"Monte Carlo: N={N}, R={R:,}, sigma_v={SV}, sigma_eta={SE}, sigma_u={SU}, sigma_z={SZ}, beta={BETA}")
single_cases = {"phi=0": [], "phi=0.1": [5], "phi=0.2": [10], "phi=0.5": [25], "phi=1": [50]}
taus = {}
keep = None
for label, sizes in single_cases.items():
    th, sim, draws = check_population("P1-P2", f"single model, {label}", sizes)
    taus[label] = (th["tau"], sim["tau"])
    if label == "phi=0.2":
        keep = draws

# Prop 1(i): monotonicity across the grid (theory) and in simulation
tau_th = [taus[k][0] for k in single_cases]
tau_mc = [taus[k][1] for k in single_cases]
add("P1(i)", "tau strictly decreasing over phi grid (1=yes)", "monotone",
    1.0, float(all(np.diff(tau_th) < 0) and all(np.diff(tau_mc) < 0)))
# Prop 1(i): phi=0 and phi=1/N coincide in Omega
add("P1(i)", "Omega(phi=1/N) / Omega(phi=0)", "ratio", 1.0, omega((1 / N) ** 2, 1 / N) / omega(0.0, 0.0))
# Prop 1(ii): uniform bound tau(phi) <= tau_inf(phi) and N->infinity limit (theory, large N)
for phi in (0.2, 0.5):
    tau_inf = SV**2 / (SV**2 + phi**2 * SE**2)
    OmN = lambda n: SE**2 * (phi**2 + (1 - phi) / n) + SU**2 / n + SZ**2 / (BETA**2 * n**2)
    tau_big = SV**2 / (SV**2 + OmN(10**7))
    add("P1(ii)", f"N->inf limit, phi={phi}", "tau_inf", tau_inf, tau_big)
    add("P1(ii)", f"bound tau(N=50)<=tau_inf, phi={phi} (1=yes)", "bound", 1.0,
        float(SV**2 / (SV**2 + OmN(N)) <= tau_inf))
add("P1(iii)", "phi=0, N=1e7", "tau->1", 1.0,
    SV**2 / (SV**2 + (SE**2 + SU**2) / 1e7 + SZ**2 / (BETA**2 * 1e14)))

# Prop 2(ii): same-model vs cross-model pair correlations in the phi=0.5 population
v, x, Y = simulate([25])
C = np.corrcoef(x, rowvar=False)
same = C[:25, :25][~np.eye(25, dtype=bool)].mean()
cross = C[:25, 25:].mean()
add("P2(ii)", "phi=0.5", "rho_same", RHO1, same)
add("P2(ii)", "phi=0.5", "rho_cross", RHO0, cross)
add("P2(iv)", "phi=0.5, same-model pair", "S(rho1)=1/2+asin(rho1)/pi", 0.5 + math.asin(RHO1) / math.pi,
    ((1 + (np.sign(x[:, :25]).T @ np.sign(x[:, :25]) / R)[~np.eye(25, dtype=bool)]) / 2).mean())
# Prop 2 placebo: sigma_eta = 0 -> no dependence on phi (theory)
add("P2 placebo", "sigma_eta=0: Omega(phi=1)/Omega(phi=0)", "ratio", 1.0,
    (SU**2 / N + SZ**2 / (BETA**2 * N**2)) / (SU**2 / N + SZ**2 / (BETA**2 * N**2)))
del v, x, Y, C

# ============================================================================= Prop 3
div_cases = {
    "Phi=1, G=1 (50)": [50],
    "Phi=1, G=2 uniform (25,25)": [25, 25],
    "Phi=1, G=5 uniform (10x5)": [10] * 5,
    "Phi=1, G=3 skewed (30,10,10)": [30, 10, 10],
    "Phi=0.6, G=1 (30)": [30],
    "Phi=0.6, G=2 uniform (15,15)": [15, 15],
    "Phi=0.6, G=3 uniform (10,10,10)": [10, 10, 10],
}
div_tau = {}
for label, sizes in div_cases.items():
    th, sim, _ = check_population("P3", label, sizes, what=("tau", "rhobar", "sync"))
    div_tau[label] = (th["H"], th["tau"], sim["tau"])
# simulated tau is decreasing in H within each Phi group (hence single <= mixed <= uniform)
ok = True
for Phi_label in ("Phi=1", "Phi=0.6"):
    grp = sorted((H, tau_s) for lbl, (H, _, tau_s) in div_tau.items() if lbl.startswith(Phi_label + ","))
    ok = ok and all(grp[k][1] > grp[k + 1][1] for k in range(len(grp) - 1))   # H ascending -> tau descending
add("P3(iii)", "simulated tau decreasing in H within each Phi (1=yes)", "ordering", 1.0, float(ok))
# H bounds
for label, sizes in div_cases.items():
    H, Phi, G = sum((n / N) ** 2 for n in sizes), sum(sizes) / N, len(sizes)
    add("P3(ii)", f"Phi^2/G <= H <= Phi^2 for {label} (1=yes)", "bounds", 1.0,
        float(Phi**2 / G - 1e-12 <= H <= Phi**2 + 1e-12))

# ============================================================================= Prop 4
v, x, Y = keep
th = closed_forms([10])
lam = th["lam"]
for lg in (0.3, 0.6):
    gamma = lg / lam
    dP = lam * Y.copy()                       # start of the iterated response
    for _ in range(400):
        dP = lam * (Y + gamma * dP)
    amp = np.var(dP) / np.var(lam * Y)
    add("P4(i)", f"lambda*gamma={lg}", "variance amplification 1/(1-lg)^2", 1 / (1 - lg) ** 2, amp)
    kappa = 1 / (1 - lg)
    add("P4(iii)", f"lambda*gamma={lg}", "Var(P1-v)=sv^2[(1-tau)+tau(kappa-1)^2]",
        SV**2 * ((1 - th["tau"]) + th["tau"] * (kappa - 1) ** 2), np.var(dP - v))
    # fixed point = kappa*lambda*Y
    add("P4(i)", f"lambda*gamma={lg}", "max|dP - kappa*lam*Y| (abs)", 0.0,
        float(np.max(np.abs(dP - kappa * lam * Y))))
# divergence when lambda*gamma >= 1
for lg in (1.0, 1.2):
    gamma = lg / lam
    dP = lam * Y[:1000].copy()
    for _ in range(60):
        dP = lam * (Y[:1000] + gamma * dP)
    # lambda*gamma = 1 diverges linearly (|dP| = 61|lam Y| after 60 rounds), > 1 geometrically
    add("P4(ii)", f"lambda*gamma={lg}: |dP| after 60 rounds >= 50 * |lam*Y| (1=yes)", "divergence",
        1.0, float(np.all(np.abs(dP) >= 50 * np.abs(lam * Y[:1000]))))
del v, x, Y

# ============================================================================= Prop 5
# (ii)-(iii): cascade fixed point and multiplier
def cascade(a: float, dP0: float, b: float, X: float, lam_: float, m: float, rounds: int = 2000) -> float:
    dP = dP0
    for _ in range(rounds):
        F = max(0.0, (-X * dP - b) / m)
        dP = dP0 - lam_ * F
        if dP < -1e6:
            return dP
    return dP


lam_, m, X = 0.02, 1.0, 30.0         # a = 0.6
a = lam_ * X / m
for b in (0.0, 5.0):
    dP0 = -1.0
    fp = cascade(a, dP0, b, X, lam_, m)
    add("P5(ii)", f"a={a:.2f}, b={b}, dP0={dP0}", "fixed point (dP0+lam*b/m)/(1-a)",
        (dP0 + lam_ * b / m) / (1 - a), fp)
add("P5(ii)", f"a={a:.2f}, b=0", "multiplier 1/(1-a)", 1 / (1 - a), cascade(a, -1.0, 0.0, X, lam_, m) / -1.0)
add("P5(i)", f"a={a:.2f}, b=5, dP0=-0.1 (inside buffer)", "dP = dP0", -0.1, cascade(a, -0.1, 5.0, X, lam_, m))
for Xs in (50.0, 60.0):                      # a = 1.0, 1.2
    fp = cascade(lam_ * Xs / m, -1.0, 0.0, Xs, lam_, m, rounds=2000)   # a=1: linear, -2001 after 2000 rounds
    add("P5(iii)", f"a={lam_ * Xs / m:.2f}: cascade below -1e3 (1=yes)", "no finite fixed point", 1.0,
        float(fp < -1e3))

# (iv): distribution of the cohort net position X = |sum_{i in L} x_i|, N_L = 10 leveraged traders
NL = 10
c_thr = None
for phiL, sizes in (("phi_L=0", []), ("phi_L=0.5", [5]), ("phi_L=1", [10])):
    v = rng.normal(0.0, SV, R)
    eps = rng.normal(0.0, SU, (R, NL))
    if sizes:
        eps[:, :sizes[0]] += rng.normal(0.0, SE, (R, 1))
        eps[:, sizes[0]:] += rng.normal(0.0, SE, (R, NL - sizes[0]))
    else:
        eps += rng.normal(0.0, SE, (R, NL))
    xL = BETA * (v[:, None] + eps)
    Xnet = np.abs(xL.sum(axis=1))
    nL = sizes[0] if sizes else 0
    wL = nL * (nL - 1) / (NL * (NL - 1))
    rhoL = RHO0 + wL * SE**2 / SS2
    sigX = BETA * math.sqrt(SS2) * math.sqrt(NL * (1 + (NL - 1) * rhoL))
    if c_thr is None:
        c_thr = 1.0 * sigX                       # fixed spiral threshold m/lambda across cases
    add("P5(iv)", phiL, "sd(X_net)", sigX, float(np.std(xL.sum(axis=1))))
    add("P5(iv)", phiL, "E|X| = sigX*sqrt(2/pi)", sigX * math.sqrt(2 / math.pi), float(Xnet.mean()))
    p_th = 2 * (1 - 0.5 * (1 + math.erf(c_thr / sigX / math.sqrt(2))))
    add("P5(iv)", f"{phiL}, threshold m/lambda={c_thr:.3f}", "P(X>=m/lambda)=2*Phibar", p_th,
        float((Xnet >= c_thr).mean()))
    del v, eps, xL, Xnet

# ============================================================================= report
df = pd.DataFrame(rows)
OUT.parent.mkdir(parents=True, exist_ok=True)
df.to_csv(OUT, index=False)
pd.set_option("display.width", 200, "display.max_rows", 500, "display.max_colwidth", 60)
print(df.to_string(index=False, float_format=lambda f: f"{f:.5f}"))
worst = df["rel_error"].max()
print(f"\n{len(df)} checks written to {OUT.relative_to(ROOT)}; max relative error = {worst:.4%}")
if worst >= TOL:
    print(f"FAIL: at least one relative error >= {TOL:.0%}")
    sys.exit(1)
print(f"OK: all relative errors < {TOL:.0%}")