"""make_figures.py, student copy.

Holds the two figures the labs draw, being the Exercise II-B audit panel
and the Exercise III-B reliability maps. The author copy carries the rest,
which draw instructor figures.
"""
import json

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

from matplotlib.collections import PolyCollection        # noqa: E402

from matplotlib.patches import FancyArrowPatch, Rectangle  # noqa: E402

sys.path.insert(0, str(Path(__file__).resolve().parent))

import common as C                                       # noqa: E402

from mapping import assign, classify, load_geometry      # noqa: E402

FIG = C.FIG

NAVY, RED, BLUE, GRAY = "#1f3864", "#b2182b", "#2166ac", "#e8e8e8"

SRC_ACS = ("Source: U.S. Census Bureau, ACS 2023 5-year summary file. "
           "EPSG:5070.")

def finish(fig, ax, source, name, tight=True):
    target = ax if not isinstance(ax, (list, np.ndarray)) else np.ravel(ax)[-1]
    # Offset in points below the axes, so the source line clears the tick
    # labels and the axis label whatever the height of the axes.
    target.annotate(source, xy=(0, 0), xycoords="axes fraction",
                    xytext=(0, -40), textcoords="offset points", fontsize=7,
                    color="#666", va="top")
    if tight:
        fig.tight_layout()
    fig.savefig(FIG / name, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  wrote {name}")

def poly_ax(ax):
    ax.autoscale_view()
    ax.set_aspect("equal")
    ax.axis("off")

def fig_before_after():
    """The consequence figure.

    The delivered analysis and the corrected one recommend the identical ten
    tracts, which is the finding rather than a bug. Income enters as one of
    four equally weighted percentile ranks, and a rank transform compresses a
    five thousand dollar error into almost nothing. The defect is still
    invalid, and it bites the moment the design shifts, so the figure shows
    the dollar error beside the design space where the error starts to matter.
    """
    import cooling_centers_analysis as CC
    bad = CC.build()
    m = C.load_acs()
    tr = m[(m.level == "tract") & (m.code.str.startswith(CC.COUNTY))].copy()
    tr = tr[(tr["pop"] > 0) & (tr.hh > 0) & (tr.pov_univ > 0)]
    j = bad.merge(tr[["code", "medinc"]], left_on="tract", right_on="code")
    err = (j["median_income"] - j["medinc"]).abs()

    def top(col, w, k):
        s = (j.pct65.rank(pct=True) + j.pct_pov.rank(pct=True)
             + j.pct_nv.rank(pct=True) + w * (-j[col]).rank(pct=True))
        return set(j.assign(s=s).nlargest(k, "s")["tract"])

    def jac(a, b):
        return len(a & b) / len(a | b)

    weights = [1, 2, 3, 5, 8, 12]
    wj = [jac(top("median_income", w, 10), top("medinc", w, 10))
          for w in weights]
    sizes = [3, 5, 10, 15, 20, 30]
    sj = [jac(top("median_income", 1, k), top("medinc", 1, k)) for k in sizes]

    fig, axes = plt.subplots(1, 3, figsize=(17, 4.9))
    ax = axes[0]
    ax.hist(err, bins=22, color="#c6dbef", edgecolor="#4292c6")
    ax.axvline(err.mean(), color=RED, lw=2,
               label=f"mean error ${err.mean():,.0f}")
    ax.set_xlabel("absolute error in the aggregated median (dollars)")
    ax.set_ylabel("tracts")
    ax.set_title("The defect is real", loc="left", fontweight="bold",
                 fontsize=11.5)
    ax.text(0, 1.02, f"{100*(err>5000).mean():.0f} percent of tracts off by "
                     f"more than $5,000, worst ${err.max():,.0f}",
            transform=ax.transAxes, fontsize=8.8, va="bottom", color="#444")
    ax.legend(frameon=False, fontsize=9)

    ax = axes[1]
    ax.plot(weights, wj, marker="o", color=NAVY, lw=1.8)
    ax.axhline(1.0, ls="--", color="#999", lw=1)
    ax.annotate("as delivered", (1, 1.0), textcoords="offset points",
                xytext=(8, -16), fontsize=9, color=RED)
    ax.set_ylim(.55, 1.05)
    ax.set_xlabel("weight given to income among four inputs")
    ax.set_ylabel("overlap of the recommended ten")
    ax.set_title("Where it starts to matter", loc="left", fontweight="bold",
                 fontsize=11.5)
    first_w = next((w for w, v in zip(weights, wj) if v < 1.0), None)
    ax.text(0, 1.02, (f"Weighting income {first_w} times changes the "
                      "recommendation" if first_w
                      else "No weighting tested changes the recommendation"),
            transform=ax.transAxes, fontsize=8.8, va="bottom", color="#444")

    ax = axes[2]
    ax.plot(sizes, sj, marker="o", color="#1b7837", lw=1.8)
    ax.axhline(1.0, ls="--", color="#999", lw=1)
    ax.set_ylim(.55, 1.05)
    ax.set_xlabel("how many places the program funds")
    ax.set_ylabel("overlap of the selected set")
    ax.set_title("And where it hides", loc="left", fontweight="bold",
                 fontsize=11.5)
    moved = [k for k, v in zip(sizes, sj) if v < 1.0]
    quiet = [k for k, v in zip(sizes, sj) if v == 1.0]
    ax.text(0, 1.02, (f"Unchanged at {quiet} places, changed at {moved}"),
            transform=ax.transAxes, fontsize=8.8, va="bottom", color="#444")

    fig.suptitle("An invalid method that changes nothing, until the design "
                 "shifts", fontsize=13, fontweight="bold")
    finish(fig, axes, SRC_ACS + f" Hinds County, {len(j)} tracts.",
           "exercise_2b_before_after.png")

def fig_reliability():
    m = C.load_acs()
    tr = m[m.level == "tract"].copy()
    tr["cv"] = 100 * (tr.medinc_moe / 1.645) / tr.medinc
    geoms = load_geometry(C.SNAP / "ms_tracts.geojson")
    fig, axes = plt.subplots(1, 3, figsize=(17.5, 6.6))

    def draw(ax, values, cmap, title, sub, k=5, mask=None):
        vals = values.to_numpy(float)
        edges = classify(vals[np.isfinite(vals)], "quantile", k)
        cls = assign(vals, edges)
        colors = plt.get_cmap(cmap, k)
        polys, fc = [], []
        for code, c, keep in zip(tr.code, cls,
                                 (mask if mask is not None
                                  else np.ones(len(tr), bool))):
            col = ("#d9d9d9" if (not np.isfinite(c) or not keep)
                   else colors(int(c)))
            for ring in geoms.get(str(code), []):
                polys.append(ring), fc.append(col)
        ax.add_collection(PolyCollection(polys, facecolors=fc,
                                         edgecolors="white", linewidths=.25))
        poly_ax(ax)
        ax.set_title(title, loc="left", fontweight="bold", fontsize=11.5)
        ax.text(0, 1.01, sub, transform=ax.transAxes, fontsize=8.8,
                va="bottom", color="#444")
        return edges

    n_fail = int((tr.cv > 30).sum())
    n_sup = int(tr.medinc.isna().sum())
    draw(axes[0], tr.medinc, "YlGnBu", "The estimate",
         "Median household income, the map everyone publishes")
    draw(axes[1], tr.cv, "OrRd", "Its reliability",
         f"Coefficient of variation. {n_fail} tracts exceed 30 percent.")
    draw(axes[2], tr.medinc, "YlGnBu", "What survives a screen",
         f"Units above 30 percent suppressed, with {n_sup} already absent",
         mask=(tr.cv <= 30).to_numpy())
    fig.suptitle("One download, three honest maps", fontsize=13,
                 fontweight="bold")
    axes[0].text(0, -0.03, SRC_ACS + " Gray is no reliable estimate.",
                 transform=axes[0].transAxes, fontsize=7, color="#666",
                 va="top")
    fig.tight_layout()
    fig.savefig(FIG / "exercise_3b_reliability.png", dpi=150,
                bbox_inches="tight")
    plt.close(fig)
    print("  wrote exercise_3b_reliability.png")