"""Phase 2 orchestrator: build the three panels, write processed files, coverage tables, missingness diagnostics,
temporal-distribution figures and the data dictionary.  Usage: python -m wp35.harmonization.build_panels
No predictor-outcome association is computed here.
"""
from __future__ import annotations
import json
from datetime import datetime, timezone
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from wp35.harmonization import seshat as S, dplace as D, modern as M

ROOT = Path(__file__).resolve().parents[3]
PROC = ROOT / "data" / "processed"; META = ROOT / "data" / "metadata"; DIAG = ROOT / "results" / "diagnostics"
FIG = ROOT / "paper" / "figures"; TAB = ROOT / "paper" / "tables"
for p in (PROC, META, DIAG, FIG, TAB):
    p.mkdir(parents=True, exist_ok=True)

K_BLOCKS = {
    "admin": ["full_time_bureaucrat", "specialized_government_building", "merit_promotion", "examination_system", "administrative_level"],
    "rules_money": ["formal_legal_code", "judge", "court", "professional_lawyer", "indigenous_coin", "foreign_coin",
                    "paper_currency", "token", "precious_metal", "article"],
    "comm_infra": ["courier", "postal_station", "general_postal_service", "road", "bridge", "canal", "port"],
    "records": ["written_record", "script", "phonetic_alphabetic_writing", "non_phonetic_writing",
                "lists_tables_and_classification", "calendar", "practical_literature", "nonwritten_record", "mnemonic_device"],
}
STANDARDS = ["length_measurement_system", "area_measurement_system", "volume_measurement_system", "weight_measurement_system", "time_measurement_system"]
D_VARS = ["polity_population", "polity_territory", "population_of_the_largest_settlement", "settlement_hierarchy"]
PROD_INFRA = ["irrigation_system", "drinking_water_supply_system", "market", "food_storage_site", "mines_or_quarry"]
MILITARY = ["professional_soldier", "professional_military_officer", "military_level"]


def seshat_panels():
    out = S.build_all()
    facts, agg, meta, crisis, pre = out["facts"], out["polity_agg"], out["polity_meta"], out["crisis"], out["precrisis_agg"]
    facts.to_parquet(PROC / "seshat_facts.parquet", index=False)
    agg.to_parquet(PROC / "seshat_polity_agg.parquet", index=False)
    meta.to_parquet(PROC / "seshat_polity_meta.parquet", index=False)
    crisis.to_parquet(PROC / "seshat_crisis.parquet", index=False)
    pre.to_parquet(PROC / "seshat_precrisis_agg.parquet", index=False)
    wide = S.wide(agg, "value_mean"); wide_log = S.wide(agg, "logv_mean")
    wide.to_parquet(PROC / "seshat_polity_wide.parquet")
    # coverage table: per variable, polities coded, facts, dated facts, inferred share
    cov = agg.groupby("variable").agg(polities_coded=("polity_id", "nunique"), facts=("n_facts", "sum"),
                                      dated_facts=("n_dated", "sum"), inferred_facts=("n_inferred", "sum"),
                                      disputed_facts=("n_disputed", "sum")).reset_index()
    sect = facts.drop_duplicates("variable").set_index("variable").section
    cov["section"] = cov.variable.map(sect)
    cov["block"] = cov.variable.map({v: b for b, vs in K_BLOCKS.items() for v in vs}).fillna(
        cov.variable.map({v: "D_scale" for v in D_VARS}).fillna(cov.variable.map({v: "standards" for v in STANDARDS})))
    cov = cov.sort_values(["block", "polities_coded"], ascending=[True, False])
    cov.to_csv(DIAG / "seshat_coverage_by_variable.csv", index=False)
    # Warfare and other sheets: coverage
    warfare_vars = facts[facts.sheet == "Warfare"].variable.unique().tolist()
    # missingness among the pre-registered indicator set, per polity
    ind = sum(K_BLOCKS.values(), []) + D_VARS + STANDARDS + PROD_INFRA + MILITARY
    pres = wide.reindex(columns=ind)
    miss = pd.DataFrame({"indicator": ind, "share_polities_missing": pres.isna().mean().values})
    miss.to_csv(DIAG / "seshat_missingness_by_indicator.csv", index=False)
    per_pol = pd.DataFrame({"polity_id": pres.index, "n_K_indicators_coded": pres[sum(K_BLOCKS.values(), [])].notna().sum(axis=1).values,
                            "n_D_coded": pres[D_VARS].notna().sum(axis=1).values})
    per_pol = per_pol.merge(meta[["polity_id", "macro_region", "era", "start_year", "end_year"]], on="polity_id")
    per_pol.to_csv(DIAG / "seshat_missingness_by_polity.csv", index=False)
    # temporal distribution figure (polity time spans by macro-region) and crisis start years
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2), gridspec_kw={"width_ratios": [2, 1]})
    regs = sorted(meta.macro_region.unique())
    for i, r in enumerate(regs):
        sub = meta[meta.macro_region == r]
        axes[0].hlines([i + np.random.RandomState(1).uniform(-0.35, 0.35) for _ in range(len(sub))], sub.start_year, sub.end_year, lw=1.2, alpha=0.6)
    axes[0].set_yticks(range(len(regs))); axes[0].set_yticklabels([r.replace("_", " ") for r in regs], fontsize=8)
    axes[0].set_xlabel("Year (negative = BCE)"); axes[0].set_title(f"Polity time spans, Polaris-2026 (n = {len(meta)})", fontsize=10)
    axes[1].hist(crisis.crisis_start.dropna(), bins=30, color="#555"); axes[1].set_xlabel("Crisis start year"); axes[1].set_title(f"CrisisDB records (n = {len(crisis)})", fontsize=10)
    fig.tight_layout(); fig.savefig(FIG / "fig_seshat_temporal_coverage.pdf"); plt.close(fig)
    # missingness heatmap for pre-registered indicators
    fig, ax = plt.subplots(figsize=(10, 5))
    order = pres[sum(K_BLOCKS.values(), []) + D_VARS].notna()
    order = order.loc[order.sum(axis=1).sort_values(ascending=False).index]
    ax.imshow(order.values.T, aspect="auto", cmap="Greys", interpolation="nearest")
    ax.set_yticks(range(order.shape[1])); ax.set_yticklabels(order.columns, fontsize=6); ax.set_xlabel("Polities (sorted by number of coded indicators)")
    ax.set_title("Coded (dark) vs missing (white) indicators, Seshat Polaris-2026", fontsize=10)
    fig.tight_layout(); fig.savefig(FIG / "fig_seshat_missingness.pdf"); plt.close(fig)
    # coverage summary numbers
    summary = {
        "polities": int(len(meta)), "regions": int(meta.seshat_region.nunique()), "macro_regions": int(meta.macro_region.nunique()),
        "facts_total": int(len(facts)), "facts_dated": int(facts.dated.sum()), "facts_unparsed": int(facts.unparsed.sum()),
        "facts_inferred": int(facts.inferred.sum()), "facts_disputed_code": int(facts.disputed.sum()),
        "polities_with_population": int(wide.polity_population.notna().sum()), "polities_with_territory": int(wide.polity_territory.notna().sum()),
        "polities_with_any_D": int(pres[D_VARS].notna().any(axis=1).sum()),
        "polities_with_ge50pct_K": int((pres[sum(K_BLOCKS.values(), [])].notna().mean(axis=1) >= 0.5).sum()),
        "crisis_records": int(len(crisis)), "crisis_polities": int(crisis.polity_id.nunique()),
        "crisis_y1_coded": int(crisis.y1_severity.notna().sum()), "crisis_y1b_coded": int(crisis.y1b_severe.notna().sum()),
        "crisis_y1b_positive": int((crisis.y1b_severe == 1).sum()),
        "y3_coded": int(meta.y3_discontinuous.notna().sum()), "y3_positive": int((meta.y3_discontinuous == 1).sum()),
        "duration_censored_ge1900": int(meta.censored.sum()), "warfare_indicators": len(warfare_vars),
        "era_counts": meta.era.value_counts().to_dict(), "macro_region_counts": meta.macro_region.value_counts().to_dict(),
    }
    return summary, cov


def dplace_panel():
    out = D.build()
    soc = out["societies"]
    soc.to_parquet(PROC / "dplace_ea_societies.parquet", index=False)
    out["codes"].to_csv(META / "dplace_ea_codes_used.csv", index=False)
    cov = pd.DataFrame({"variable": [c for c in soc.columns if c.startswith("EA")], })
    cov["societies_coded"] = [int(soc[c].notna().sum()) for c in cov.variable]
    cov.to_csv(DIAG / "dplace_coverage_by_variable.csv", index=False)
    fig, ax = plt.subplots(figsize=(6, 3.5))
    ax.hist(soc.main_focal_year.dropna(), bins=40, color="#555"); ax.set_xlabel("Main focal year of the ethnographic description"); ax.set_title(f"D-PLACE EA v3.3 societies (n = {len(soc)})", fontsize=10)
    fig.tight_layout(); fig.savefig(FIG / "fig_dplace_focal_years.pdf"); plt.close(fig)
    return {"societies": int(len(soc)), "with_EA031": int(soc.ea031_code.notna().sum()), "with_EA202": int(soc.ea202_code.notna().sum()),
            "with_EA033": int(soc.EA033.notna().sum()), "with_glottolog_family": int(soc.Family_ID.notna().sum()),
            "families": int(soc.Family_ID.nunique()), "regions": int(soc.region.nunique()),
            "focal_year_range": [int(soc.main_focal_year.min()), int(soc.main_focal_year.max())]}


def modern_panel():
    out = M.build(); p = out["panel"]
    p.to_parquet(PROC / "modern_panel.parquet", index=False)
    cov = pd.DataFrame({"variable": [c for c in p.columns if c not in ("country_name", "country_text_id", "country_id", "year")]})
    cov["country_years_coded"] = [int(p[c].notna().sum()) for c in cov.variable]
    cov["first_year"] = [int(p.year[p[c].notna()].min()) if p[c].notna().any() else None for c in cov.variable]
    cov["last_year"] = [int(p.year[p[c].notna()].max()) if p[c].notna().any() else None for c in cov.variable]
    cov["countries"] = [int(p.country_text_id[p[c].notna()].nunique()) for c in cov.variable]
    cov.to_csv(DIAG / "modern_coverage_by_variable.csv", index=False)
    fig, ax = plt.subplots(figsize=(7, 3.5))
    for c, lab in [("maddison_pop", "Maddison population"), ("v2clrspct", "V-Dem K items"), ("e_civil_war", "civil war (COW)"), ("e_coups", "coups"), ("v2regendtype", "regime end type")]:
        ax.plot(p.groupby("year")[c].apply(lambda s: s.notna().sum()), label=lab, lw=1.2)
    ax.set_ylabel("countries coded"); ax.set_xlabel("year"); ax.legend(fontsize=7); ax.set_title("Modern panel coverage by year", fontsize=10)
    fig.tight_layout(); fig.savefig(FIG / "fig_modern_coverage.pdf"); plt.close(fig)
    cw = out["crosswalk"]
    return {"country_years": int(len(p)), "countries": int(p.country_text_id.nunique()), "years": [int(p.year.min()), int(p.year.max())],
            "with_maddison_pop": int(p.maddison_pop.notna().sum()), "countries_matched_maddison": int(cw[cw.in_vdem & cw.in_maddison].shape[0]),
            "vdem_countries_unmatched": int(cw[cw.in_vdem & ~cw.in_maddison].shape[0]),
            "y4_onsets": int((p.y4_civilwar_onset == 1).sum()), "y4_coded": int(p.y4_civilwar_onset.notna().sum()),
            "y5_coups": int((p.y5_coup == 1).sum()), "y5_coded": int(p.y5_coup.notna().sum()),
            "y6_irregular_ends": int((p.y6_irregular_regime_end == 1).sum()), "y6_coded": int(p.y6_irregular_regime_end.notna().sum())}


def main():
    s, cov = seshat_panels(); d = dplace_panel(); m = modern_panel()
    summary = {"built_at_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"), "seshat": s, "dplace": d, "modern": m}
    (DIAG / "coverage_summary.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()