"""Independent re-computation of the Vellora cardiovascular disparities study.

Recomputes every statistic in public/data/research/cardio-disparities-results.v1.json
from the raw input files with SciPy and NumPy (a separate implementation from
src/lib/research/cardio-study.ts) and exits non-zero on any mismatch.

Run from the repository root, in a throwaway environment:
    uvx --with scipy --with numpy python scripts/verify-cardio-study.py
"""

import json
import math
import sys

import numpy as np
from scipy import stats

load = lambda p: json.load(open(p, encoding="utf-8"))
inp = load("public/data/research/cardio-disparities-inputs.v1.json")
heart = load("public/data/cdc-heart-mortality-tx.2022-2024.v1.json")
counties = load("src/data/counties.json")
res = load("public/data/research/cardio-disparities-results.v1.json")

order = inp["countyOrder"]
assert order == heart["countyOrder"] == [c["fips"] for c in counties]
RISK = ["BPHIGH", "DIABETES", "OBESITY", "CSMOKING", "LPA"]
OUT = ["CHD", "STROKE", "HDMORT"]
VARS = RISK + OUT
metro = np.array([c["metro"] for c in inp["counties"]])


def series(v, kind="aa"):
    raw = heart["observations"]["hd-mort"]["hd-2022-2024|all"]["values"] if v == "HDMORT" else inp["measures"][v][kind]["values"]
    return np.array([np.nan if x is None else x for x in raw], dtype=float)


failures = []


def check(label, got, want, tol):
    if not (abs(got - want) <= tol):
        failures.append(f"{label}: results file {want} vs SciPy/NumPy {got}")


# Descriptives (NumPy percentile "linear" is the same type-7 definition).
for v in VARS:
    x = series(v)
    x = x[~np.isnan(x)]
    d = res["descriptives"][v]
    check(f"{v} n", len(x), d["n"], 0)
    for q, key in [(0.5, "median"), (0.25, "q1"), (0.75, "q3"), (0.1, "p10"), (0.9, "p90")]:
        check(f"{v} {key}", float(np.percentile(x, q * 100)), d[key], 1e-4)
    check(f"{v} mean", float(x.mean()), d["mean"], 1e-4)
    check(f"{v} sd", float(x.std(ddof=1)), d["sd"], 1e-4)

# Spearman correlations, two-sided p, Benjamini–Hochberg.
def corr_table(pairs, kind, rows):
    ps = []
    for (a, b), row in zip(pairs, rows):
        x, y = series(a, kind), series(b, kind)
        ok = ~np.isnan(x) & ~np.isnan(y)
        r = stats.spearmanr(x[ok], y[ok])
        check(f"{kind} {a}-{b} n", int(ok.sum()), row["n"], 0)
        check(f"{kind} {a}-{b} rho", float(r.statistic), row["rho"], 1e-4)
        check(f"{kind} {a}-{b} p", float(r.pvalue), row["p"], 1e-9 + 1e-6 * row["p"])
        ps.append(float(r.pvalue))
    adj = stats.false_discovery_control(ps, method="bh")
    for (a, b), row, q in zip(pairs, rows, adj):
        check(f"{kind} {a}-{b} pAdj", float(q), row["pAdj"], 1e-9 + 1e-6 * row["pAdj"])
        n, rho = row["n"], row["rho"]
        se = math.sqrt((1 + rho * rho / 2) / (n - 3))
        z = math.atanh(rho)
        check(f"{kind} {a}-{b} ci lo", math.tanh(z - 1.959963984540054 * se), row["ci"][0], 2e-4)
        check(f"{kind} {a}-{b} ci hi", math.tanh(z + 1.959963984540054 * se), row["ci"][1], 2e-4)


grid = [(a, b) for a in RISK for b in OUT]
corr_table(grid, "aa", res["correlations"])
corr_table(grid, "crude", res["correlationsCrude"])
corr_table([("CHD", "STROKE"), ("CHD", "HDMORT"), ("STROKE", "HDMORT")], "aa", res["outcomePairs"])

# Mann–Whitney U, nonmetro vs metro.
ps = []
for row in res["metroComparison"]:
    x = series(row["id"])
    non, met = x[~metro & ~np.isnan(x)], x[metro & ~np.isnan(x)]
    t = stats.mannwhitneyu(non, met, alternative="two-sided", method="asymptotic", use_continuity=True)
    check(f"MW {row['id']} U", float(t.statistic), row["U"], 1e-9)
    check(f"MW {row['id']} p", float(t.pvalue), row["p"], 1e-12 + 1e-6 * row["p"])
    check(f"MW {row['id']} median non", float(np.median(non)), row["medianNonmetro"], 1e-4)
    check(f"MW {row['id']} median metro", float(np.median(met)), row["medianMetro"], 1e-4)
    ps.append(float(t.pvalue))
for row, q in zip(res["metroComparison"], stats.false_discovery_control(ps, method="bh")):
    check(f"MW {row['id']} pAdj", float(q), row["pAdj"], 1e-12 + 1e-6 * row["pAdj"])

# Moran's I (row-standardised contiguity) and an independent permutation test.
nbrs = {c["fips"]: c["neighbors"] for c in counties}
rng = np.random.default_rng(20261022)
for v in VARS:
    x = series(v)
    ids = [f for f, val in zip(order, x) if not np.isnan(val)]
    idx = {f: i for i, f in enumerate(ids)}
    vals = np.array([val for val in x if not np.isnan(val)])
    W = np.zeros((len(ids), len(ids)))
    for f in ids:
        nb = [idx[n] for n in nbrs[f] if n in idx]
        W[idx[f], nb] = 1 / len(nb)
    z = vals - vals.mean()
    moran = lambda zz: float(zz @ W @ zz / (zz @ zz))
    I = moran(z)
    check(f"Moran {v} I", I, res["spatial"][v]["I"], 1e-4)
    perms = np.array([moran(rng.permutation(z)) for _ in range(9999)])
    p = (np.sum(perms >= I) + 1) / 10000
    if p > 0.001:
        failures.append(f"Moran {v}: SciPy-side permutation p = {p}, results file {res['spatial'][v]['p']}")

if failures:
    print("VERIFICATION FAILED")
    print("\n".join(failures))
    sys.exit(1)
print(f"VERIFIED: descriptives for {len(VARS)} measures, {len(grid) * 2 + 3} Spearman correlations with BH and CIs, "
      f"{len(res['metroComparison'])} Mann–Whitney tests with BH, {len(VARS)} Moran's I statistics and permutation tests "
      f"(SciPy {__import__('scipy').__version__}, NumPy {np.__version__}).")