"""Independent re-computation of the Vellora smoking and lung cancer study.

Recomputes every statistic in public/data/research/lung-smoking-results.v1.json
from the raw verified input files with SciPy and NumPy (a separate
implementation from src/lib/research/lung-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-lung-study.py
"""

import json
import math
import sys

import numpy as np
import scipy
from scipy import stats

load = lambda p: json.load(open(p, encoding="utf-8"))
places = load("public/data/places-tx-cardio.v2025.1.json")
cancer = load("public/data/scp-cancer-incidence-tx.2019-2023.v1.json")
facts = load("public/data/research/cardio-disparities-inputs.v1.json")
res = load("public/data/research/lung-smoking-results.v1.json")

order = places["countyOrder"]
assert order == cancer["countyOrder"] == facts["countyOrder"]
arr = lambda cell: np.array([np.nan if v is None else v for v in cell["values"]], dtype=float)
smoke = arr(places["observations"]["smoking-aa"]["places-2025|all"])
crude = arr(places["observations"]["smoking"]["places-2025|all"])
lung = arr(cancer["observations"]["cancer-inc-lung"]["scp-2019-2023|all"])
metro = np.array([c["metro"] for c in facts["counties"]])
pop = np.array([c["pop2020"] for c in facts["counties"]])
ok = ~np.isnan(smoke) & ~np.isnan(lung)

failures = []


def check(label, got, want, tol):
    if not math.isclose(got, want, rel_tol=0, abs_tol=tol):
        failures.append(f"{label}: results file {got} vs recomputed {want}")


def spearman(x, y, label, r):
    rho, p = stats.spearmanr(x, y)
    n = len(x)
    se = math.sqrt((1 + rho**2 / 2) / (n - 3))
    z = stats.norm.ppf(0.975)
    check(f"{label} n", r["n"], n, 0)
    check(f"{label} rho", r["rho"], rho, 5e-5)
    check(f"{label} ci lo", r["ci"][0], math.tanh(math.atanh(rho) - z * se), 5e-5)
    check(f"{label} ci hi", r["ci"][1], math.tanh(math.atanh(rho) + z * se), 5e-5)
    check(f"{label} p", r["p"], p, max(1e-12, 1e-6 * p))


x, y = smoke[ok], lung[ok]
check("included", res["counties"]["included"], int(ok.sum()), 0)
check("excluded", res["counties"]["excluded"], int((~ok).sum()), 0)
if [e["fips"] for e in res["excluded"]] != [f for f, k in zip(order, ok) if not k]:
    failures.append("excluded county list differs")
spearman(x, y, "primary", res["primary"])
for name, v in (("smoking", x), ("lung", y)):
    d = res["descriptives"][name]
    check(f"{name} mean", d["mean"], v.mean(), 5e-5)
    check(f"{name} sd", d["sd"], v.std(ddof=1), 5e-5)
    check(f"{name} median", d["median"], np.median(v), 5e-5)
    check(f"{name} q1", d["q1"], np.quantile(v, 0.25), 5e-5)
    check(f"{name} q3", d["q3"], np.quantile(v, 0.75), 5e-5)
    check(f"{name} min", d["min"], v.min(), 0)
    check(f"{name} max", d["max"], v.max(), 0)

e = res["exploratory"]
pr = stats.pearsonr(x, y)
ci = pr.confidence_interval(0.95)
check("pearson r", e["pearson"]["r"], pr.statistic, 5e-5)
check("pearson p", e["pearson"]["p"], pr.pvalue, max(1e-12, 1e-6 * pr.pvalue))
check("pearson ci lo", e["pearson"]["ci"][0], ci.low, 5e-5)
check("pearson ci hi", e["pearson"]["ci"][1], ci.high, 5e-5)
m = metro[ok]
spearman(x[m], y[m], "metro", e["metro"])
spearman(x[~m], y[~m], "nonmetro", e["nonmetro"])
okc = ok & ~np.isnan(crude)
spearman(crude[okc], lung[okc], "crude smoking", e["crudeSmoking"])
loo = [stats.spearmanr(np.delete(x, k), np.delete(y, k))[0] for k in range(len(x))]
check("leave-one-out min", e["leaveOneOut"]["min"]["rho"], min(loo), 5e-5)
check("leave-one-out max", e["leaveOneOut"]["max"]["rho"], max(loo), 5e-5)
ex = e["excludedProfile"]
check("median pop included", ex["medianPop2020Included"], np.median(pop[ok]), 1e-9)
check("median pop excluded", ex["medianPop2020Excluded"], np.median(pop[~ok]), 1e-9)
check("max pop excluded", ex["maxPop2020Excluded"], pop[~ok].max(), 0)
check("metro excluded", ex["metroExcluded"], int(metro[~ok].sum()), 0)
check("median smoking excluded", ex["medianSmokingExcluded"], np.nanmedian(smoke[~ok]), 5e-5)

if failures:
    print(f"LUNG STUDY CHECK FAILED ({len(failures)}):\n" + "\n".join(failures))
    sys.exit(1)
p = res["primary"]
print(f"OK: every statistic in the lung study results file reproduced with SciPy {scipy.__version__}: {p['n']} counties included, {res['counties']['excluded']} excluded; Spearman rho = {p['rho']:.4f} (95% CI {p['ci'][0]:.4f} to {p['ci'][1]:.4f}), p = {p['p']:.3g}.")