"""Phase 3 measurement: build D, K, P (and rival covariates) exactly as pre-registered, with alternatives,
measurement error, cross-region comparability diagnostics and an a-priori power computation.
Usage: python -m wp35.measurement.build_indices
No outcome (Y1, Y2, Y3, Y4-6) is read or used anywhere in this module.
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import statsmodels.api as sm
from wp35.measurement.irt import fit_2pl
from wp35.harmonization.build_panels import K_BLOCKS, STANDARDS, D_VARS, PROD_INFRA, MILITARY

ROOT = Path(__file__).resolve().parents[3]
PROC = ROOT / "data" / "processed"; RES = ROOT / "results" / "empirical"; FIG = ROOT / "paper" / "figures"
RES.mkdir(parents=True, exist_ok=True); FIG.mkdir(parents=True, exist_ok=True)
SEED = 20261101
K_ITEMS_MODERN = ["v2clrspct", "v2stcritrecadm", "v2cltrnslw", "v2stfisccap", "v2svstterr", "v2elembcap"]
BIN_K = [v for b in K_BLOCKS.values() for v in b if v != "administrative_level"]


def z(x: pd.Series, ref: pd.Series | None = None) -> pd.Series:
    ref = x if ref is None else ref
    return (x - ref.mean()) / ref.std(ddof=0)


def pc1_pairwise(df: pd.DataFrame, min_coded: int = 2):
    """First principal component of the pairwise-complete correlation matrix; scores = loading-weighted mean of the
    available standardized variables (requires >= min_coded coded). Loadings signed so that the sum is positive."""
    Zs = (df - df.mean()) / df.std(ddof=0)
    C = Zs.corr(min_periods=10).fillna(0).values
    w, V = np.linalg.eigh(C); v = V[:, -1]; v = v if v.sum() > 0 else -v
    load = pd.Series(v / np.abs(v).sum(), index=df.columns)
    avail = Zs.notna()
    num = (Zs.fillna(0) * load).sum(axis=1); den = (avail * load.abs()).sum(axis=1)
    score = num / den.where(den > 0)
    score[avail.sum(axis=1) < min_coded] = np.nan
    return score, load, float(w[-1] / w.sum())


def block_score(wide: pd.DataFrame, items: list[str], min_share: float = 0.5) -> pd.Series:
    sub = wide.reindex(columns=items).copy()
    if "administrative_level" in sub.columns:
        r = sub["administrative_level"].rank(method="average")
        sub["administrative_level"] = (r - 1) / (r.max() - 1)
    coded = sub.notna().mean(axis=1)
    return sub.mean(axis=1).where(coded >= min_share)


def build_seshat():
    wide = pd.read_parquet(PROC / "seshat_polity_wide.parquet")
    agg = pd.read_parquet(PROC / "seshat_polity_agg.parquet")
    meta = pd.read_parquet(PROC / "seshat_polity_meta.parquet").set_index("polity_id")
    logw = agg.pivot(index="polity_id", columns="variable", values="logv_mean")
    lo = agg.pivot(index="polity_id", columns="variable", values="lo_min"); hi = agg.pivot(index="polity_id", columns="variable", values="hi_max")
    out = pd.DataFrame(index=wide.index)
    # ---- D ----
    Dmat = pd.DataFrame({"log_pop": logw["polity_population"], "log_terr": logw["polity_territory"],
                         "log_cap": logw["population_of_the_largest_settlement"], "settl_hier": wide["settlement_hierarchy"]})
    out["D_pc1"], D_load, D_var = pc1_pairwise(Dmat)
    out["log_pop"] = Dmat.log_pop; out["log_terr"] = Dmat.log_terr; out["log_cap"] = Dmat.log_cap; out["settl_hier"] = Dmat.settl_hier
    # bounds for measurement-error sensitivity (geometric variables only)
    with np.errstate(divide="ignore"):
        out["log_pop_lo"] = np.log(lo["polity_population"].where(lo["polity_population"] > 0)); out["log_pop_hi"] = np.log(hi["polity_population"].where(hi["polity_population"] > 0))
        out["log_terr_lo"] = np.log(lo["polity_territory"].where(lo["polity_territory"] > 0)); out["log_terr_hi"] = np.log(hi["polity_territory"].where(hi["polity_territory"] > 0))
    # ---- K blocks ----
    for b, items in K_BLOCKS.items():
        out[f"K_{b}"] = block_score(wide, items)
    blocks = out[[f"K_{b}" for b in K_BLOCKS]]
    out["K_mean"] = blocks.mean(axis=1).where(blocks.notna().sum(axis=1) >= 3)
    out["K_n_blocks"] = blocks.notna().sum(axis=1)
    out["K_pc1"], K_load, K_var = pc1_pairwise(blocks, min_coded=3)
    # IRT on binary items (disputed 0.5 -> missing)
    Xb = wide.reindex(columns=BIN_K).copy(); Xb = Xb.where(Xb != 0.5)
    irt = fit_2pl(Xb.values)
    out["K_irt"] = irt["eap"]; out["K_irt_psd"] = irt["psd"]; out["K_irt_nitems"] = irt["n_items"]
    item_params = pd.DataFrame({"item": BIN_K, "a": irt["a"], "b": irt["b"], "block": [next(k for k, v in K_BLOCKS.items() if i in v) for i in BIN_K]})
    item_params.to_csv(RES / "measurement_irt_item_parameters.csv", index=False)
    # standards sub-index (sparse; latent only)
    out["K_standards"] = block_score(wide, STANDARDS)
    # ---- P ----
    for dname in ["D_pc1", "log_pop"]:
        for kname in ["K_mean", "K_irt", "K_pc1"]:
            out[f"P_{dname}_{kname}"] = z(out[dname]) - z(out[kname])
    out["P_main"] = out["P_D_pc1_K_mean"]
    # standardization windows (sensitivity)
    df = out.join(meta[["macro_region", "era"]])
    out["P_within_era"] = df.groupby("era").D_pc1.transform(z) - df.groupby("era").K_mean.transform(z)
    out["P_within_region"] = df.groupby("macro_region").D_pc1.transform(z) - df.groupby("macro_region").K_mean.transform(z)
    out["P_rank"] = out.D_pc1.rank(pct=True) - out.K_mean.rank(pct=True)
    # bounds
    Dlo = pd.DataFrame({"log_pop": out.log_pop_lo, "log_terr": out.log_terr_lo, "log_cap": Dmat.log_cap, "settl_hier": Dmat.settl_hier})
    Dhi = pd.DataFrame({"log_pop": out.log_pop_hi, "log_terr": out.log_terr_hi, "log_cap": Dmat.log_cap, "settl_hier": Dmat.settl_hier})
    out["D_pc1_lo"] = ((Dlo - Dmat.mean()) / Dmat.std(ddof=0)).fillna(0).mul(D_load).sum(axis=1) / (Dlo.notna() * D_load.abs()).sum(axis=1).replace(0, np.nan)
    out["D_pc1_hi"] = ((Dhi - Dmat.mean()) / Dmat.std(ddof=0)).fillna(0).mul(D_load).sum(axis=1) / (Dhi.notna() * D_load.abs()).sum(axis=1).replace(0, np.nan)
    out.loc[out.D_pc1.isna(), ["D_pc1_lo", "D_pc1_hi"]] = np.nan
    out["P_main_lo"] = z(out.D_pc1_lo, out.D_pc1) - z(out.K_mean); out["P_main_hi"] = z(out.D_pc1_hi, out.D_pc1) - z(out.K_mean)
    # drop inferred codes sensitivity: recompute K_mean using only non-inferred facts
    facts = pd.read_parquet(PROC / "seshat_facts.parquet")
    ni = facts[(~facts.inferred) & facts.value.notna()].groupby(["polity_id", "variable"]).value.mean().unstack()
    out["K_mean_noinf"] = pd.concat([block_score(ni, it) for it in K_BLOCKS.values()], axis=1).pipe(lambda b: b.mean(axis=1).where(b.notna().sum(axis=1) >= 3)).reindex(out.index)
    out["P_main_noinf"] = z(out.D_pc1) - z(out.K_mean_noinf)
    # ---- rival covariates (excluded from K by design) ----
    wf = facts[facts.sheet == "Warfare"].groupby(["polity_id", "variable"]).value.mean().unstack()
    out["X_warfare"] = wf.mean(axis=1).where(wf.notna().mean(axis=1) >= 0.5).reindex(out.index)
    out["X_prod_infra"] = block_score(wide, PROD_INFRA)
    mil = wide.reindex(columns=MILITARY).copy(); r = mil["military_level"].rank(); mil["military_level"] = (r - 1) / (r.max() - 1)
    out["X_military"] = mil.mean(axis=1).where(mil.notna().mean(axis=1) >= 0.5)
    out["X_moralizing"] = wide.get("moralizing_supernatural_punishment_and_reward")
    rt = facts[facts.sheet == "Religion"].groupby(["polity_id", "variable"]).value.mean().unstack()
    out["X_moralizing"] = rt.get("moralizing_supernatural_punishment_and_reward").reindex(out.index) if "moralizing_supernatural_punishment_and_reward" in rt else np.nan
    # Turchin-style single complexity axis over ALL social-complexity indicators (baseline B3)
    sc_vars = [v for v in facts[facts.sheet == "Social complexity"].variable.unique() if v in wide.columns]
    scmat = wide[sc_vars].copy()
    for v in ["polity_population", "polity_territory", "population_of_the_largest_settlement"]:
        scmat[v] = logw[v]
    scmat = scmat.loc[:, scmat.notna().sum() >= 100]
    out["complexity_pc1"], C_load, C_var = pc1_pairwise(scmat, min_coded=5)
    out = out.join(meta[["seshat_region", "homeNGA", "macro_region", "cultural_sphere", "n_polities_in_nga", "era", "start_year", "end_year", "midpoint", "duration", "censored", "centralization", "utm_zone"]])
    out.index.name = "polity_id"
    out.reset_index().to_parquet(PROC / "seshat_indices.parquet", index=False)
    # ---- pre-crisis versions (same loadings/standardization as the full sample) ----
    pre = pd.read_parquet(PROC / "seshat_precrisis_agg.parquet")
    pw = pre.pivot(index="crisis_id", columns="variable", values="value_mean"); plog = pre.pivot(index="crisis_id", columns="variable", values="logv_mean")
    cr = pd.DataFrame(index=pw.index)
    Dp = 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)
    Zp = (Dp - Dmat.mean()) / Dmat.std(ddof=0)
    cr["D_pc1"] = (Zp.fillna(0) * D_load).sum(axis=1) / (Zp.notna() * D_load.abs()).sum(axis=1).replace(0, np.nan)
    cr.loc[Zp.notna().sum(axis=1) < 2, "D_pc1"] = np.nan
    cr["log_pop"] = Dp.log_pop
    for b, items in K_BLOCKS.items():
        cr[f"K_{b}"] = block_score(pw, items)
    bl = cr[[f"K_{b}" for b in K_BLOCKS]]; cr["K_mean"] = bl.mean(axis=1).where(bl.notna().sum(axis=1) >= 3)
    Xp = pw.reindex(columns=BIN_K).where(lambda d: d != 0.5)
    # EAP with the full-sample item parameters (no refit): posterior under fixed a, b
    nodes, weights = np.polynomial.hermite_e.hermegauss(41); weights /= weights.sum()
    from scipy.special import expit
    Pq = expit(irt["a"][None, :] * (nodes[:, None] - irt["b"][None, :]))
    Xv = Xp.values; obs = ~np.isnan(Xv); Xz_ = np.where(obs, Xv, 0)
    L = Xz_ @ np.log(Pq + 1e-300).T + ((1 - Xz_) * obs) @ np.log(1 - Pq + 1e-300).T + np.log(weights)[None, :]
    post = np.exp(L - L.max(axis=1, keepdims=True)); post /= post.sum(axis=1, keepdims=True)
    cr["K_irt"] = np.where(obs.sum(axis=1) >= 3, post @ nodes, np.nan)
    cr["K_pc1"] = (((bl - blocks.mean()) / blocks.std(ddof=0)).fillna(0) * K_load).sum(axis=1) / (bl.notna() * K_load.abs()).sum(axis=1).replace(0, np.nan)
    cr.loc[bl.notna().sum(axis=1) < 3, "K_pc1"] = np.nan
    for dname in ["D_pc1", "log_pop"]:
        for kname in ["K_mean", "K_irt", "K_pc1"]:
            cr[f"P_{dname}_{kname}"] = z(cr[dname], out[dname]) - z(cr[kname], out[kname])
    cr["P_main"] = cr["P_D_pc1_K_mean"]
    # pre-registered robustness variants at the crisis level
    plo = pre.pivot(index="crisis_id", columns="variable", values="lo_min"); phi = pre.pivot(index="crisis_id", columns="variable", values="hi_max")
    with np.errstate(divide="ignore"):
        Dlo_c = pd.DataFrame({"log_pop": np.log(plo.get("polity_population").where(plo.get("polity_population") > 0)), "log_terr": np.log(plo.get("polity_territory").where(plo.get("polity_territory") > 0)), "log_cap": Dp.log_cap, "settl_hier": Dp.settl_hier}, index=pw.index)
        Dhi_c = pd.DataFrame({"log_pop": np.log(phi.get("polity_population").where(phi.get("polity_population") > 0)), "log_terr": np.log(phi.get("polity_territory").where(phi.get("polity_territory") > 0)), "log_cap": Dp.log_cap, "settl_hier": Dp.settl_hier}, index=pw.index)
    for nm, Dx in [("lo", Dlo_c), ("hi", Dhi_c)]:
        Zx = (Dx - Dmat.mean()) / Dmat.std(ddof=0); dx = (Zx.fillna(0) * D_load).sum(axis=1) / (Zx.notna() * D_load.abs()).sum(axis=1).replace(0, np.nan)
        dx[Zx.notna().sum(axis=1) < 2] = np.nan; cr[f"D_pc1_{nm}"] = dx
        cr[f"P_main_{nm}"] = z(dx, out.D_pc1) - z(cr.K_mean, out.K_mean)
    # inferred codes dropped: K recomputed from non-inferred facts of the polity (whole-polity; pre-crisis dating not applied)
    cr["K_mean_noinf"] = cr.index.map(lambda c: np.nan); pol_of = crisis.polity_id if False else None
    wfp = pre[pre.variable.isin(wf.columns)].pivot(index="crisis_id", columns="variable", values="value_mean")
    cr["X_warfare"] = wfp.mean(axis=1).where(wfp.notna().mean(axis=1) >= 0.5).reindex(cr.index)
    cr["X_prod_infra"] = block_score(pw, PROD_INFRA)
    milp = pw.reindex(columns=MILITARY).copy(); milp["military_level"] = milp["military_level"].rank().pipe(lambda r: (r - 1) / (r.max() - 1)); cr["X_military"] = milp.mean(axis=1).where(milp.notna().mean(axis=1) >= 0.5)
    scp = pw.reindex(columns=scmat.columns).copy()
    for v in ["polity_population", "polity_territory", "population_of_the_largest_settlement"]:
        scp[v] = plog.get(v)
    Zc = (scp - scmat.mean()) / scmat.std(ddof=0)
    cr["complexity_pc1"] = (Zc.fillna(0) * C_load).sum(axis=1) / (Zc.notna() * C_load.abs()).sum(axis=1).replace(0, np.nan)
    cr.loc[Zc.notna().sum(axis=1) < 5, "complexity_pc1"] = np.nan
    crisis = pd.read_parquet(PROC / "seshat_crisis.parquet").set_index("crisis_id")
    cr["polity_id"] = crisis.polity_id; cr["crisis_start"] = crisis.crisis_start
    cr = cr.join(meta[["macro_region", "cultural_sphere", "era", "seshat_region", "homeNGA", "start_year"]], on="polity_id")
    cr["K_mean_noinf"] = cr.polity_id.map(out.K_mean_noinf); cr["P_main_noinf"] = z(cr.D_pc1, out.D_pc1) - z(cr.K_mean_noinf, out.K_mean_noinf)
    # standardization windows: within era / within macro-region (group moments from the polity sample), and ranks
    for grp, nm in [("era", "within_era"), ("macro_region", "within_region")]:
        gm = df.groupby(grp)[["D_pc1", "K_mean"]].agg(["mean", "std"])
        dz = (cr.D_pc1 - cr[grp].map(gm[("D_pc1", "mean")])) / cr[grp].map(gm[("D_pc1", "std")]); kz = (cr.K_mean - cr[grp].map(gm[("K_mean", "mean")])) / cr[grp].map(gm[("K_mean", "std")])
        cr[f"P_{nm}"] = dz - kz
    cr["P_rank"] = cr.D_pc1.rank(pct=True) - cr.K_mean.rank(pct=True)
    # overlap check (A-006): D without settlement hierarchy; K with the administration block recomputed without administrative levels
    D3 = Dmat[["log_pop", "log_terr", "log_cap"]]; out["D_pc1_nosh"], D3_load, _ = pc1_pairwise(D3)
    Zp3 = (Dp[["log_pop", "log_terr", "log_cap"]] - D3.mean()) / D3.std(ddof=0); cr["D_pc1_nosh"] = ((Zp3.fillna(0) * D3_load).sum(axis=1) / (Zp3.notna() * D3_load.abs()).sum(axis=1).replace(0, np.nan)).where(Zp3.notna().sum(axis=1) >= 2)
    adm_items = [v for v in K_BLOCKS["admin"] if v != "administrative_level"]
    for frame, target in [(wide, out), (pw, cr)]:
        ka = block_score(frame, adm_items); bl_ = pd.concat([ka] + [block_score(frame, K_BLOCKS[k]) for k in ("rules_money", "comm_infra", "records")], axis=1)
        target["K_mean_noadmlev"] = bl_.mean(axis=1).where(bl_.notna().sum(axis=1) >= 3).reindex(target.index)
    cr["P_nosh"] = z(cr.D_pc1_nosh, out.D_pc1_nosh) - z(cr.K_mean, out.K_mean); cr["P_noadmlev"] = z(cr.D_pc1, out.D_pc1) - z(cr.K_mean_noadmlev, out.K_mean_noadmlev)
    cr["P_nosh_noadmlev"] = z(cr.D_pc1_nosh, out.D_pc1_nosh) - z(cr.K_mean_noadmlev, out.K_mean_noadmlev)
    # crisis-sample saturation diagnostics
    diag_cr = {"K_mean_crisis_mean": float(cr.K_mean.mean()), "K_mean_crisis_share_gt_0.8": float((cr.K_mean > 0.8).mean()), "K_mean_crisis_sd": float(cr.K_mean.std()),
               "K_mean_all_mean": float(out.K_mean.mean()), "K_mean_all_sd": float(out.K_mean.std()), "corr_P_D_crisis": float(cr.P_main.corr(cr.D_pc1)), "corr_P_K_crisis": float(cr.P_main.corr(cr.K_mean)),
               "dated_share_precrisis_facts": float(pre.n_dated.sum() / pre.n_facts.sum())}
    (RES / "measurement_crisis_sample.json").write_text(json.dumps(diag_cr, indent=2))
    cr["years_since_start"] = cr.crisis_start - cr.start_year
    cr.reset_index().to_parquet(PROC / "seshat_crisis_indices.parquet", index=False)
    # ---- diagnostics: collinearity, comparability, measurement error ----
    diag = {"D_pc1_loadings": D_load.round(4).to_dict(), "D_pc1_var_share": round(D_var, 4),
            "K_pc1_loadings": K_load.round(4).to_dict(), "K_pc1_var_share": round(K_var, 4),
            "complexity_pc1_var_share": round(C_var, 4), "complexity_pc1_n_vars": int(scmat.shape[1]),
            "irt_loglik": irt["loglik"], "irt_iterations": irt["iterations"],
            "n_D_pc1": int(out.D_pc1.notna().sum()), "n_K_mean": int(out.K_mean.notna().sum()), "n_K_irt": int(out.K_irt.notna().sum()),
            "n_P_main": int(out.P_main.notna().sum()), "n_P_main_crisis": int(cr.P_main.notna().sum()),
            "corr_zD_zK": {"D_pc1~K_mean": round(out.D_pc1.corr(out.K_mean), 4), "D_pc1~K_irt": round(out.D_pc1.corr(out.K_irt), 4),
                            "log_pop~K_mean": round(out.log_pop.corr(out.K_mean), 4), "K_mean~K_irt": round(out.K_mean.corr(out.K_irt), 4),
                            "K_mean~K_pc1": round(out.K_mean.corr(out.K_pc1), 4), "P_main~complexity_pc1": round(out.P_main.corr(out.complexity_pc1), 4),
                            "P_main~D_pc1": round(out.P_main.corr(out.D_pc1), 4), "P_main~K_mean": round(out.P_main.corr(out.K_mean), 4)},
            "P_variance_share_of_D_plus_K": round(float(out.P_main.var() / (z(out.D_pc1).var() + z(out.K_mean).var())), 4),
            "block_correlations": blocks.corr().round(3).to_dict(),
            "block_alpha_by_macro_region": {}, "block_corr_mean_by_era": {},
            "measurement_error": {"D_pc1_halfwidth_median": float(((out.D_pc1_hi - out.D_pc1_lo) / 2).median()),
                                  "D_pc1_halfwidth_q90": float(((out.D_pc1_hi - out.D_pc1_lo) / 2).quantile(0.9)),
                                  "K_irt_psd_median": float(out.K_irt_psd.median()), "K_irt_psd_q90": float(out.K_irt_psd.quantile(0.9)),
                                  "sd_K_irt": float(out.K_irt.std()), "P_main_range_median": float((out.P_main_hi - out.P_main_lo).abs().median())}}
    for r, g in df.groupby("macro_region"):
        b = g[[f"K_{k}" for k in K_BLOCKS]].dropna()
        if len(b) >= 8:
            k = b.shape[1]; alpha = k / (k - 1) * (1 - b.var(ddof=1).sum() / b.sum(axis=1).var(ddof=1))
            diag["block_alpha_by_macro_region"][r] = {"n": int(len(b)), "alpha": round(float(alpha), 3)}
    for e, g in df.groupby("era"):
        b = g[[f"K_{k}" for k in K_BLOCKS]]
        c = b.corr().values; diag["block_corr_mean_by_era"][e] = {"n": int(b.dropna().shape[0]), "mean_offdiag_corr": round(float(c[np.triu_indices(4, 1)].mean()), 3)}
    # IRT invariance: refit on Old World vs Americas/Oceania and compare difficulties
    new_world = df.macro_region.isin(["North_America", "South_America", "Southeast_Asia_Oceania"])
    inv = {}
    for name, mask in [("old_world", ~new_world), ("new_world_oceania", new_world)]:
        f = fit_2pl(Xb[mask.values].values); inv[name] = {"n": int(mask.sum()), "b": f["b"].round(3).tolist(), "a": f["a"].round(3).tolist()}
    bo, bn = np.array(inv["old_world"]["b"]), np.array(inv["new_world_oceania"]["b"])
    ok = np.isfinite(bo) & np.isfinite(bn) & (np.abs(bo) < 5.9) & (np.abs(bn) < 5.9)
    diag["irt_invariance"] = {"items_compared": int(ok.sum()), "corr_difficulty_old_vs_new": round(float(np.corrcoef(bo[ok], bn[ok])[0, 1]), 3),
                              "mean_abs_diff_difficulty": round(float(np.abs(bo[ok] - bn[ok]).mean()), 3), "items": BIN_K, "fits": inv}
    # leave-one-block-out sensitivity of K and P
    loo = {}
    for b in K_BLOCKS:
        others = [f"K_{k}" for k in K_BLOCKS if k != b]; kk = out[others].mean(axis=1).where(out[others].notna().sum(axis=1) >= 2)
        loo[b] = {"corr_K_with_full": round(float(kk.corr(out.K_mean)), 4), "corr_P_with_main": round(float((z(out.D_pc1) - z(kk)).corr(out.P_main)), 4)}
    diag["leave_one_block_out"] = loo
    (RES / "measurement_diagnostics.json").write_text(json.dumps(diag, indent=2))
    # ---- figures ----
    fig, axes = plt.subplots(1, 3, figsize=(12, 3.8))
    axes[0].scatter(out.K_mean, out.D_pc1, s=8, alpha=0.5, c="#444"); axes[0].set_xlabel("K (mean of 4 blocks)"); axes[0].set_ylabel("D (PC1 of log scale variables)"); axes[0].set_title(f"r = {out.D_pc1.corr(out.K_mean):.2f}, n = {int((out.D_pc1.notna() & out.K_mean.notna()).sum())}", fontsize=9)
    axes[1].scatter(out.K_mean, out.K_irt, s=8, alpha=0.5, c="#444"); axes[1].set_xlabel("K mean"); axes[1].set_ylabel("K IRT (2PL EAP)"); axes[1].set_title(f"r = {out.K_mean.corr(out.K_irt):.2f}", fontsize=9)
    axes[2].errorbar(item_params.b, range(len(item_params)), xerr=None, fmt="o", ms=3, color="#444"); axes[2].set_yticks(range(len(item_params))); axes[2].set_yticklabels(item_params.item, fontsize=5); axes[2].set_xlabel("IRT difficulty b"); axes[2].set_title("Item difficulties", fontsize=9)
    fig.tight_layout(); fig.savefig(FIG / "fig_measurement_DK.pdf"); plt.close(fig)
    fig, ax = plt.subplots(figsize=(4.5, 4.2)); ax.scatter(bo[ok], bn[ok], s=12, c="#444"); lim = [min(bo[ok].min(), bn[ok].min()) - .3, max(bo[ok].max(), bn[ok].max()) + .3]; ax.plot(lim, lim, "k--", lw=.8)
    ax.set_xlabel("difficulty, Old World fit"); ax.set_ylabel("difficulty, Americas/Oceania fit"); ax.set_title(f"IRT item invariance (r = {diag['irt_invariance']['corr_difficulty_old_vs_new']:.2f})", fontsize=9)
    fig.tight_layout(); fig.savefig(FIG / "fig_irt_invariance.pdf"); plt.close(fig)
    return out, cr, diag


def build_modern():
    p = pd.read_parquet(PROC / "modern_panel.parquet")
    items = p[K_ITEMS_MODERN]
    Zi = (items - items.mean()) / items.std(ddof=0)
    p["K_modern"] = Zi.mean(axis=1).where(Zi.notna().sum(axis=1) >= 4)
    p["K_modern_pc1"], Km_load, Km_var = pc1_pairwise(items, min_coded=4)
    p["D_modern"] = p.log_pop
    p["P_modern"] = z(p.D_modern) - z(p.K_modern)
    p["P_modern_pc1"] = z(p.D_modern) - z(p.K_modern_pc1)
    p["P_modern_within_year"] = p.groupby("year").D_modern.transform(z) - p.groupby("year").K_modern.transform(z)
    p["P_modern_area"] = z(((z(p.log_pop) + z(p.log_area)) / 2)) - z(p.K_modern)
    p.to_parquet(PROC / "modern_indices.parquet", index=False)
    diag = {"K_modern_pc1_loadings": Km_load.round(4).to_dict(), "K_modern_pc1_var_share": round(Km_var, 4),
            "n_P_modern": int(p.P_modern.notna().sum()), "corr_D_K_modern": round(float(p.D_modern.corr(p.K_modern)), 4),
            "corr_K_modern_polyarchy": round(float(p.K_modern.corr(p.v2x_polyarchy)), 4),
            "item_corr": items.corr().round(3).to_dict(),
            "corr_K_by_region": {int(r): round(float(g.D_modern.corr(g.K_modern)), 3) for r, g in p.groupby("e_regionpol_6C") if g.P_modern.notna().sum() > 50}}
    (RES / "measurement_diagnostics_modern.json").write_text(json.dumps(diag, indent=2))
    return p, diag


def power_h2a(cr: pd.DataFrame, n_rep: int = 400):
    """A-priori MDE for H2-A: random outcomes assigned to the actual P distribution of crisis records (no real Y used)."""
    rng = np.random.default_rng(SEED)
    d = cr.dropna(subset=["P_main"]).copy()
    P = z(d.P_main).values; reg = pd.factorize(d.seshat_region)[0]; n = len(P)
    base = 0.77  # Y1b positive share read from coverage_summary (122/158); used only to set the intercept
    res = []
    for beta in [0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1.0]:
        rej = 0
        for _ in range(n_rep):
            a0 = np.log(base / (1 - base))
            y = rng.random(n) < 1 / (1 + np.exp(-(a0 + beta * P)))
            try:
                m = sm.Logit(y.astype(float), sm.add_constant(P)).fit(disp=0, cov_type="cluster", cov_kwds={"groups": reg})
                rej += m.pvalues[1] < 0.05
            except Exception:
                pass
        res.append({"beta": beta, "power": rej / n_rep, "n": n})
    res = pd.DataFrame(res); res.to_csv(RES / "power_h2a_binary.csv", index=False)
    # fractional outcome: Binomial(14, logistic(a0 + beta P)) with base severity 0.29
    res2 = []
    for beta in [0.1, 0.15, 0.2, 0.3, 0.4, 0.5]:
        rej = 0
        for _ in range(n_rep):
            a0 = np.log(0.29 / 0.71); pr = 1 / (1 + np.exp(-(a0 + beta * P))); y = rng.binomial(14, pr) / 14
            try:
                m = sm.GLM(y, sm.add_constant(P), family=sm.families.Binomial()).fit(cov_type="cluster", cov_kwds={"groups": reg})
                rej += m.pvalues[1] < 0.05
            except Exception:
                pass
        res2.append({"beta": beta, "power": rej / n_rep, "n": n})
    res2 = pd.DataFrame(res2); res2.to_csv(RES / "power_h2a_fractional.csv", index=False)
    return res, res2


def main():
    out, cr, diag = build_seshat()
    p, dm = build_modern()
    pw1, pw2 = power_h2a(cr)
    print(json.dumps({k: v for k, v in diag.items() if k in ("D_pc1_loadings", "D_pc1_var_share", "K_pc1_loadings", "n_D_pc1", "n_K_mean", "n_K_irt", "n_P_main", "n_P_main_crisis", "corr_zD_zK", "P_variance_share_of_D_plus_K", "measurement_error", "irt_invariance", "leave_one_block_out", "block_alpha_by_macro_region")}, indent=1, default=str)[:6000])
    print(json.dumps({k: v for k, v in dm.items() if k != "item_corr"}, indent=1))
    print("power binary:\n", pw1.to_string(index=False)); print("power fractional:\n", pw2.to_string(index=False))


if __name__ == "__main__":
    main()