"""mapping.py

Draws a choropleth from a GeoJSON file and a table of values, using only
`json` and `matplotlib`. No GIS package, no shapefile reader, no projection
library. The geometry ships pre-projected to EPSG:5070 Albers Equal Area
Conic, so the drawing code sets an equal aspect and stops.

The classification schemes below all describe the same numbers. They produce
different maps, and choosing among them is a cartographic judgment the
student makes and defends. The program will not make it.

    from mapping import choropleth, classify, CLASS_SCHEMES
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib
import numpy as np

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

CLASS_SCHEMES = ("equal_interval", "quantile", "natural_breaks", "std_dev")


# --------------------------------------------------------------------------
# Classification. Every scheme returns k+1 break edges.
# --------------------------------------------------------------------------
def classify(values, scheme: str = "quantile", k: int = 5):
    v = np.asarray(values, dtype=float)
    v = v[np.isfinite(v)]
    if scheme == "equal_interval":
        return np.linspace(v.min(), v.max(), k + 1)
    if scheme == "quantile":
        return np.quantile(v, np.linspace(0, 1, k + 1))
    if scheme == "std_dev":
        # k classes need k+1 edges, so k-1 interior breaks sit at half
        # standard deviation steps either side of the mean, clipped into
        # the observed range so the outer classes stay drawable.
        m, s = v.mean(), v.std()
        interior = [m + (i - k / 2) * s for i in range(1, k)]
        interior = [min(max(e, v.min()), v.max()) for e in interior]
        return np.array([v.min()] + interior + [v.max()])
    if scheme == "natural_breaks":
        return _kmeans_breaks(v, k)
    raise ValueError(scheme)


def _kmeans_breaks(v, k, iters=60, seed=20260824):
    """One dimensional k-means, the usual stand-in for Fisher-Jenks."""
    rng = np.random.default_rng(seed)
    c = np.quantile(v, np.linspace(0.1, 0.9, k))
    for _ in range(iters):
        lab = np.abs(v[:, None] - c[None, :]).argmin(1)
        new = np.array([v[lab == i].mean() if (lab == i).any() else c[i]
                        for i in range(k)])
        if np.allclose(new, c):
            break
        c = np.sort(new)
    mids = (c[:-1] + c[1:]) / 2
    return np.concatenate([[v.min()], mids, [v.max()]])


def assign(values, edges):
    """Class index 0..k-1 for each value; NaN stays NaN."""
    v = np.asarray(values, dtype=float)
    out = np.full(len(v), np.nan)
    ok = np.isfinite(v)
    idx = np.clip(np.searchsorted(edges[1:-1], v[ok], side="right"),
                  0, len(edges) - 2)
    out[ok] = idx
    return out


# --------------------------------------------------------------------------
# Drawing.
# --------------------------------------------------------------------------
def load_geometry(path):
    obj = json.loads(Path(path).read_text())
    geoms = {}
    for f in obj["features"]:
        geoms[f["properties"]["id"]] = [
            np.asarray(ring[0], dtype=float)
            for ring in f["geometry"]["coordinates"]]
    return geoms


def choropleth(geojson_path, ids, values, *, scheme="quantile", k=5,
               title="", subtitle="", legend_label="", source="",
               cmap="YlOrBr", missing="#d9d9d9", out=None, ax=None):
    geoms = load_geometry(geojson_path)
    values = np.asarray(values, dtype=float)
    edges = classify(values, scheme, k)
    cls = assign(values, edges)
    colors = plt.get_cmap(cmap, k)

    polys, facecolors = [], []
    for gid, c in zip(ids, cls):
        for ring in geoms.get(str(gid), []):
            polys.append(ring)
            facecolors.append(missing if not np.isfinite(c)
                              else colors(int(c)))

    own = ax is None
    if own:
        fig, ax = plt.subplots(figsize=(9, 6.5))
    ax.add_collection(PolyCollection(polys, facecolors=facecolors,
                                     edgecolors="white", linewidths=0.25))
    ax.autoscale_view()
    ax.set_aspect("equal")          # Albers is metric; never squash it
    ax.axis("off")
    if title:
        ax.set_title(title, fontsize=12, loc="left", fontweight="bold")
    if subtitle:
        ax.text(0, 1.005, subtitle, transform=ax.transAxes, fontsize=9,
                va="bottom", color="#444444")
    handles = [plt.Rectangle((0, 0), 1, 1, fc=colors(i), ec="white")
               for i in range(k)]
    labels = [f"{edges[i]:,.0f} to {edges[i+1]:,.0f}" for i in range(k)]
    if not np.isfinite(values).all():
        handles.append(plt.Rectangle((0, 0), 1, 1, fc=missing, ec="white"))
        labels.append("no estimate")
    # The legend sits outside the data extent. Placing it inside collides
    # with the polygons, and the collision moves as the study area changes,
    # so anchoring beyond the axes is the only placement that survives a
    # change of state.
    ax.legend(handles, labels, title=legend_label, loc="upper left",
              bbox_to_anchor=(1.01, 1.0), fontsize=8, title_fontsize=8,
              frameon=False)
    if source:
        ax.text(0, -0.04, source, transform=ax.transAxes, fontsize=7,
                va="top", color="#666666")
    if own and out:
        fig.savefig(out, dpi=150, bbox_inches="tight")
        plt.close(fig)
    return edges, cls