"""Leakage check for the Seshat out-of-sample comparison: the D principal-component loadings and the standardization of
P were estimated on the full polity sample (without outcomes) before cross-validation. Here both are re-estimated INSIDE
each leave-one-macro-region-out training fold (loadings from polities outside the held-out region; z-scores from
training crisis records) and the A2-vs-B1 comparison is repeated. Usage: python -m wp35.econometrics.leakage_check
"""
from __future__ import annotations
import numpy as np, pandas as pd, statsmodels.api as sm, statsmodels.formula.api as smf
from pathlib import Path
from wp35.econometrics.common import oos_loglik_binomial
from wp35.measurement.build_indices import pc1_pairwise

ROOT = Path(__file__).resolve().parents[3]; PROC = ROOT / "data/processed"; RES = ROOT / "results/empirical"
BIN = sm.families.Binomial(); XS = "X_warfare + X_prod_infra + epidemic_v"


def main():
    ind = pd.read_parquet(PROC / "seshat_indices.parquet"); cr = pd.read_parquet(PROC / "seshat_h2_sample.parquet")
    agg = pd.read_parquet(PROC / "seshat_precrisis_agg.parquet")
    pw = agg.pivot(index="crisis_id", columns="variable", values="value_mean"); plog = agg.pivot(index="crisis_id", columns="variable", values="logv_mean")
    Dc = pd.DataFrame({"log_pop": plog.get("polity_population"), "log_terr": plog.get("polity_territory"), "log_cap": plog.get("population_of_the_largest_settlement"), "settl_hier": pw.get("settlement_hierarchy")}, index=pw.index)
    Dp = ind.set_index("polity_id")[["log_pop", "log_terr", "log_cap", "settl_hier"]]
    rows = []
    for reg in sorted(cr.macro_region.dropna().unique()):
        tr_pol = ind[ind.macro_region != reg]
        _, load, _ = pc1_pairwise(Dp.loc[tr_pol.polity_id])          # loadings from training polities only
        mu, sd = Dp.loc[tr_pol.polity_id].mean(), Dp.loc[tr_pol.polity_id].std(ddof=0)
        Z = (Dc - mu) / sd; D_new = (Z.fillna(0) * load).sum(axis=1) / (Z.notna() * load.abs()).sum(axis=1).replace(0, np.nan)
        D_new[Z.notna().sum(axis=1) < 2] = np.nan
        d = cr.set_index("crisis_id").copy(); d["D_new"] = D_new.reindex(d.index)
        tr, te = d[d.macro_region != reg], d[d.macro_region == reg]
        zd = lambda s: (s - tr.D_new.mean()) / tr.D_new.std(); zk = lambda s: (s - tr.K_mean.mean()) / tr.K_mean.std()
        for part in (tr, te):
            part["P_new"] = zd(part.D_new) - zk(part.K_mean); part["D_new_z"] = zd(part.D_new)
        need = ["y1_severity", "P_new", "D_new_z", "X_warfare", "X_prod_infra", "epidemic_v", "era"]
        tr2, te2 = tr.dropna(subset=need), te.dropna(subset=need)
        if len(te2) == 0:
            continue
        try:
            mA = smf.glm(f"y1_severity ~ P_new + {XS} + C(era)", tr2, family=BIN).fit(); mB = smf.glm(f"y1_severity ~ D_new_z + {XS} + C(era)", tr2, family=BIN).fit()
            pA, pB = mA.predict(te2), mB.predict(te2); ok = np.isfinite(pA) & np.isfinite(pB)
            rows.append({"held_out": reg, "n_test": int(ok.sum()), "ll_A2_P": oos_loglik_binomial(te2.y1_severity[ok], pA[ok]), "ll_B1_D": oos_loglik_binomial(te2.y1_severity[ok], pB[ok])})
        except Exception as e:
            rows.append({"held_out": reg, "error": str(e)[:60]})
    out = pd.DataFrame(rows); out["P_wins"] = out.ll_A2_P > out.ll_B1_D; out.to_csv(RES / "h2_seshat_cv_noleakage.csv", index=False)
    print(out.round(4).to_string(index=False)); print("win share of P over D-only with fold-internal loadings/standardization:", round(float(out.P_wins.mean()), 3))


if __name__ == "__main__":
    main()