"""H1 — association between coordination capacity and social scale (no direction claimed).
Seshat (polity), Ethnographic Atlas (society), modern panel (country-year). Usage: python -m wp35.econometrics.h1
"""
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 statsmodels.miscmodels.ordinal_model import OrderedModel
from scipy.stats import norm
from wp35.econometrics.common import fit_formula, coef_row, conley_ols

ROOT = Path(__file__).resolve().parents[3]
PROC = ROOT / "data" / "processed"; RES = ROOT / "results" / "empirical"
RES.mkdir(parents=True, exist_ok=True)


def seshat():
    d = pd.read_parquet(PROC / "seshat_indices.parquet")
    d["quasi"] = d.centralization.eq("quasi-polity")
    rows = []
    base = "C(macro_region) + C(era)"
    specs = {
        "H1-S1 primary: D_pc1 ~ K_mean + FE": ("D_pc1 ~ K_mean + " + base, "K_mean"),
        "H1-S2: D_pc1 ~ K_mean + X + FE": ("D_pc1 ~ K_mean + X_warfare + X_prod_infra + X_military + " + base, "K_mean"),
        "H1-S3: log_pop ~ K_mean + X + FE": ("log_pop ~ K_mean + X_warfare + X_prod_infra + X_military + " + base, "K_mean"),
        "H1-S4: log_terr ~ K_mean + X + FE": ("log_terr ~ K_mean + X_warfare + X_prod_infra + X_military + " + base, "K_mean"),
        "H1-S5: D_pc1 ~ K_irt + X + FE": ("D_pc1 ~ K_irt + X_warfare + X_prod_infra + X_military + " + base, "K_irt"),
        "H1-S6: D_pc1 ~ K blocks + X + FE": ("D_pc1 ~ K_admin + K_rules_money + K_comm_infra + K_records + X_warfare + X_prod_infra + X_military + " + base, "K_admin"),
        "H1-S7: D_pc1 ~ K_mean + X + FE, no quasi-polities": ("D_pc1 ~ K_mean + X_warfare + X_prod_infra + X_military + " + base, "K_mean"),
    }
    for label, (f, term) in specs.items():
        dd = d[~d.quasi] if "no quasi" in label else d
        r = fit_formula(f, dd, None, groups="seshat_region")
        rows.append(coef_row(r, term, label, {"r2": r.rsquared}))
        if "K blocks" in label:
            for t in ["K_rules_money", "K_comm_infra", "K_records"]:
                rows.append(coef_row(r, t, label, {"r2": r.rsquared}))
    # region random effects (MixedLM) on the primary spec
    dd = d.dropna(subset=["D_pc1", "K_mean", "macro_region", "era", "seshat_region"])
    try:
        mm = smf.mixedlm("D_pc1 ~ K_mean + C(era)", dd, groups=dd["seshat_region"]).fit(reml=True)
        rows.append({"model": "H1-S8: D_pc1 ~ K_mean + era, region random intercept", "term": "K_mean", "coef": mm.params["K_mean"], "se": mm.bse["K_mean"],
                     "ci_lo": mm.params["K_mean"] - 1.96 * mm.bse["K_mean"], "ci_hi": mm.params["K_mean"] + 1.96 * mm.bse["K_mean"], "p": mm.pvalues["K_mean"], "n": len(dd), "llf": mm.llf})
    except Exception as e:
        rows.append({"model": "H1-S8 mixedlm", "error": str(e)[:100]})
    # standardized effect: K_mean has range 0-1; report SD-scaled coefficient
    out = pd.DataFrame(rows); out["coef_per_sd_K"] = out.coef * d.K_mean.std()
    out.to_csv(RES / "h1_seshat.csv", index=False)
    return out


def ea():
    d = pd.read_parquet(PROC / "dplace_ea_societies.parquet")
    d["hier"] = d.jurisdictional_hierarchy_beyond; d["cls"] = d.EA066; d["agri"] = d.EA005; d["intens"] = d.EA028; d["subs"] = d.EA042
    rows = []
    # (a) OLS of log community size on hierarchy levels with Conley SE and family-cluster SE
    f = "log_community_size ~ hier + C(subs) + intens + cls + C(region)"
    for cut in [500, 1000]:
        res, se, n = conley_ols(f, d, "Latitude", "Longitude", cut)
        rows.append({"model": f"H1-E1 OLS log community size ~ hierarchy + X + region, Conley {cut} km", "term": "hier", "coef": res.params["hier"], "se": se["hier"],
                     "ci_lo": res.params["hier"] - 1.96 * se["hier"], "ci_hi": res.params["hier"] + 1.96 * se["hier"], "p": 2 * (1 - norm.cdf(abs(res.params["hier"] / se["hier"]))), "n": n, "r2": res.rsquared})
    r = fit_formula(f, d, None, groups="Family_ID"); rows.append(coef_row(r, "hier", "H1-E2 same, SE clustered by language family", {"r2": r.rsquared}))
    f2 = "log_population ~ hier + C(subs) + intens + cls + C(region)"
    r = fit_formula(f2, d, None, groups="Family_ID"); rows.append(coef_row(r, "hier", "H1-E3 log population (EA202) ~ hierarchy + X + region, family-clustered", {"r2": r.rsquared}))
    # (b) ordered logit of hierarchy on log community size
    dd = d.dropna(subset=["hier", "log_community_size", "intens", "cls", "region", "Family_ID"]).copy()
    X = pd.get_dummies(dd[["log_community_size", "intens", "cls", "region"]], columns=["region"], drop_first=True, dtype=float)
    try:
        om = OrderedModel(dd.hier.astype(int), X, distr="logit").fit(method="bfgs", disp=0, maxiter=500)
        # cluster-robust SE by family via sandwich on scores
        groups = pd.factorize(dd.Family_ID)[0]
        from statsmodels.stats.sandwich_covariance import cov_cluster
        try:
            V = cov_cluster(om, groups); se = float(np.sqrt(V[0, 0]))
        except Exception:
            se = float(om.bse.iloc[0])
        b = float(om.params.iloc[0])
        rows.append({"model": "H1-E4 ordered logit hierarchy ~ log community size + X + region (family-clustered)", "term": "log_community_size", "coef": b, "se": se, "ci_lo": b - 1.96 * se, "ci_hi": b + 1.96 * se,
                     "p": 2 * (1 - norm.cdf(abs(b / se))), "n": len(dd), "llf": om.llf})
    except Exception as e:
        rows.append({"model": "H1-E4 ordered logit", "error": str(e)[:120]})
    out = pd.DataFrame(rows); out.to_csv(RES / "h1_ea.csv", index=False)
    return out


def modern():
    p = pd.read_parquet(PROC / "modern_indices.parquet")
    d = p.dropna(subset=["log_pop", "K_modern", "log_gdppc", "v2x_polyarchy"]).copy()
    # within-country (demeaned) regression with year FE
    for c in ["log_pop", "K_modern", "log_gdppc", "v2x_polyarchy"]:
        d[c + "_dm"] = d[c] - d.groupby("country_text_id")[c].transform("mean")
    rows = []
    r = fit_formula("log_pop_dm ~ K_modern_dm + log_gdppc_dm + v2x_polyarchy_dm + C(year)", d, None, groups="country_text_id")
    rows.append(coef_row(r, "K_modern_dm", "H1-M1 within-country: log pop ~ K + log gdppc + polyarchy + year FE", {"r2": r.rsquared}))
    r = fit_formula("log_pop ~ K_modern + log_gdppc + v2x_polyarchy + C(year) + C(e_regionpol_6C)", d, None, groups="country_text_id")
    rows.append(coef_row(r, "K_modern", "H1-M2 pooled: log pop ~ K + log gdppc + polyarchy + year FE + region FE", {"r2": r.rsquared}))
    # long differences 50 years
    d50 = d.merge(d[["country_text_id", "year", "log_pop", "K_modern", "log_gdppc", "v2x_polyarchy"]].assign(year=lambda x: x.year + 50), on=["country_text_id", "year"], suffixes=("", "_lag50"))
    for c in ["log_pop", "K_modern", "log_gdppc", "v2x_polyarchy"]:
        d50["d" + c] = d50[c] - d50[c + "_lag50"]
    r = fit_formula("dlog_pop ~ dK_modern + dlog_gdppc + dv2x_polyarchy + C(year)", d50, None, groups="country_text_id")
    rows.append(coef_row(r, "dK_modern", "H1-M3 50-year differences: Δlog pop ~ ΔK + Δlog gdppc + Δpolyarchy + year FE", {"r2": r.rsquared}))
    out = pd.DataFrame(rows); out.to_csv(RES / "h1_modern.csv", index=False)
    return out


def main():
    a = seshat(); b = ea(); c = modern()
    pd.set_option("display.width", 250); pd.set_option("display.max_colwidth", 90)
    for name, t in [("SESHAT", a), ("EA", b), ("MODERN", c)]:
        print(f"\n=== H1 {name} ===\n", t[[c for c in t.columns if c in ("model", "term", "coef", "se", "p", "n", "r2", "coef_per_sd_K", "error")]].round(4).to_string(index=False))


if __name__ == "__main__":
    main()