"""Parse Seshat Polaris-2026 into (i) a long facts table with explicit uncertainty, (ii) polity-level
aggregates, (iii) the CrisisDB outcome table, (iv) polity metadata (macro-region, era, duration, succession).

Coding rules are those frozen in research/preregistration.md §4:
  P, IP -> 1 (IP flagged inferred); A, IA -> 0 (IA flagged inferred); A~P / P~A -> 0.5 (disputed);
  SU, U, unknown, blank, anything else -> NaN (logged in results/diagnostics/seshat_unparsed_values.csv).
Numeric ranges: midpoint (geometric for scale variables that are logged later; arithmetic for levels), bounds kept.
Nothing is forward-filled; nothing is imputed; undated facts are kept as "whole-polity" facts.
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import pandas as pd
import yaml

ROOT = Path(__file__).resolve().parents[3]
RAW = ROOT / "data" / "raw" / "seshat_polaris2026" / "Polaris2026.xlsx"
DIAG = ROOT / "results" / "diagnostics"

# Variables that are numeric in Polaris (everything else in the SC/WF/RT sheets is a presence code)
GEOMETRIC = ["polity_population", "polity_territory", "population_of_the_largest_settlement"]
ARITHMETIC = ["settlement_hierarchy", "administrative_level", "military_level", "religious_level",
              "largest_communication_distance", "fastest_individual_communication"]
NUMERIC = GEOMETRIC + ARITHMETIC

PRESENT = {"P": 1.0, "IP": 1.0, "present": 1.0, "inferred present": 1.0, "inferred P": 1.0}
ABSENT = {"A": 0.0, "IA": 0.0, "absent": 0.0, "inferred absent": 0.0, "inferred A": 0.0}
DISPUTED = {"A~P": 0.5, "P~A": 0.5, "IA~IP": 0.5, "IP~IA": 0.5, "A~IP": 0.5, "IP~A": 0.5, "P~IA": 0.5, "IA~P": 0.5}
INFERRED = {"IP", "IA", "inferred present", "inferred absent", "inferred P", "inferred A"}

# CrisisDB items (pre-registration §3.5)
INSTABILITY_ITEMS = ["decline", "collapse", "downward_mobility", "extermination", "uprising", "revolution",
                     "successful_revolution", "civil_war", "century_plus", "fragmentation", "capital", "conquest",
                     "assassination", "depose"]
SEVERE_ITEMS = ["collapse", "fragmentation", "century_plus", "conquest", "civil_war"]
REFORM_ITEMS = ["constitution", "labor", "unfree_labor", "suffrage", "public_goods", "religion"]


def load_sheets() -> dict[str, pd.DataFrame]:
    xl = pd.ExcelFile(RAW)
    return {s: xl.parse(s) for s in xl.sheet_names}


def code_presence(v) -> tuple[float, bool, bool, bool]:
    """-> (value, inferred, disputed, unparsed)"""
    if v is None or (isinstance(v, float) and np.isnan(v)):
        return np.nan, False, False, False
    s = str(v).strip()
    if s in PRESENT:
        return 1.0, s in INFERRED, False, False
    if s in ABSENT:
        return 0.0, s in INFERRED, False, False
    if s in DISPUTED:
        return 0.5, False, True, False
    if s in {"SU", "U", "unknown", "suspected unknown", "SA", "SP", "DIS", "uncoded", ""}:
        return np.nan, False, False, False
    return np.nan, False, False, True


def _num(v):
    try:
        x = float(str(v).replace(",", ""))
        return x if np.isfinite(x) else np.nan
    except Exception:
        return np.nan


def build_facts(sheets: dict[str, pd.DataFrame], write_diag: bool = True) -> pd.DataFrame:
    frames = []
    for sheet in ["Social complexity", "Warfare", "Religion", "General"]:
        df = sheets[sheet].copy()
        df["sheet"] = sheet
        frames.append(df)
    f = pd.concat(frames, ignore_index=True)
    f = f.rename(columns={"variable_name": "variable"})
    f["dated"] = f.year_from.notna()
    f["year_from"] = pd.to_numeric(f.year_from, errors="coerce")
    f["year_to"] = pd.to_numeric(f.year_to, errors="coerce")
    f.loc[f.year_to.isna() & f.year_from.notna(), "year_to"] = f.year_from
    f["is_numeric"] = f.variable.isin(NUMERIC)
    # presence codes
    coded = f.loc[~f.is_numeric, "value_from"].map(code_presence)
    f["value"] = np.nan; f["inferred"] = False; f["disputed"] = False; f["unparsed"] = False
    f.loc[~f.is_numeric, "value"] = [c[0] for c in coded]
    f.loc[~f.is_numeric, "inferred"] = [c[1] for c in coded]
    f.loc[~f.is_numeric, "disputed"] = [c[2] for c in coded]
    f.loc[~f.is_numeric, "unparsed"] = [c[3] for c in coded]
    # General sheet is categorical text (centralization, succession...): keep raw text in `text`
    f["text"] = np.where(f.sheet == "General", f.value_from.astype(str), None)
    f.loc[f.sheet == "General", ["value", "unparsed"]] = [np.nan, False]
    # numeric
    lo = f.loc[f.is_numeric, "value_from"].map(_num)
    hi = f.loc[f.is_numeric, "value_to"].map(_num)
    hi = hi.where(hi.notna(), lo)
    lo2, hi2 = np.minimum(lo, hi), np.maximum(lo, hi)
    f.loc[f.is_numeric, "value_lo"] = lo2
    f.loc[f.is_numeric, "value_hi"] = hi2
    geo = f.is_numeric & f.variable.isin(GEOMETRIC)
    ari = f.is_numeric & f.variable.isin(ARITHMETIC)
    with np.errstate(divide="ignore", invalid="ignore"):
        f.loc[geo, "value"] = np.sqrt(f.loc[geo, "value_lo"].clip(lower=0) * f.loc[geo, "value_hi"].clip(lower=0))
    f.loc[ari, "value"] = (f.loc[ari, "value_lo"] + f.loc[ari, "value_hi"]) / 2
    # a strictly positive scale variable with value 0 is not a measurement -> NaN (logged)
    zero_scale = geo & (f.value <= 0)
    f.loc[zero_scale, "value"] = np.nan
    f["is_uncertain"] = f.is_uncertain.fillna(False).astype(bool)
    f["is_disputed"] = f.is_disputed.fillna(False).astype(bool)
    if write_diag:
        DIAG.mkdir(parents=True, exist_ok=True)
        f[f.unparsed][["sheet", "polity_id", "variable", "value_from"]].to_csv(DIAG / "seshat_unparsed_values.csv", index=False)
        f[zero_scale][["polity_id", "variable", "value_from", "value_to", "year_from"]].to_csv(DIAG / "seshat_zero_scale_values.csv", index=False)
    keep = ["sheet", "section", "seshat_region", "polity_number", "polity_id", "variable", "value", "value_lo", "value_hi",
            "text", "year_from", "year_to", "dated", "is_numeric", "inferred", "disputed", "unparsed", "is_uncertain", "is_disputed"]
    return f[keep]


def aggregate_polity(facts: pd.DataFrame, before_year: pd.Series | None = None) -> pd.DataFrame:
    """One row per polity x variable. If `before_year` (indexed by polity_id) is given, dated facts are kept only
    when year_to < before_year[polity]; undated facts are always kept (flagged via `n_undated`)."""
    f = facts[facts.value.notna()].copy()
    if before_year is not None:
        by = f.polity_id.map(before_year)
        f = f[(~f.dated) | (f.year_to < by)]
    f["logv"] = np.where(f.variable.isin(GEOMETRIC), np.log(f.value.where(f.value > 0)), np.nan)
    g = f.groupby(["polity_id", "variable"])
    out = g.agg(value_mean=("value", "mean"), value_min=("value", "min"), value_max=("value", "max"),
                logv_mean=("logv", "mean"), lo_min=("value_lo", "min"), hi_max=("value_hi", "max"),
                n_facts=("value", "size"), n_dated=("dated", "sum"), n_inferred=("inferred", "sum"),
                n_disputed=("disputed", "sum"), n_flag_uncertain=("is_uncertain", "sum")).reset_index()
    out["n_undated"] = out.n_facts - out.n_dated
    return out


def wide(agg: pd.DataFrame, stat: str = "value_mean") -> pd.DataFrame:
    return agg.pivot(index="polity_id", columns="variable", values=stat)


def crisis_table(sheets: dict[str, pd.DataFrame]) -> pd.DataFrame:
    cc = sheets["CrisisDB - Crisis consequences"].copy()
    cc["crisis_id"] = [f"{p}_{i}" for i, p in enumerate(cc.polity_id)]
    items = INSTABILITY_ITEMS + REFORM_ITEMS + ["epidemic"]
    for it in items:
        coded = cc[it].map(code_presence)
        cc[it + "_v"] = [c[0] for c in coded]
    inst = cc[[i + "_v" for i in INSTABILITY_ITEMS]]
    cc["y1_n_coded"] = inst.notna().sum(axis=1)
    cc["y1_n_present"] = (inst == 1).sum(axis=1)
    cc["y1_severity"] = np.where(cc.y1_n_coded > 0, cc.y1_n_present / cc.y1_n_coded, np.nan)
    sev = cc[[i + "_v" for i in SEVERE_ITEMS]]
    cc["y1b_severe"] = np.where(sev.notna().all(axis=1), (sev == 1).any(axis=1).astype(float), np.nan)
    ref = cc[[i + "_v" for i in REFORM_ITEMS]]
    cc["y_reform"] = np.where(ref.notna().sum(axis=1) > 0, (ref == 1).sum(axis=1) / ref.notna().sum(axis=1), np.nan)
    cc["epidemic_v"] = cc["epidemic_v"]
    cc["crisis_start"] = pd.to_numeric(cc.year_from, errors="coerce")
    cc["crisis_end"] = pd.to_numeric(cc.year_to, errors="coerce")
    keep = ["crisis_id", "polity_id", "seshat_region", "other_polity_id", "crisis_start", "crisis_end", "y1_n_coded",
            "y1_n_present", "y1_severity", "y1b_severe", "y_reform", "epidemic_v"] + [i + "_v" for i in INSTABILITY_ITEMS]
    return cc[keep]


def polity_metadata(sheets: dict[str, pd.DataFrame], facts: pd.DataFrame) -> pd.DataFrame:
    pol = sheets["Polities"].copy()
    cfg = yaml.safe_load((ROOT / "configs" / "regions.yaml").read_text())
    r2m = {r: m for m, rs in cfg["macro_regions"].items() for r in rs}
    pol["macro_region"] = pol.seshat_region.map(r2m)
    r2c = {r: m for m, rs in cfg["cultural_spheres"].items() for r in rs}
    pol["cultural_sphere"] = pol.seshat_region.map(r2c)
    pol["n_polities_in_nga"] = pol.groupby("homeNGA").polity_id.transform("count")
    pol["midpoint"] = (pol.start_year + pol.end_year) / 2
    def era(y):
        for b in cfg["era_blocks"]:
            if y >= b.get("min", -1e9) and y <= b.get("max", 1e9):
                return b["name"]
    pol["era"] = np.floor(pol.midpoint).map(era)
    pol["duration"] = pol.end_year - pol.start_year
    pol["censored"] = pol.end_year >= 1900
    gen = facts[facts.sheet == "General"]
    def first_text(var):
        s = gen[gen.variable == var].groupby("polity_id").text.first()
        return pol.polity_id.map(s)
    pol["centralization"] = first_text("polity_degree_of_centralization")
    pol["relationship_to_preceding"] = first_text("polity_relationship_to_preceding_entity")
    pol["preceding_entity"] = first_text("polity_preceding_entity")
    pol["succeeding_entity"] = first_text("polity_succeeding_entity")
    pol["suprapolity_relations"] = first_text("polity_suprapolity_relations")
    pol["utm_zone"] = first_text("polity_utm_zone")
    pol["language_family"] = first_text("polity_linguistic_family")
    # Y3: discontinuous succession = relationship_to_preceding of the successor polity (matched by name or id)
    name2id = {str(n).strip().lower(): i for n, i in zip(pol.long_name, pol.polity_id)}
    id_set = set(pol.polity_id)
    def succ_id(s):
        if s is None or str(s) == "nan" or str(s).strip() == "":
            return None
        key = str(s).strip()
        if key in id_set:
            return key
        return name2id.get(key.lower())
    pol["successor_id"] = pol.succeeding_entity.map(succ_id)
    # fallback: the polity that names this one as preceding entity
    prec_map = {}
    for pid, prec in zip(pol.polity_id, pol.preceding_entity):
        sid = succ_id(prec)
        if sid and sid not in prec_map:
            prec_map[sid] = pid
    pol["successor_id"] = pol.successor_id.where(pol.successor_id.notna(), pol.polity_id.map(prec_map))
    rel = dict(zip(pol.polity_id, pol.relationship_to_preceding))
    succ_rel = pol.successor_id.map(rel)
    cont = {"continuity", "continuation"}
    pol["successor_relationship"] = succ_rel
    pol["y3_discontinuous"] = np.where(succ_rel.isna() | succ_rel.isin(["SU", "suspected unknown", "nan", "None"]), np.nan,
                                        (~succ_rel.isin(cont)).astype(float))
    return pol


def build_all() -> dict[str, pd.DataFrame]:
    sheets = load_sheets()
    facts = build_facts(sheets)
    agg = aggregate_polity(facts)
    meta = polity_metadata(sheets, facts)
    crisis = crisis_table(sheets)
    # pre-crisis aggregates: one per crisis record
    pre = []
    for _, r in crisis.iterrows():
        sub = facts[facts.polity_id == r.polity_id]
        a = aggregate_polity(sub, before_year=pd.Series({r.polity_id: r.crisis_start}))
        a["crisis_id"] = r.crisis_id
        pre.append(a)
    pre_agg = pd.concat(pre, ignore_index=True) if pre else pd.DataFrame()
    return {"facts": facts, "polity_agg": agg, "polity_meta": meta, "crisis": crisis, "precrisis_agg": pre_agg}