"""classification_demo.py

Four maps of one column. Same tracts, same numbers, four classification
schemes, four different pictures. Reports how much of the map changes class
and how much of the top class survives the change, which is the cartographic
form of the decision stability measure used in Part II.
"""
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))
from mapping import CLASS_SCHEMES, assign, choropleth, classify  # noqa: E402

SNAP = Path(__file__).resolve().parent.parent / "snapshots"
FIG = Path(__file__).resolve().parent.parent / "figures"
GEO = SNAP / "ms_tracts.geojson"
K = 5

NICE = {"equal_interval": "Equal interval", "quantile": "Quantile",
        "natural_breaks": "Natural breaks", "std_dev": "Standard deviation"}


def load():
    df = pd.read_csv(SNAP / "acs2023_mississippi_asi.csv",
                     dtype={"code": str, "sumlev": str})
    tr = df[df.sumlev == "140"].copy()
    tr["pct_pov"] = 100 * (tr.C17002_E002 + tr.C17002_E003) / tr.C17002_E001
    return tr[["code", "pct_pov"]].reset_index(drop=True)


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


if __name__ == "__main__":
    FIG.mkdir(exist_ok=True)
    tr = load()
    vals = tr.pct_pov.to_numpy()
    n_ok = int(np.isfinite(vals).sum())

    fig, axes = plt.subplots(2, 2, figsize=(15, 11))
    classes, tops = {}, {}
    for ax, scheme in zip(axes.ravel(), CLASS_SCHEMES):
        edges, cls = choropleth(
            GEO, tr.code, vals, scheme=scheme, k=K, ax=ax,
            title=f"{NICE[scheme]}",
            subtitle="Population below poverty, Mississippi census tracts, "
                     "ACS 2023 5-year",
            legend_label="Percent below poverty",
            source="Source: U.S. Census Bureau, ACS 2023 5-year summary "
                   f"file. {n_ok} tracts. EPSG:5070.")
        classes[scheme] = cls
        tops[scheme] = set(tr.code[cls == K - 1])
    fig.suptitle("One column, four defensible classifications",
                 fontsize=15, fontweight="bold")
    fig.tight_layout()
    fig.savefig(FIG / "classification_four_ways.png", dpi=140,
                bbox_inches="tight")
    plt.close(fig)

    print(f"Mississippi tracts with a poverty estimate: {n_ok}\n")
    print("Highest class membership by scheme")
    for s in CLASS_SCHEMES:
        e = classify(vals, s, K)
        print(f"  {NICE[s]:<22s} n={len(tops[s]):>4d}  "
              f"top class begins at {e[-2]:6.2f} percent")

    print("\nAgreement between schemes")
    print(f"  {'pair':<40s} {'same class':>11s} {'top-class Jaccard':>19s}")
    for i, a in enumerate(CLASS_SCHEMES):
        for b in CLASS_SCHEMES[i + 1:]:
            ok = np.isfinite(classes[a]) & np.isfinite(classes[b])
            same = 100 * float((classes[a][ok] == classes[b][ok]).mean())
            print(f"  {NICE[a]+' vs '+NICE[b]:<40s} {same:>10.1f}% "
                  f"{jaccard(tops[a], tops[b]):>19.3f}")

    allc = np.vstack([classes[s] for s in CLASS_SCHEMES])
    ok = np.isfinite(allc).all(0)
    unstable = int((allc[:, ok].max(0) != allc[:, ok].min(0)).sum())
    ever = set().union(*tops.values())
    always = set.intersection(*tops.values())
    print(f"\n  tracts placed in a different class by at least one scheme: "
          f"{unstable} of {int(ok.sum())} ({100*unstable/ok.sum():.1f}%)")
    print(f"  tracts in the worst class under at least one scheme: {len(ever)}")
    print(f"  tracts in the worst class under every scheme:        "
          f"{len(always)}")
    print(f"\nwrote {FIG / 'classification_four_ways.png'}")