"""exercise_3a.py

Exercise III-A. The Spatial Exploitability Assessment.

Builds a need index from documented Census inputs, treats it as a system
under audit, finds the cheapest single change that moves a tract out of the
highest scoring decile, then applies that change to every top decile tract at
once. Because the index is a percentile, the universal case reorders the list
in a way no individual actor predicts.

    python exercise_3a.py
"""
from __future__ import annotations

import sys
from pathlib import Path

import matplotlib
import numpy as np
import pandas as pd

matplotlib.use("Agg")
import matplotlib.pyplot as plt                        # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parent))
import common as C                                     # noqa: E402
from mapping import choropleth, load_geometry          # noqa: E402
from matplotlib.collections import PolyCollection      # noqa: E402

GEO = C.SNAP / "ms_tracts.geojson"
INPUTS = ["pct65", "pct_pov", "pct_nv"]
DECILE = 0.90
BUDGETS = [0.05, 0.10, 0.20, 0.30, 0.50]   # fractional reduction allowed


def index_of(frame):
    return sum(frame[c].rank(pct=True) for c in INPUTS)


def min_change(d, i, col, threshold):
    """Smallest fractional cut in one input that drops tract i below the bar.

    Bisection on the multiplier. The index is recomputed across every tract
    at each step, because a percentile rank is relative and moving one tract
    moves the others.
    """
    lo, hi = 0.0, 1.0
    base = d[col].to_numpy(float).copy()
    # Feasibility is tested at the largest cut available, which drives the
    # input to zero. A tract that still clears the bar with one input at zero
    # cannot escape by moving that input at all.
    if not _below(d, i, col, base, 1.0, threshold):
        return None
    for _ in range(24):
        mid = 0.5 * (lo + hi)
        if _below(d, i, col, base, mid, threshold):
            hi = mid
        else:
            lo = mid
    return hi


def _below(d, i, col, base, frac, threshold):
    v = base.copy()
    v[i] = base[i] * (1.0 - frac)
    tmp = d.copy()
    tmp[col] = v
    return index_of(tmp).iloc[i] < threshold


def main():
    C.ensure_dirs()
    m = C.load_acs()
    tr = (m[m.level == "tract"].dropna(subset=INPUTS)
          .reset_index(drop=True).copy())
    C.banner("exercise III-A: the spatial exploitability assessment",
             "ACS 2023 5-year, Mississippi, keyless summary file",
             "acs2023_mississippi_asi.csv", {"tracts": len(tr)})
    out = {"tracts": int(len(tr))}

    tr["idx"] = index_of(tr)
    thresh = tr["idx"].quantile(DECILE)
    top = tr.index[tr["idx"] >= thresh].tolist()
    print("\n[1] THE INDEX")
    print(f"    inputs           {', '.join(INPUTS)}")
    print(f"    construction     sum of three percentile ranks, equal weight")
    print(f"    top decile bar   {thresh:.4f}")
    print(f"    tracts in top decile  {len(top)}")
    out["threshold"] = round(float(thresh), 4)
    out["top_decile_n"] = len(top)

    # ---- 2. leverage -----------------------------------------------------
    print("\n[2] LEVERAGE BY INPUT")
    print(f"    {'input':<10}{'escapable':>11}{'median cut':>13}"
          f"{'mean cut':>11}")
    lev = {}
    for col in INPUTS:
        cuts = [min_change(tr, i, col, thresh) for i in top]
        got = [c for c in cuts if c is not None]
        lev[col] = cuts
        if not got:
            print(f"    {col:<10}{0:>11d}      no tract escapes on this input")
            out.setdefault("leverage", {})[col] = {
                "escapable": 0, "median_cut_pct": None, "mean_cut_pct": None}
            continue
        print(f"    {col:<10}{len(got):>11d}{np.median(got)*100:>12.1f}%"
              f"{np.mean(got)*100:>10.1f}%")
        out.setdefault("leverage", {})[col] = {
            "escapable": len(got),
            "median_cut_pct": round(float(np.median(got)) * 100, 1),
            "mean_cut_pct": round(float(np.mean(got)) * 100, 1)}
    usable = {k: v for k, v in out["leverage"].items()
              if v["median_cut_pct"] is not None}
    if not usable:
        raise SystemExit("no input can move a tract out of the decile")
    best = min(usable, key=lambda c: usable[c]["median_cut_pct"])
    print(f"    cheapest lever   {best} at a median cut of "
          f"{out['leverage'][best]['median_cut_pct']}%")
    out["cheapest_lever"] = best

    # ---- 3. everybody does it -------------------------------------------
    print("\n[3] UNIVERSAL GAMING")
    cut = out["leverage"][best]["median_cut_pct"] / 100.0
    g = tr.copy()
    v = g[best].to_numpy(float).copy()
    for i in top:
        v[i] = v[i] * (1.0 - cut)
    g[best] = v
    g["idx2"] = index_of(g)
    thresh2 = g["idx2"].quantile(DECILE)
    top2 = set(g.index[g["idx2"] >= thresh2].tolist())
    left = sorted(set(top) - top2)
    entered = sorted(top2 - set(top))
    stayed = sorted(set(top) & top2)
    j = C.jaccard(top, top2)
    print(f"    every top decile tract cuts {best} by {cut*100:.1f}%")
    print(f"    left the decile      {len(left)}")
    print(f"    entered the decile   {len(entered)}  (these tracts did "
          f"nothing at all)")
    print(f"    stayed               {len(stayed)}")
    print(f"    Jaccard overlap      {j:.3f}")
    out["universal"] = {"cut_pct": round(cut * 100, 1), "left": len(left),
                        "entered": len(entered), "stayed": len(stayed),
                        "jaccard": round(j, 3)}

    # ---- 4. budget sweep -------------------------------------------------
    print("\n[4] BUDGET SWEEP")
    print(f"    {'budget':>8}{'left':>7}{'entered':>9}{'jaccard':>9}")
    out["budget_sweep"] = []
    for b in BUDGETS:
        gg = tr.copy()
        vv = gg[best].to_numpy(float).copy()
        for i in top:
            vv[i] = vv[i] * (1.0 - b)
        gg[best] = vv
        s2 = gg.assign(i2=index_of(gg))
        t2 = set(s2.index[s2["i2"] >= s2["i2"].quantile(DECILE)].tolist())
        print(f"    {b*100:>7.0f}%{len(set(top)-t2):>7d}"
              f"{len(t2-set(top)):>9d}{C.jaccard(top, t2):>9.3f}")
        out["budget_sweep"].append({
            "budget_pct": b * 100, "left": len(set(top) - t2),
            "entered": len(t2 - set(top)),
            "jaccard": round(C.jaccard(top, t2), 3)})

    # ---- 5. one mitigation ----------------------------------------------
    print("\n[5] MITIGATION: A THRESHOLD FIXED IN ADVANCE")
    fixed_bar = thresh                       # frozen at the pre-gaming value
    t3 = set(g.index[g["idx2"] >= fixed_bar].tolist())
    print(f"    holding the bar at {fixed_bar:.4f} instead of recomputing it")
    print(f"    left {len(set(top)-t3)}   entered {len(t3-set(top))}   "
          f"Jaccard {C.jaccard(top, t3):.3f}")
    print("    A fixed bar admits fewer bystanders, lets more acting tracts "
          "escape, and shrinks the decile.")
    out["mitigation_fixed_bar"] = {
        "left": len(set(top) - t3), "entered": len(t3 - set(top)),
        "jaccard": round(C.jaccard(top, t3), 3)}

    # ---- map --------------------------------------------------------------
    geoms = load_geometry(GEO)
    cat = np.zeros(len(tr), int)             # 0 never in, 1 stayed, 2 left, 3 entered
    for i in stayed:
        cat[i] = 1
    for i in left:
        cat[i] = 2
    for i in entered:
        cat[i] = 3
    colors = {0: "#ececec", 1: "#4d4d4d", 2: "#2166ac", 3: "#b2182b"}
    polys, fc = [], []
    for code, c in zip(tr.code, cat):
        for ring in geoms.get(str(code), []):
            polys.append(ring)
            fc.append(colors[c])
    fig, ax = plt.subplots(figsize=(9, 7))
    ax.add_collection(PolyCollection(polys, facecolors=fc, edgecolors="white",
                                     linewidths=0.2))
    ax.autoscale_view(), ax.set_aspect("equal"), ax.axis("off")
    ax.set_title("Who moves when everybody games the same lever",
                 loc="left", fontweight="bold")
    ax.text(0, 1.005, f"Top decile before and after every top decile tract "
                      f"cuts {best} by {cut*100:.0f} percent",
            transform=ax.transAxes, fontsize=9, va="bottom", color="#444")
    handles = [plt.Rectangle((0, 0), 1, 1, fc=colors[k]) for k in (1, 2, 3, 0)]
    ax.legend(handles, [f"stayed ({len(stayed)})", f"left ({len(left)})",
                        f"entered without acting ({len(entered)})",
                        "never in the decile"],
              loc="upper left", bbox_to_anchor=(1.01, 1.0), fontsize=9,
              frameon=False)
    ax.text(0, -0.03, "Source: U.S. Census Bureau, ACS 2023 5-year summary "
                      f"file. {len(tr)} tracts. EPSG:5070.",
            transform=ax.transAxes, fontsize=7, va="top", color="#666")
    fig.savefig(C.FIG / "exercise_3a_churn.png", dpi=140, bbox_inches="tight")
    plt.close(fig)

    tbl = tr[["code"] + INPUTS + ["idx"]].copy()
    tbl["in_top_decile"] = [i in set(top) for i in tr.index]
    tbl["status_after_universal_gaming"] = [
        {0: "never in", 1: "stayed", 2: "left", 3: "entered"}[c] for c in cat]
    for col in INPUTS:
        tbl[f"min_cut_{col}"] = [lev[col][top.index(i)] if i in set(top)
                                 else None for i in tr.index]
    C.write_handoff("exercise_3a", tbl, {
        "index_inputs": INPUTS,
        "standardization": "percentile rank, equal weight",
        "threshold": f"{DECILE:.0%} of the index distribution, recomputed",
        "perturbation": f"proportional cut in {best}",
        "search": "bisection, 24 iterations",
        "color_ramp": "categorical, unjustified",
    }, [
        "The index weights all three inputs equally. Who decided that, and "
        "what would a defensible weighting look like?",
        "A proportional cut assumes an actor can move the input at all. Which "
        "of the three inputs can a jurisdiction actually change, and over "
        "what timescale?",
        "Tracts entered the decile without acting. Is that acceptable in a "
        "funding formula, and what would you tell one of them?",
        "The fixed threshold let more acting tracts escape and shrank the "
        "decile. Is that a mitigation in your application, and what would "
        "you use in its place?",
    ])
    C.save_checkvalues("exercise_3a", out)
    print(f"\nwrote {C.FIG / 'exercise_3a_churn.png'}")


if __name__ == "__main__":
    main()