"""Parse the D-PLACE Ethnographic Atlas CLDF release (v3.3) into one row per society with the variables fixed in
research/preregistration.md §2.2, plus Glottolog language family for Galton controls.

Ordinal categories of EA031 (mean size of local communities) and EA202 (population size) are converted to
log-midpoints of the category ranges printed in `codes.csv` (parsed from the code descriptions, not typed by hand);
the parsed table is written to results/diagnostics/dplace_code_midpoints.csv for inspection.
"""
from __future__ import annotations
import io, re, zipfile
from pathlib import Path
import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[3]
ZIP = ROOT / "data" / "raw" / "dplace_ea_v3_3" / "dplace-dataset-ea-v3.3.zip"
GLOTTO = ROOT / "data" / "raw" / "glottolog_cldf_languages" / "glottolog_v5.3_languages.csv"
DIAG = ROOT / "results" / "diagnostics"

VARS = ["EA031", "EA202", "EA030", "EA033", "EA032", "EA072", "EA073", "EA090", "EA066", "EA067", "EA068",
        "EA070", "EA005", "EA028", "EA042", "EA034", "EA113"]


def _read(z: zipfile.ZipFile, name: str) -> pd.DataFrame:
    path = [n for n in z.namelist() if n.endswith("cldf/" + name)][0]
    return pd.read_csv(io.BytesIO(z.read(path)), low_memory=False)


def _range_midpoint(desc: str):
    """'1,000-5,000' -> geometric midpoint; '50,000 or more' -> 50,000 (lower bound, flagged); '<50' -> 25."""
    s = desc.replace(",", "")
    m = re.search(r"(\d+)\s*[-–]\s*(\d+)", s)
    if m:
        lo, hi = float(m.group(1)), float(m.group(2))
        return np.sqrt(lo * hi), lo, hi, "range"
    m = re.search(r"(\d+)\s*(or more|\+|and over|or over)", s)
    if m:
        lo = float(m.group(1)); return lo, lo, np.nan, "open_upper"
    m = re.search(r"(fewer than|less than|under|<)\s*(\d+)", s, flags=re.I)
    if m:
        hi = float(m.group(2)); return hi / 2, np.nan, hi, "open_lower"
    return np.nan, np.nan, np.nan, "none"


def build() -> dict[str, pd.DataFrame]:
    z = zipfile.ZipFile(ZIP)
    soc = _read(z, "societies.csv"); var = _read(z, "variables.csv"); codes = _read(z, "codes.csv"); data = _read(z, "data.csv")
    d = data[data.Var_ID.isin(VARS)].copy()
    # numeric code value (the ordinal code number) and its description
    # ordinal code number = `ord` column of codes.csv; the "Missing data" code (ord 99) is NaN, never a value
    codes["code_num"] = pd.to_numeric(codes["ord"], errors="coerce")
    codes.loc[codes.Name.eq("Missing data") | codes.ID.str.endswith("-NA"), "code_num"] = np.nan
    cmap = codes.set_index("ID")[["code_num", "Description"]]
    d = d.join(cmap, on="Code_ID")
    wide = d.pivot_table(index="Soc_ID", columns="Var_ID", values="code_num", aggfunc="first")
    desc = d.pivot_table(index="Soc_ID", columns="Var_ID", values="Description", aggfunc="first")
    # midpoints for EA031 and EA202 from code descriptions
    mids = []
    # EA202 (population size) is a continuous variable in the CLDF release: numeric Value, no codes
    ea202 = d[d.Var_ID == "EA202"].set_index("Soc_ID").Value.map(lambda x: pd.to_numeric(x, errors="coerce"))
    wide["EA202_mid"] = ea202.reindex(wide.index) if len(ea202) else np.nan
    wide["EA202"] = wide["EA202_mid"]
    for v in ["EA031"]:
        cc = codes[codes.Var_ID == v].copy()
        parsed = [(_range_midpoint(n) if _range_midpoint(n)[3] != "none" else _range_midpoint(desc)) for n, desc in zip(cc.Name.astype(str), cc.Description.astype(str))]
        cc["mid"] = [p[0] for p in parsed]; cc["lo"] = [p[1] for p in parsed]; cc["hi"] = [p[2] for p in parsed]; cc["kind"] = [p[3] for p in parsed]
        mids.append(cc[["Var_ID", "ID", "Name", "Description", "mid", "lo", "hi", "kind"]])
        m = cc.set_index("ID")["mid"]
        wide[v + "_mid"] = d[d.Var_ID == v].set_index("Soc_ID").Code_ID.map(m).reindex(wide.index)
    DIAG.mkdir(parents=True, exist_ok=True)
    pd.concat(mids).to_csv(DIAG / "dplace_code_midpoints.csv", index=False)
    wide.columns = [str(c) for c in wide.columns]
    out = soc.set_index("ID")[["Name", "Latitude", "Longitude", "Glottocode", "region", "main_focal_year", "xd_id"]].join(wide, how="left")
    g = pd.read_csv(GLOTTO, low_memory=False)[["Glottocode", "Family_ID", "Macroarea"]].drop_duplicates("Glottocode")
    out = out.join(g.set_index("Glottocode"), on="Glottocode")
    out["Family_ID"] = out.Family_ID.fillna(out.Glottocode)  # isolates: own family
    out.index.name = "soc_id"
    out = out.rename(columns={"EA031": "ea031_code", "EA202": "ea202_code"})
    out["log_community_size"] = np.log(out["EA031_mid"])
    out["log_population"] = np.log(out["EA202_mid"])
    out["jurisdictional_hierarchy_beyond"] = out["EA033"] - 1  # ord 1 = acephalous -> 0 levels beyond the community
    out["jurisdictional_hierarchy_local"] = out["EA032"]
    return {"societies": out.reset_index(), "variables": var[var.ID.isin(VARS)][["ID", "Name", "Description", "type"]],
            "codes": codes[codes.Var_ID.isin(VARS)][["Var_ID", "ID", "Name", "Description"]], "desc": desc.reset_index()}