"""exercise_5a.py

Exercise V-A. The Proxy Integrity Assessment.

Median household income is the usual stand-in for need. This measures where
the substitution holds and where it fails, and it finds that a strong global
correlation coexists with substantial disagreement among exactly the units a
program would fund.

    python exercise_5a.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 load_geometry                      # noqa: E402
from matplotlib.collections import PolyCollection      # noqa: E402

GEO = C.SNAP / "ms_tracts.geojson"
TOPN = 50
CV_LIMIT = 30.0


def main():
    C.ensure_dirs()
    m = C.load_acs()
    tr = (m[m.level == "tract"]
          .dropna(subset=["pct65", "pct_pov", "pct_nv", "medinc"])
          .reset_index(drop=True).copy())
    C.banner("exercise V-A: the proxy integrity assessment",
             "ACS 2023 5-year, Mississippi, keyless summary file",
             "acs2023_mississippi_asi.csv", {"tracts": len(tr)})
    out = {"tracts": int(len(tr)), "top_n": TOPN}

    tr["referent"] = C.need_index(tr)
    # The proxy is low income, so it is inverted to point the same way as need.
    tr["proxy"] = -tr["medinc"]

    r, n = C.corr(tr.proxy, tr.referent)
    print("\n[1] THE PROXY AGAINST ITS REFERENT")
    print(f"    Pearson r across all {n} tracts       {r:+.4f}")
    rs = float(tr.proxy.rank().corr(tr.referent.rank(), method="pearson"))
    print(f"    Spearman rank correlation             {rs:+.4f}")
    print("    On this evidence the substitution looks defensible.")
    out["global"] = {"pearson": round(r, 4), "spearman": round(rs, 4), "n": n}

    print("\n[2] WHERE THE DECISION HAPPENS")
    by_proxy = set(tr.nsmallest(TOPN, "medinc")["code"])
    by_ref = set(tr.nlargest(TOPN, "referent")["code"])
    j = C.jaccard(by_proxy, by_ref)
    print(f"    lowest {TOPN} by income and highest {TOPN} by the index")
    print(f"    tracts chosen by both                 {len(by_proxy & by_ref)}")
    print(f"    chosen by the proxy alone             "
          f"{len(by_proxy - by_ref)}")
    print(f"    chosen by the referent alone          "
          f"{len(by_ref - by_proxy)}")
    print(f"    Jaccard overlap                       {j:.3f}")
    print(f"    A correlation of {r:+.3f} and an overlap of {j:.3f} describe "
          f"the same two variables.")
    out["decision"] = {"both": len(by_proxy & by_ref),
                       "proxy_only": len(by_proxy - by_ref),
                       "referent_only": len(by_ref - by_proxy),
                       "jaccard": round(j, 3)}

    print("\n[3] WHAT THE PROXY CANNOT SEE")
    only_ref = tr[tr.code.isin(by_ref - by_proxy)].copy()
    only_prx = tr[tr.code.isin(by_proxy - by_ref)].copy()
    print(f"    {'group':<28}{'pct65':>9}{'pct_pov':>10}{'pct_nv':>9}"
          f"{'medinc':>10}")
    for lab, g in (("chosen by referent alone", only_ref),
                   ("chosen by proxy alone", only_prx),
                   ("all tracts", tr)):
        print(f"    {lab:<28}{g.pct65.mean():>9.1f}{g.pct_pov.mean():>10.1f}"
              f"{g.pct_nv.mean():>9.1f}{g.medinc.mean():>10,.0f}")
        out.setdefault("group_means", {})[lab] = {
            "pct65": round(float(g.pct65.mean()), 1),
            "pct_pov": round(float(g.pct_pov.mean()), 1),
            "pct_nv": round(float(g.pct_nv.mean()), 1),
            "medinc": int(g.medinc.mean())}
    print("    The referent-only group is the one the income proxy misses.")

    print("\n[4] IS THE DIVERGENCE REAL OR MEASUREMENT ERROR")
    tr["cv"] = 100 * (tr.medinc_moe / 1.645) / tr.medinc
    ok = tr[tr.cv <= CV_LIMIT].copy()
    by_proxy_r = set(ok.nsmallest(TOPN, "medinc")["code"])
    by_ref_r = set(ok.nlargest(TOPN, "referent")["code"])
    jr = C.jaccard(by_proxy_r, by_ref_r)
    print(f"    restricting to tracts with a coefficient of variation at or "
          f"below {CV_LIMIT:.0f}%")
    print(f"    tracts surviving the screen           {len(ok)} of {len(tr)}")
    print(f"    Jaccard on the screened set           {jr:.3f} "
          f"against {j:.3f} unscreened")
    verdict = ("measurement error explains part of it" if jr > j + 0.05
               else "the divergence survives the screen and is substantive")
    print(f"    {verdict}")
    out["screened"] = {"n": len(ok), "jaccard": round(jr, 3),
                       "verdict": verdict}

    # ---- map ---------------------------------------------------------------
    geoms = load_geometry(GEO)
    colors, polys, fc = {0: "#ececec", 1: "#4d4d4d", 2: "#b2182b",
                          3: "#2166ac"}, [], []
    cat = []
    for code in tr.code:
        if code in by_proxy & by_ref:
            cat.append(1)
        elif code in by_ref - by_proxy:
            cat.append(2)
        elif code in by_proxy - by_ref:
            cat.append(3)
        else:
            cat.append(0)
    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("Where the stand-in and the thing disagree", loc="left",
                 fontweight="bold")
    ax.text(0, 1.005, f"Lowest {TOPN} tracts by median income against highest "
                      f"{TOPN} by a three variable need index",
            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"both ({len(by_proxy & by_ref)})",
                        f"index only ({len(by_ref - by_proxy)})",
                        f"income only ({len(by_proxy - by_ref)})",
                        "neither"],
              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_5a_proxy.png", dpi=140, bbox_inches="tight")
    plt.close(fig)

    tbl = tr[["code", "medinc", "cv", "pct65", "pct_pov", "pct_nv",
              "referent"]].copy()
    tbl["selected_by_proxy"] = tr.code.isin(by_proxy)
    tbl["selected_by_referent"] = tr.code.isin(by_ref)
    C.write_handoff("exercise_5a", tbl, {
        "proxy": "median household income, inverted",
        "referent": "sum of three percentile ranks, equal weight",
        "selection_size": TOPN,
        "reliability_limit": CV_LIMIT,
        "assumption": "the three variable index is treated as ground truth, "
                      "which it is not",
        "color_ramp": "categorical, unjustified",
    }, [
        "The index is treated here as the referent, and it is itself a proxy "
        "for need. What would a real referent be, and could you measure it?",
        "The income-only group and the index-only group differ in composition. "
        "Which group would a program rather serve, and who decides?",
        "A proxy is chosen for availability as much as validity. What did "
        "income cost to obtain against the three variable index?",
        "If the divergence survives the reliability screen, the proxy is "
        "wrong rather than noisy. Does your program have a way to notice "
        "that on its own?",
    ])
    C.save_checkvalues("exercise_5a", out)
    print(f"\nwrote {C.FIG / 'exercise_5a_proxy.png'}")


if __name__ == "__main__":
    main()