"""exercise_1a.py

Exercise I-A. Where uncertainty enters, and the fallacy that follows.

Traces one pipeline from survey response to published ranking, then computes
the same relationship at three geographies, a measure of the scale effect that
bounds what any figure can say about a household, and finds the counties
where the within-county pattern runs opposite to the between-county pattern.

    python exercise_1a.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                         # noqa: E402

GEO_T = C.SNAP / "ms_tracts.geojson"
GEO_C = C.SNAP / "ms_counties.geojson"
MIN_BG = 8          # a county needs this many block groups to be tested

STAGES = [
    ("1. Survey response", "A household answers the American Community "
     "Survey", "sampling error, non-response, recall error"),
    ("2. Weighted estimate", "Responses weight up to a population estimate",
     "weighting model error; the individual response is gone"),
    ("3. Published estimate", "The Bureau publishes an estimate and a margin "
     "of error per unit", "suppression where the sample is too thin"),
    ("4. Geographic unit", "Values attach to a block group, tract or county",
     "the modifiable areal unit problem enters here"),
    ("5. Derived index", "Several variables combine into one score",
     "weighting and standardization choices; the margins stop traveling"),
    ("6. Published ranking", "The score orders places and a decision follows",
     "rank instability; nothing downstream can recover the distribution"),
]


def main():
    C.ensure_dirs()
    m = C.load_acs()
    C.banner("exercise I-A: where uncertainty enters, and the fallacy",
             "ACS 2023 5-year, Mississippi, keyless summary file",
             "acs2023_mississippi_asi.csv")
    bg = m[m.level == "block group"].dropna(subset=["pct_pov", "pct_nv"]).copy()
    tr = m[m.level == "tract"].dropna(subset=["pct_pov", "pct_nv"]).copy()
    co = m[m.level == "county"].dropna(subset=["pct_pov", "pct_nv"]).copy()
    out = {}

    print("\n[1] THE PIPELINE")
    print(f"    {'stage':<24}{'what happens':<52}what enters")
    for name, what, unc in STAGES:
        print(f"    {name:<24}{what:<52}{unc}")
    out["stages"] = [{"stage": a, "what": b, "uncertainty": c}
                     for a, b, c in STAGES]

    print("\n[2] THE SAME RELATIONSHIP AT THREE GEOGRAPHIES")
    rows = []
    for lab, d in (("block group", bg), ("tract", tr), ("county", co)):
        r, n = C.corr(d.pct_pov, d.pct_nv)
        rows.append({"level": lab, "n": n, "r": round(r, 4),
                     "r2": round(r * r, 4)})
        print(f"    {lab:<12} n={n:>5d}  r={r:+.4f}  r squared {r*r:.4f}")
    out["scale"] = rows

    print("\n[3] COUNTIES WHERE THE WITHIN-COUNTY PATTERN REVERSES")
    bg["cty"] = bg["code"].str[:5]
    state_r = rows[2]["r"]
    rev = []
    for cty, g in bg.groupby("cty"):
        if len(g) < MIN_BG:
            continue
        r, n = C.corr(g.pct_pov, g.pct_nv)
        if np.isfinite(r) and np.sign(r) != np.sign(state_r):
            rev.append({"county": cty, "n_bg": int(n), "within_r": round(r, 4)})
    rev.sort(key=lambda a: a["within_r"])
    print(f"    counties tested (at least {MIN_BG} block groups): "
          f"{bg.groupby('cty').size().ge(MIN_BG).sum()}")
    print(f"    counties whose internal relationship runs opposite to the "
          f"between-county relationship of {state_r:+.4f}: {len(rev)}")
    for e in rev[:6]:
        print(f"      county {e['county']}  {e['n_bg']:>3d} block groups  "
              f"r = {e['within_r']:+.4f}")
    out["reversal_count"] = len(rev)
    out["reversals"] = rev
    out["counties_tested"] = int(bg.groupby("cty").size().ge(MIN_BG).sum())

    print("\n[4] WHAT EACH FIGURE SUPPORTS")
    print(f"    The county figure of {state_r:+.4f} describes counties.")
    print(f"    The block group figure of {rows[0]['r']:+.4f} describes "
          f"block groups.")
    print("    Neither describes a household, and the reversal counties show "
          "why.")

    # ---- maps -------------------------------------------------------------
    src = ("Source: U.S. Census Bureau, ACS 2023 5-year summary file. "
           "EPSG:5070.")
    fig, axes = plt.subplots(1, 2, figsize=(15, 6))
    choropleth(GEO_C, co.code, co.pct_pov.to_numpy(), scheme="quantile", k=5,
               ax=axes[0], title="County level",
               subtitle="Percent below poverty",
               legend_label="Percent", source=src)
    choropleth(GEO_T, tr.code, tr.pct_pov.to_numpy(), scheme="quantile", k=5,
               ax=axes[1], title="Tract level",
               subtitle="Percent below poverty",
               legend_label="Percent", source=src)
    fig.suptitle("The same variable at two geographies",
                 fontsize=14, fontweight="bold")
    fig.tight_layout()
    fig.savefig(C.FIG / "exercise_1a_two_geographies.png", dpi=140,
                bbox_inches="tight")
    plt.close(fig)

    tbl = tr[["code", "pct_pov", "pct_nv", "pct65", "medinc"]].copy()
    C.write_handoff("exercise_1a", tbl, {
        "geographies_compared": "block group, tract, county",
        "minimum_block_groups_for_a_county_test": MIN_BG,
        "correlation": "Pearson, unweighted by population",
        "classification": "quantile, five classes",
        "color_ramp": "YlOrBr, unjustified",
        "suppressed_units": "dropped silently before correlating",
    }, [
        "The correlation is unweighted, so a block group of 300 people counts "
        "as much as one of 3,000. Should it?",
        "Units with a suppressed estimate were dropped. Are they dropped at "
        "random, or are the poorest units the ones most often suppressed?",
        "Which of the two maps would you send to a legislator, and what would "
        "you say about the other?",
        "The reversal counties carry a story the state figure hides. Does "
        "local knowledge explain any of them?",
    ])
    C.save_checkvalues("exercise_1a", out)
    print(f"\nwrote {C.FIG / 'exercise_1a_two_geographies.png'}")


if __name__ == "__main__":
    main()