#!/usr/bin/env python3
"""Compact reproducibility checks for Online Resource 1.

This script verifies the main manuscript-level numerical claims from the
flat JSON file `ESM_2_core_data_results.json`. It is intentionally compact
for journal submission systems that unpack ZIP archives into many files.

Run:
    python ESM_3_core_reproducibility_check.py
"""

from __future__ import annotations

import json
import math
from pathlib import Path


DATA = Path(__file__).with_name("ESM_2_core_data_results.json")


def as_bool(value):
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        return value.lower() == "true"
    return bool(value)


def approx(value, target, tol=1e-6):
    return abs(float(value) - float(target)) <= tol


def row(rows, **where):
    for item in rows:
        ok = True
        for key, expected in where.items():
            value = item.get(key)
            if isinstance(expected, float):
                ok = ok and approx(value, expected, 1e-9)
            else:
                ok = ok and value == expected
        if ok:
            return item
    raise AssertionError(f"missing row: {where}")


def check_data(payload):
    reports = payload["data"]["reports_geocoded_real"]
    n_zhengzhou = sum(as_bool(r.get("in_zhengzhou")) for r in reports)
    assert len(reports) == 175
    assert n_zhengzhou == 52

    priors = payload["data"]["credibility_priors"]
    for source in ["eyewitness", "relay", "low_cred"]:
        assert source in priors
        assert "alpha_mean" in priors[source]
        assert "beta_mean" in priors[source]
    print(f"PASS data: {len(reports)} de-identified reports; {n_zhengzhou} Zhengzhou reports; credibility priors present")


def check_synthetic(payload):
    tables = payload["results"]

    e1 = tables["results__E1_collapse"]
    nr0 = row(e1, kind="rumor", tau=0, policy="nonrobust")["regret_mean"]
    nr5 = row(e1, kind="rumor", tau=0.5, policy="nonrobust")["regret_mean"]
    robust_vals = [
        r["regret_mean"] for r in e1
        if r["kind"] == "rumor" and r["policy"] == "robust"
    ]
    assert approx(nr0, 142.1, 0.05)
    assert approx(nr5, 297.18, 0.05)
    assert round(min(robust_vals), 2) == 105.29
    assert round(max(robust_vals), 2) == 168.19

    e3 = tables["results__E3_spillover"]
    spill = row(e3, mode="verify_plus_calibration")["reports_cleaned_mean"]
    assert approx(spill, 67.6, 0.05)

    e2 = tables["results__E2_threshold"]
    road_01 = row(e2, tau=0.1, strategy="verify_road")["regret_mean"]
    cal_01 = row(e2, tau=0.1, strategy="calibrate_source")["regret_mean"]
    road_02 = row(e2, tau=0.2, strategy="verify_road")["regret_mean"]
    cal_02 = row(e2, tau=0.2, strategy="calibrate_source")["regret_mean"]
    assert road_01 < cal_01
    assert cal_02 < road_02

    print("PASS synthetic: non-robust regret 142.1 -> 297.18; robust range 105.29-168.19; threshold switch bracketed")


def check_realcase(payload):
    results = payload["results"]["realcase_out__realcase_results"]
    assert results["n_reports"] == 52
    assert results["occupied_cells"] == 22
    assert approx(results["calibration"]["reports_cleaned"], 59.0, 1e-9)

    by_eta = {item["eta"]: item for item in results["by_eta"]}
    assert approx(by_eta[0.5]["nonrobust"]["dev"], 61.875, 1e-9)
    assert approx(by_eta[0.5]["robust"]["dev"], 0.0, 1e-9)

    tail = payload["results"]["realcase_out__realcase_tail"]
    assert approx(row(tail, policy="robust")["shift_CVaR20"], 0.0, 1e-9)
    print("PASS Zhengzhou replay: 52 reports / 22 cells; non-robust footprint about 62 cells; calibration cleans 59 reports")


def check_realnet(payload):
    results = payload["results"]["realnet_out__realnet_results"]
    regret = results["regret"]
    assert regret["robust (proposed)"] == [0.0] * 6
    assert [round(x) for x in regret["non-robust"]] == [0, 0, 33, 175, 267, 300]

    verify = results["verify"]
    nonrobust_no_verify = verify["non-robust"]["no-verify"]
    proposed_no_verify = verify["robust (proposed)"]["no-verify"]
    assert approx(nonrobust_no_verify, 283.3333333333333, 1e-6)
    assert approx(proposed_no_verify, 0.0, 1e-9)

    sensitivity = payload["results"]["realnet_out__realnet_sensitivity"]
    assert len(sensitivity) == 14
    assert all(approx(r["proposed_regret@0.4"], 0.0, 1e-9) for r in sensitivity)
    nonrobust_positive = sum(float(r["nonrobust_regret@0.4"]) > 0 for r in sensitivity)
    assert nonrobust_positive == 11

    print("PASS real-network: proposed regret is zero over eta grid and all 14 sensitivity configurations")


def sigmoid(x):
    return 1.0 / (1.0 + math.exp(-x))


def logit(p):
    return math.log(p / (1.0 - p))


def tau_star(kappa=1.0, nc=12.0, kbar=1.0, slope=18.0, tau0=0.30):
    return tau0 + (1.0 / slope) * logit(kappa / (nc * kbar + kappa))


def check_theory_anchors():
    tau = tau_star()
    assert 0.15 <= tau <= 0.17

    spill = [n * sigmoid(14.0 * (0.40 - 0.30)) for n in [5, 10, 20, 40, 60]]
    slopes = [(spill[i + 1] - spill[i]) / ([5, 10, 20, 40, 60][i + 1] - [5, 10, 20, 40, 60][i]) for i in range(4)]
    assert max(slopes) - min(slopes) < 1e-12
    print(f"PASS theory anchors: closed-form tau*={tau:.3f}; calibration spillover is linear in N_c")


def main():
    payload = json.loads(DATA.read_text(encoding="utf-8"))
    check_data(payload)
    check_synthetic(payload)
    check_realcase(payload)
    check_realnet(payload)
    check_theory_anchors()
    print("ALL CORE CHECKS PASSED")


if __name__ == "__main__":
    main()