"""H2 (and H4/H5) on the modern panel (V-Dem v16 + Maddison 2023), pre-registration §5.4.
Primary outcome Y456 = any instability event in (t, t+5]; pooled logit with region and year FE, country-clustered SE.
Usage: python -m wp35.econometrics.h2_modern
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf
from wp35.econometrics.common import fit_formula, coef_row, binary_metrics, oos_loglik_binomial, hinge_fit, fold_win_share

ROOT = Path(__file__).resolve().parents[3]
PROC = ROOT / "data" / "processed"; RES = ROOT / "results" / "empirical"
BIN = sm.families.Binomial()
FE = "C(e_regionpol_6C) + C(year)"
XS = "log_gdppc + v2x_polyarchy + t_since + t_since2 + t_since3 + no_prior_event"


def window_outcome(df: pd.DataFrame, col: str, h: int, forward: bool = True) -> pd.Series:
    """any(event in (t, t+h]) (forward) or any(event in [t-h, t-1]) (backward); NaN if no coded year in the window."""
    out = pd.Series(np.nan, index=df.index)
    for cid, g in df.groupby("country_text_id"):
        s = g.set_index("year")[col]
        yrs = g.year.values
        vals = []
        for t in yrs:
            rng = range(t + 1, t + h + 1) if forward else range(t - h, t)
            w = s.reindex(rng)
            vals.append(np.nan if w.notna().sum() == 0 else float((w == 1).any()))
        out.loc[g.index] = vals
    return out


def build_sample(h: int = 5):
    p = pd.read_parquet(PROC / "modern_indices.parquet").sort_values(["country_text_id", "year"]).reset_index(drop=True)
    p["any_event"] = p[["y4_civilwar_onset", "y5_coup", "y6_irregular_regime_end"]].max(axis=1, skipna=True)
    p.loc[p[["y4_civilwar_onset", "y5_coup", "y6_irregular_regime_end"]].isna().all(axis=1), "any_event"] = np.nan
    for col, name in [("any_event", "y456"), ("y4_civilwar_onset", "y4"), ("y5_coup", "y5"), ("y6_irregular_regime_end", "y6")]:
        p[f"{name}_f{h}"] = window_outcome(p, col, h, forward=True)
    p["y456_b5"] = window_outcome(p, "any_event", 5, forward=False)   # persistence / negative control
    # time since last event (years), cubic polynomial (Carter & Signorino 2010)
    p["t_since"] = np.nan
    for cid, g in p.groupby("country_text_id"):
        last = None; vals = []
        for y, e in zip(g.year, g.any_event):
            vals.append(np.nan if last is None else y - last)
            if e == 1:
                last = y
        p.loc[g.index, "t_since"] = vals
    # no prior event observed: spell measured from the country's first observed year (left-censored), with an indicator (A-004)
    first = p.groupby("country_text_id").year.transform("min")
    p["no_prior_event"] = p.t_since.isna().astype(float)
    p["t_since"] = p.t_since.fillna(p.year - first + 1) / 100.0
    p["t_since2"] = p.t_since ** 2; p["t_since3"] = p.t_since ** 3
    for c in ["P_modern", "D_modern", "K_modern", "P_modern_pc1", "P_modern_within_year", "P_modern_area"]:
        p[c + "_z"] = (p[c] - p[c].mean()) / p[c].std()
    p["P_placebo_z"] = ((p.D_modern - p.D_modern.mean()) / p.D_modern.std() - (p.v2x_polyarchy - p.v2x_polyarchy.mean()) / p.v2x_polyarchy.std())
    p["period"] = pd.cut(p.year, [1700, 1899, 1945, 1989, 2100], labels=["pre-1900", "1900-1945", "1946-1989", "1990+"])
    return p


def primary(d: pd.DataFrame, y: str = "y456_f5"):
    rows = []
    specs = {
        "C1 primary: Y456 ~ P + X + region FE + year FE": (f"{y} ~ P_modern_z + {XS} + {FE}", "P_modern_z"),
        "C2: Y456 ~ P + FE only": (f"{y} ~ P_modern_z + {FE}", "P_modern_z"),
        "C3: Y456 ~ P + K + X + FE (reparametrization of B5: beta_P = beta_D)": (f"{y} ~ P_modern_z + K_modern_z + {XS} + {FE}", "P_modern_z"),
        "B1 baseline population only": (f"{y} ~ D_modern_z + {XS} + {FE}", "D_modern_z"),
        "B2 baseline capacity only": (f"{y} ~ K_modern_z + {XS} + {FE}", "K_modern_z"),
        "B3 baseline development (gdppc) only": (f"{y} ~ log_gdppc + v2x_polyarchy + t_since + t_since2 + t_since3 + {FE}", "log_gdppc"),
        "B4 baseline persistence": (f"{y} ~ y456_b5 + {XS} + {FE}", "y456_b5"),
        "B5 D and K separately": (f"{y} ~ D_modern_z + K_modern_z + {XS} + {FE}", "D_modern_z"),
        "Placebo: P with polyarchy in place of K (polyarchy not separately controlled)": (f"{y} ~ P_placebo_z + log_gdppc + t_since + t_since2 + t_since3 + {FE}", "P_placebo_z"),
        "Placebo comparison: P_main with the same controls as the placebo": (f"{y} ~ P_modern_z + log_gdppc + t_since + t_since2 + t_since3 + {FE}", "P_modern_z"),
        "Negative control: past events (t-5..t-1) ~ P_t + gdppc + polyarchy + FE (no duration terms)": (f"y456_b5 ~ P_modern_z + log_gdppc + v2x_polyarchy + {FE}", "P_modern_z"),
        "Forward comparison for the negative control: future events ~ P_t + gdppc + polyarchy + FE": (f"{y} ~ P_modern_z + log_gdppc + v2x_polyarchy + {FE}", "P_modern_z"),
    }
    for label, (f, term) in specs.items():
        try:
            r = fit_formula(f, d, BIN, groups="country_text_id"); rows.append(coef_row(r, term, label, {"events": int(r._data_used[f.split(" ~")[0]].sum())}))
            if label.startswith("B5"):
                rows.append(coef_row(r, "K_modern_z", label))
                from wp35.econometrics.h2_seshat import wald_sum_zero
                wt = wald_sum_zero(r, "D_modern_z", "K_modern_z"); rows.append({"model": "B5 Wald test of beta_D + beta_K = 0 (restriction implied by P)", "term": "D+K", "coef": wt["sum_coef"], "se": wt["sum_se"], "p": wt["wald_p"], "n": r._n_used})
        except Exception as e:
            rows.append({"model": label, "error": str(e)[:100]})
    # country fixed effects: conditional logit via statsmodels (drops never-event countries)
    try:
        from statsmodels.discrete.conditional_models import ConditionalLogit
        dd = d.dropna(subset=[y, "P_modern_z", "log_gdppc", "v2x_polyarchy", "t_since"]).copy()
        dd["decade"] = (dd.year // 10) * 10
        X = pd.concat([dd[["P_modern_z", "log_gdppc", "v2x_polyarchy", "t_since", "t_since2", "t_since3", "no_prior_event"]], pd.get_dummies(dd.decade, prefix="dec", drop_first=True, dtype=float)], axis=1)
        cl = ConditionalLogit(dd[y].values, X.values, groups=pd.factorize(dd.country_text_id)[0]).fit(disp=0, maxiter=200)
        b, se = cl.params[0], cl.bse[0]
        rows.append({"model": "C4: conditional logit, country FE + decade FE", "term": "P_modern_z", "coef": b, "se": se, "ci_lo": b - 1.96 * se, "ci_hi": b + 1.96 * se, "p": cl.pvalues[0], "n": len(dd), "events": int(dd[y].sum())})
    except Exception as e:
        rows.append({"model": "C4 conditional logit", "error": str(e)[:120]})
    # LPM two-way FE
    try:
        dd = d.dropna(subset=[y, "P_modern_z", "log_gdppc", "v2x_polyarchy", "t_since"]).copy()
        for c in [y, "P_modern_z", "log_gdppc", "v2x_polyarchy", "t_since", "t_since2", "t_since3", "no_prior_event"]:
            dd[c + "_dm"] = dd[c] - dd.groupby("country_text_id")[c].transform("mean")
        r = fit_formula(f"{y}_dm ~ P_modern_z_dm + log_gdppc_dm + v2x_polyarchy_dm + t_since_dm + t_since2_dm + t_since3_dm + no_prior_event_dm + C(year)", dd, None, groups="country_text_id")
        rows.append(coef_row(r, "P_modern_z_dm", "C5: linear probability, country + year FE", {"events": int(dd[y].sum())}))
    except Exception as e:
        rows.append({"model": "C5 LPM", "error": str(e)[:100]})
    return pd.DataFrame(rows)


def oos(d: pd.DataFrame, y: str = "y456_f5"):
    formulas = {"C1 P": f"{y} ~ P_modern_z + {XS} + {FE}", "B0 FE+X": f"{y} ~ {XS} + {FE}", "B1 D": f"{y} ~ D_modern_z + {XS} + {FE}", "B2 K": f"{y} ~ K_modern_z + {XS} + {FE}",
                "B4 persistence": f"{y} ~ y456_b5 + {XS} + {FE}", "B5 D+K": f"{y} ~ D_modern_z + K_modern_z + {XS} + {FE}"}
    rows = []; preds = []
    need = [y, "P_modern_z", "D_modern_z", "K_modern_z", "y456_b5", "log_gdppc", "v2x_polyarchy", "t_since", "e_regionpol_6C", "year"]
    dd = d.dropna(subset=need).copy()
    # forward chaining (year FE replaced by decade FE so that test years have a fitted level: pre-registered deviation logged as A-001)
    fc = {k: v.replace("C(year)", "C(decade)") for k, v in formulas.items()}
    dd["decade"] = (dd.year // 10) * 10
    splits = [("train<=1945,test1946-1975", dd.year <= 1945, dd.year.between(1946, 1975)), ("train<=1975,test1976-2005", dd.year <= 1975, dd.year.between(1976, 2005)), ("train<=2005,test>=2006", dd.year <= 2005, dd.year >= 2006)]
    for name, f in fc.items():
        for sname, tr, te in splits:
            try:
                m = smf.glm(f, dd[tr], family=BIN).fit()
                te_d = dd[te].copy(); te_d["decade"] = np.minimum(te_d.decade, dd[tr].decade.max())  # unseen decades mapped to the last training decade
                seen = set(dd[tr].e_regionpol_6C.unique()); te_d.loc[~te_d.e_regionpol_6C.isin(seen), "e_regionpol_6C"] = dd[tr].e_regionpol_6C.mode().iloc[0]  # unseen regions (A-004)
                p = m.predict(te_d); r = binary_metrics(te_d[y], p); r.update({"model": name, "fold": sname, "scheme": "forward"}); rows.append(r)
                preds.append(pd.DataFrame({"scheme": "forward", "fold": sname, "model": name, "y": te_d[y].values, "p": np.asarray(p)}))
            except Exception as e:
                rows.append({"model": name, "fold": sname, "scheme": "forward", "error": str(e)[:80]})
    # leave-one-region-out (region FE removed)
    lr = {k: v.replace("C(e_regionpol_6C) + ", "") for k, v in formulas.items()}
    for name, f in lr.items():
        for reg in sorted(dd.e_regionpol_6C.unique()):
            try:
                m = smf.glm(f, dd[dd.e_regionpol_6C != reg], family=BIN).fit(); te_d = dd[dd.e_regionpol_6C == reg]
                p = m.predict(te_d); r = binary_metrics(te_d[y], p); r.update({"model": name, "fold": f"loro_{int(reg)}", "scheme": "loro"}); rows.append(r)
                preds.append(pd.DataFrame({"scheme": "loro", "fold": f"loro_{int(reg)}", "model": name, "y": te_d[y].values, "p": np.asarray(p)}))
            except Exception as e:
                rows.append({"model": name, "fold": f"loro_{int(reg)}", "scheme": "loro", "error": str(e)[:80]})
    cv = pd.DataFrame(rows); cv.to_csv(RES / "h2_modern_cv_folds.csv", index=False)
    pd.concat(preds).to_parquet(RES / "h2_modern_cv_predictions.parquet", index=False)
    wins = []
    for scheme in ["forward", "loro"]:
        c = cv[cv.scheme == scheme]
        for b in ["B0 FE+X", "B1 D", "B2 K", "B4 persistence", "B5 D+K"]:
            for metric in ["oos_loglik", "auc"]:
                w = fold_win_share(c, "C1 P", b, metric); w.update({"scheme": scheme, "metric": metric}); wins.append(w)
    wins = pd.DataFrame(wins); wins.to_csv(RES / "h2_modern_cv_wins.csv", index=False)
    summ = cv.groupby(["scheme", "model"])[["oos_loglik", "auc", "brier", "calibration_slope"]].mean().reset_index(); summ.to_csv(RES / "h2_modern_cv_summary.csv", index=False)
    return wins, summ


def h4(d: pd.DataFrame, y: str = "y456_f5"):
    base = f"{y} ~ P_modern_z + {XS} + C(e_regionpol_6C) + C(decade)"
    from wp35.econometrics.common import map_unseen_levels
    dd = d.dropna(subset=[y, "P_modern_z", "log_gdppc", "v2x_polyarchy", "t_since", "e_regionpol_6C"]).copy(); dd["decade"] = (dd.year // 10) * 10
    rng = np.random.default_rng(20261101); rows = []
    countries = dd.country_text_id.unique()
    for rep in range(3):
        perm = rng.permutation(countries); parts = np.array_split(perm, 5)   # country-blocked folds
        for i, te_c in enumerate(parts):
            te = dd.country_text_id.isin(te_c); tr, ts = dd[~te], map_unseen_levels(dd[~te], dd[te], base)
            try:
                lin = smf.glm(base, tr, family=BIN).fit(); ll_lin = oos_loglik_binomial(ts[y], lin.predict(ts))
                h = hinge_fit(tr, y, base, "P_modern_z", BIN, grid=np.quantile(tr.P_modern_z, np.linspace(0.2, 0.8, 7)))
                ts["_hinge"] = np.clip(ts.P_modern_z - h[1], 0, None); ll_h = oos_loglik_binomial(ts[y], h[2].predict(ts))
                sp = smf.glm(base.replace("P_modern_z", "cr(P_modern_z, df=3)"), tr, family=BIN).fit(); ll_s = oos_loglik_binomial(ts[y], sp.predict(ts))
                rows.append({"fold": f"{rep}_{i}", "ll_linear": ll_lin, "ll_hinge": ll_h, "ll_spline": ll_s, "knot": h[1]})
            except Exception as e:
                rows.append({"fold": f"{rep}_{i}", "error": str(e)[:80]})
    r = pd.DataFrame(rows); r.to_csv(RES / "h4_modern_folds.csv", index=False); ok = r.dropna(subset=["ll_linear"])
    # in-sample shape (for the figure): predicted probability vs P at covariate means, linear vs spline
    sp = smf.glm(base.replace("P_modern_z", "cr(P_modern_z, df=3)"), dd, family=BIN).fit(); lin = smf.glm(base, dd, family=BIN).fit()
    grid = np.linspace(dd.P_modern_z.quantile(0.02), dd.P_modern_z.quantile(0.98), 40)
    ref = dd.iloc[[0]].copy()
    for c in ["log_gdppc", "v2x_polyarchy", "t_since", "t_since2", "t_since3"]:
        ref[c] = dd[c].mean()
    ref["e_regionpol_6C"] = dd.e_regionpol_6C.mode()[0]; ref["decade"] = 1970
    g = pd.concat([ref] * len(grid), ignore_index=True); g["P_modern_z"] = grid
    pd.DataFrame({"P_z": grid, "p_linear": lin.predict(g), "p_spline": sp.predict(g)}).to_csv(RES / "h4_modern_curve.csv", index=False)
    return {"folds": int(len(ok)), "hinge_beats_linear_share": float((ok.ll_hinge > ok.ll_linear).mean()), "spline_beats_linear_share": float((ok.ll_spline > ok.ll_linear).mean()),
            "mean_ll_linear": float(ok.ll_linear.mean()), "mean_ll_hinge": float(ok.ll_hinge.mean()), "mean_ll_spline": float(ok.ll_spline.mean()), "median_knot": float(ok.knot.median())}


def h5(d: pd.DataFrame, y: str = "y456_f5"):
    rows = []
    for col, name in [("e_regionpol_6C", "region"), ("period", "period")]:
        for g in sorted(d[col].dropna().unique()):
            held = d[d[col] != g]; inb = d[d[col] == g]
            fe = "C(year)" if col == "e_regionpol_6C" else "C(e_regionpol_6C) + C(year)"
            try:
                r = fit_formula(f"{y} ~ P_modern_z + {XS} + {fe}", held, BIN, groups="country_text_id"); row = coef_row(r, "P_modern_z", f"H5 leave-{name}-out: {g}")
                rb = fit_formula(f"{y} ~ P_modern_z + {XS} + " + ("C(year)" if col == "e_regionpol_6C" else "C(e_regionpol_6C)"), inb, BIN, groups="country_text_id")
                row.update({"within_block_coef": rb.params["P_modern_z"], "within_block_se": rb.bse["P_modern_z"], "within_block_n": rb._n_used, "within_block_events": int(rb._data_used[y].sum())})
                rows.append(row)
            except Exception as e:
                rows.append({"model": f"H5 {name} {g}", "error": str(e)[:80]})
    out = pd.DataFrame(rows); out.to_csv(RES / "h5_modern_blocks.csv", index=False); return out


def robustness(d: pd.DataFrame):
    rows = []
    for h in [1, 3, 10]:
        dh = build_sample(h) if h != 5 else d
        r = fit_formula(f"y456_f{h} ~ P_modern_z + {XS} + {FE}", dh, BIN, groups="country_text_id"); rows.append(coef_row(r, "P_modern_z", f"R: horizon h={h}", {"events": int(r._data_used[f"y456_f{h}"].sum())}))
    for comp in ["y4_f5", "y5_f5", "y6_f5"]:
        r = fit_formula(f"{comp} ~ P_modern_z + {XS} + {FE}", d, BIN, groups="country_text_id"); rows.append(coef_row(r, "P_modern_z", f"R: component {comp}", {"events": int(r._data_used[comp].sum())}))
    for v, lab in [("P_modern_pc1_z", "K = PC1"), ("P_modern_within_year_z", "standardized within year"), ("P_modern_area_z", "D = (log pop + log area)/2")]:
        r = fit_formula(f"y456_f5 ~ {v} + {XS} + {FE}", d, BIN, groups="country_text_id"); rows.append(coef_row(r, v, f"R: {lab}", {"events": int(r._data_used.y456_f5.sum())}))
    r = fit_formula(f"y456_f5 ~ P_modern_z + {XS} + {FE}", d[d.year >= 1900], BIN, groups="country_text_id"); rows.append(coef_row(r, "P_modern_z", "R: years ≥ 1900", {"events": int(r._data_used.y456_f5.sum())}))
    r = fit_formula(f"y456_f5 ~ P_modern_z + log_gdppc + t_since + t_since2 + t_since3 + {FE}", d, BIN, groups="country_text_id"); rows.append(coef_row(r, "P_modern_z", "R: polyarchy not controlled", {"events": int(r._data_used.y456_f5.sum())}))
    r = fit_formula(f"y456_f5 ~ P_modern_z + {XS} + y456_b5 + {FE}", d, BIN, groups="country_text_id"); rows.append(coef_row(r, "P_modern_z", "R: + persistence control", {"events": int(r._data_used.y456_f5.sum())}))
    # Oster / Cinelli-Hazlett style sensitivity on the LPM version of the primary
    try:
        dd = d.dropna(subset=["y456_f5", "P_modern_z", "log_gdppc", "v2x_polyarchy", "t_since", "e_regionpol_6C"]).copy()
        # distinct event-years (not windows) in the estimation sample, for the text
        (RES / "h2_modern_event_counts.json").write_text(json.dumps({"country_years": int(len(dd)), "event_windows": int(dd.y456_f5.sum()), "distinct_event_years_in_sample": int((dd.any_event == 1).sum()), "distinct_event_years_panel": int((d.any_event == 1).sum())}))
        full = smf.ols(f"y456_f5 ~ P_modern_z + {XS} + {FE}", dd).fit(); short = smf.ols(f"y456_f5 ~ P_modern_z + {FE}", dd).fit()
        restricted = smf.ols(f"y456_f5 ~ {XS} + {FE}", dd).fit()
        r2_full, r2_short = full.rsquared, short.rsquared; b_full, b_short = full.params["P_modern_z"], short.params["P_modern_z"]
        rmax = min(1.0, 1.3 * r2_full); delta1 = b_full - (b_short - b_full) * (rmax - r2_full) / max(r2_full - r2_short, 1e-9)
        t = full.tvalues["P_modern_z"]; dfres = full.df_resid; r2_yd = t ** 2 / (t ** 2 + dfres)
        rv = 0.5 * (np.sqrt(r2_yd ** 2 + 4 * r2_yd) - r2_yd)
        rows.append({"model": "Sensitivity (LPM): Oster bias-adjusted β at Rmax=1.3·R² (δ=1); Cinelli–Hazlett RV", "term": "P_modern_z", "coef": delta1, "oster_beta_full": b_full, "oster_beta_short": b_short, "r2_full": r2_full, "r2_short": r2_short, "partial_r2_P": r2_yd, "robustness_value": rv, "n": len(dd)})
    except Exception as e:
        rows.append({"model": "Sensitivity", "error": str(e)[:100]})
    out = pd.DataFrame(rows); out.to_csv(RES / "h2_modern_robustness.csv", index=False); return out


def main():
    d = build_sample(5); d.to_parquet(PROC / "modern_h2_sample.parquet", index=False)
    prim = primary(d); prim.to_csv(RES / "h2_modern_primary.csv", index=False)
    wins, summ = oos(d)
    h4s = h4(d); (RES / "h4_modern_summary.json").write_text(json.dumps(h4s, indent=2))
    h5s = h5(d); rob = robustness(d)
    pd.set_option("display.width", 250); pd.set_option("display.max_colwidth", 70)
    cols = ["model", "term", "coef", "se", "p", "n", "events", "error"]
    print("=== PRIMARY ===\n", prim[[c for c in cols if c in prim.columns]].round(3).to_string(index=False))
    print("\n=== OOS wins ===\n", wins.round(3).to_string(index=False)); print("\n=== OOS summary ===\n", summ.round(3).to_string(index=False))
    print("\n=== H4 ===", json.dumps(h4s, indent=1))
    print("\n=== H5 ===\n", h5s[[c for c in h5s.columns if c in ("model", "coef", "se", "p", "n", "within_block_coef", "within_block_se", "within_block_n", "within_block_events", "error")]].round(3).to_string(index=False))
    print("\n=== ROBUSTNESS ===\n", rob[[c for c in rob.columns if c in cols + ["robustness_value", "partial_r2_P", "oster_beta_full", "oster_beta_short"]]].round(4).to_string(index=False))


if __name__ == "__main__":
    main()