"""common.py

Shared plumbing for the exercise programs: snapshot loading, the provenance
block every exercise prints, and the handoff package every exercise writes.

Keeping these in one place means an exercise script contains its analysis and
almost nothing else, which matters because students are asked to read those
scripts.
"""
from __future__ import annotations

import hashlib
import json
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parent.parent
SNAP, FIG, OUT = ROOT / "snapshots", ROOT / "figures", ROOT / "handoff"
SEED = 20260824
ALPHA = 0.05


def ensure_dirs():
    for d in (FIG, OUT):
        d.mkdir(exist_ok=True)


def sha256(path) -> str:
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def banner(title, source, snapshot, extra=None):
    print("=" * 74)
    print(title.upper())
    print(f"  source   {source}")
    p = SNAP / snapshot
    print(f"  snapshot {snapshot}  sha256 {sha256(p)[:32]}")
    print(f"  seed     {SEED}")
    for k, v in (extra or {}).items():
        print(f"  {k:<8s} {v}")
    print("=" * 74)


def load_acs():
    """Mississippi ACS snapshot with the derived shares every exercise uses."""
    m = pd.read_csv(SNAP / "acs2023_mississippi_asi.csv",
                    dtype={"code": str, "sumlev": str})
    lev = {"050": "county", "140": "tract", "150": "block group",
           "860": "ZCTA"}
    age = [c for c in m.columns
           if c.startswith("B01001_E") and c != "B01001_E001"]
    m["level"] = m["sumlev"].map(lev)
    m["pop"] = m["B01001_E001"]
    m["pop65"] = m[age].sum(axis=1, min_count=1)
    m["hh"] = m["B25044_E001"]
    m["hh_nv"] = m["B25044_E003"] + m["B25044_E010"]
    m["pov_univ"] = m["C17002_E001"]
    m["pov_below"] = m["C17002_E002"] + m["C17002_E003"]
    m["pct65"] = 100 * m["pop65"] / m["pop"]
    m["pct_nv"] = 100 * m["hh_nv"] / m["hh"]
    m["pct_pov"] = 100 * m["pov_below"] / m["pov_univ"]
    m["medinc"] = m["B19013_E001"].where(m["B19013_E001"] > 0)
    m["medinc_moe"] = m["B19013_M001"].where(m["B19013_M001"] > 0)
    return m


def load_azmet():
    return pd.read_csv(SNAP / "azmet_2024_annual.csv")


def need_index(d, method="rank"):
    cols = ["pct65", "pct_pov", "pct_nv"]
    if method == "rank":
        parts = [d[c].rank(pct=True) for c in cols]
    elif method == "z":
        parts = [(d[c] - d[c].mean()) / d[c].std() for c in cols]
    else:
        parts = [(d[c] - d[c].min()) / (d[c].max() - d[c].min()) for c in cols]
    return sum(parts)


def jaccard(a, b) -> float:
    a, b = set(a), set(b)
    return len(a & b) / len(a | b) if (a | b) else float("nan")


def corr(a, b):
    a, b = np.asarray(a, float), np.asarray(b, float)
    ok = np.isfinite(a) & np.isfinite(b)
    if ok.sum() < 3:
        return float("nan"), int(ok.sum())
    return float(np.corrcoef(a[ok], b[ok])[0, 1]), int(ok.sum())


def write_handoff(stem, table: pd.DataFrame, decisions: dict, questions):
    """The last mile package. Every exercise ends here."""
    ensure_dirs()
    csv = OUT / f"{stem}_for_review.csv"
    table.to_csv(csv, index=False)
    payload = {
        "decisions_the_machine_made_without_being_asked": decisions,
        "what_a_human_must_check_before_this_is_published": list(questions),
        "review_table": csv.name,
        "rows": int(len(table)),
    }
    (OUT / f"{stem}_decisions.json").write_text(json.dumps(payload, indent=2))
    print(f"\nLAST MILE HANDOFF written to {OUT}")
    print(f"    {csv.name}   every unit, every statistic")
    print(f"    {stem}_decisions.json   what the machine chose, and what a "
          f"human must check")


def save_checkvalues(stem, values: dict):
    p = SNAP / f"{stem}_checkvalues.json"
    p.write_text(json.dumps(values, indent=2, default=str))
    print(f"\ncheck values written to {p}")
    return p