"""Aggregate police.uk street crime to postcodes by spatial proximity.

Instead of attributing each incident to its published LSOA code, this transform
counts the anonymised incident *points* that fall within ``buffer_m`` (default
50m) of each postcode's boundary polygon (the polygon buffered outward). A point
inside several overlapping buffers counts for each postcode -- the same
multiplicity the tree-density filter uses for features near more than one
postcode. The 50m buffer deliberately smooths police.uk's snap-to-grid
coordinates, which would otherwise make the count hypersensitive to which side of
a narrow line a shared "map point" anchor happened to land on.

One figure is produced for every postcode and crime type, averaged over the
**last 7 years** and the **last 2 years**:

* the **average annual incident count** -- ``sum(counts in covered window) * 12 /
  covered_months`` -- the raw, absolute number of recorded incidents per year,
  with no per-area or per-capita normalisation. A covered postcode with no
  incidents of a type gets ``0``; a postcode whose force never published in the
  window, or whose geometry is unusable, gets a *null* (unknown, never a zero).

This figure is what the server exposes as the filterable crime feature -- the
headline metric in the right pane. (An earlier version divided it by an ambient
daytime population to get a per-1,000-people rate; that was hard to read, so the
absolute per-year count is now used directly.)

**Force-coverage calendar.** police.uk has multi-year publication gaps for whole
forces (Greater Manchester has published nothing between 2019-07 and the present
except 2022-08; BTP, Gloucestershire, Devon & Cornwall and others have shorter
gaps). A missing month is *no data*, not zero crime, so every figure here is
computed against the months the postcode's own force actually published:

* Each postcode is assigned a home force by majority vote of the incidents that
  matched it (BTP, which reports nationwide, is excluded from the vote);
  postcodes with no incidents inherit their outcode's majority force, then the
  national modal force.
* A window's average pools the counts over the force's *covered* months inside
  that window and annualises by those months, so a coverage gap shrinks the data
  rather than reading as a low-crime dip.
* The by-year series only emits bars for years with at least ``min_bar_months``
  covered months (default 6).
* ``covered_years`` (list[struct{year, months}]) is written for every postcode
  so the server can tell "covered, zero crime" from "no data".
* Postcodes whose boundary buffer is unusable (broken geometry) get null
  figures and an empty ``covered_years`` -- unknown, not zero.

Outputs, all keyed on ``postcode``:

* ``crime_by_postcode.parquet`` -- ``"{type} (/yr, 7y)"`` / ``"{type} (/yr, 2y)"``
  average-annual-count columns (the filterable crime features).
* ``crime_by_postcode_by_year.parquet`` -- ``covered_years`` + nested
  ``"{type} (by year)"`` ``list[struct{year, count}]`` per-year raw counts.
* ``crime_records.parquet`` -- one row per counted incident over the last 7
  years (``postcode`` + month/type/location/outcome/lat/lon), sorted by
  postcode so the server can slice each postcode's incidents directly.

Caveat: police.uk coordinates are snapped to a fixed set of anonymous "map
points", not true locations, and a share of rows have no coordinate at all
(dropped here). Spatial totals are therefore fuzzier than the old LSOA-tagged
counts -- by design, not a regression.
"""

from __future__ import annotations

import argparse
import re
import sys
import tempfile
from pathlib import Path

import numpy as np
import polars as pl
import shapely
from pyproj import Transformer

from pipeline.transform.crime import (
    LEGACY_CRIME_TYPE_ALIASES,
    MINOR_CRIME_TYPES,
    SERIOUS_CRIME_TYPES,
    find_street_crime_csvs,
)
from pipeline.transform.postcode_boundaries.loader import load_postcode_polygons

# Serious types first so column order is stable and self-documenting.
ALL_CRIME_TYPES: tuple[str, ...] = SERIOUS_CRIME_TYPES + MINOR_CRIME_TYPES
# Output type axis = the 14 leaf types plus the two rollups, in that order.
ROLLUP_TYPES: tuple[str, ...] = ("Serious crime", "Minor crime")
ALL_OUTPUT_TYPES: tuple[str, ...] = ALL_CRIME_TYPES + ROLLUP_TYPES

DEFAULT_BUFFER_M = 50.0
MONTH_DIR_RE = re.compile(r"^\d{4}-\d{2}$")
STREET_CSV_NAME_RE = re.compile(r"^(\d{4}-\d{2})-(.+)-street\.csv$")

# Trailing-window definitions, in (label, years) form. Each window's average is
# pooled over the force's covered months inside the window; one average-annual-
# count column (`(/yr, <label>)`) — the filterable crime feature — is emitted per
# window.
WINDOWS: tuple[tuple[str, int], ...] = (("7y", 7), ("2y", 2))

# Per-incident records cover the longest window's calendar years so the
# "individual crimes" list reconciles exactly with the headline counts: every
# record is an incident inside that window and vice versa. The headline counts
# are pooled per calendar year, so the records window must be calendar-year
# aligned too -- a rolling month span (e.g. a fixed 84 months) would include
# incidents from a year the headline excludes when the latest month is mid-year.
RECORDS_WINDOW_YEARS = max(win_years for _, win_years in WINDOWS)

# Minimum covered months for a year to get a by-year chart bar (and to be
# listed in `covered_years`). Annualising fewer observed months (x12 from a
# single month at the worst) produces bars dominated by noise. Six months keeps
# the annualisation factor <= 2.
MIN_BAR_MONTHS = 6

# Forces that report nationwide rather than policing a territory. They never
# define a postcode's home force, but their incidents still count.
NON_TERRITORIAL_FORCES = frozenset({"btp"})

COVERAGE_COLUMN = "covered_years"

# Generous GB bounds; points outside fall in no English postcode anyway, but
# filtering first keeps the WGS84->BNG transform out of its undefined region.
LON_BOUNDS = (-9.5, 2.5)
LAT_BOUNDS = (49.0, 61.5)

# Read CSVs in chunks of files to bound peak memory while keeping the STRtree
# query vectorised over a useful number of points. Kept modest because the
# in-window batches also materialise the per-incident record strings.
_CSV_BATCH = 32


def crime_column(type_name: str, window: str) -> str:
    """Filterable average-annual-count column, e.g. ``"Burglary (/yr, 7y)"``."""
    return f"{type_name} (/yr, {window})"


def _force_calendar(
    csvs: list[Path],
) -> tuple[list[int], list[str], np.ndarray]:
    """Derive the per-force publication calendar from the CSV paths.

    File presence IS the coverage signal: a (force, month) with no file
    published nothing. Returns the sorted distinct years, the force slugs
    (sorted), and ``months_in_year_force`` of shape (n_forces, n_years).
    """
    month_force: set[tuple[str, str]] = set()
    for path in csvs:
        if not MONTH_DIR_RE.fullmatch(path.parent.name):
            continue
        m = STREET_CSV_NAME_RE.fullmatch(path.name)
        if m is None or m.group(1) != path.parent.name:
            continue
        month_force.add((m.group(1), m.group(2)))
    if not month_force:
        raise ValueError("No valid YYYY-MM street crime CSVs found")

    years = sorted({int(month[:4]) for month, _ in month_force})
    forces = sorted({force for _, force in month_force})
    year_to_idx = {year: idx for idx, year in enumerate(years)}
    force_to_idx = {force: idx for idx, force in enumerate(forces)}

    months_in_year_force = np.zeros((len(forces), len(years)), dtype=np.int32)
    for month, force in month_force:
        months_in_year_force[force_to_idx[force], year_to_idx[int(month[:4])]] += 1

    all_months = {month for month, _ in month_force}
    for force in forces:
        published = {m for m, f in month_force if f == force}
        missing = len(all_months) - len(published)
        if missing:
            print(
                f"  coverage gap: {force} missing {missing}/{len(all_months)} months"
            )

    return years, forces, months_in_year_force


def _build_tree(
    polygons: np.ndarray, buffer_m: float
) -> tuple[shapely.STRtree, np.ndarray]:
    """Buffer postcode polygons outward by ``buffer_m`` and index them.

    Buffer index == postcode index. Returns the STRtree and a ``usable`` boolean
    mask: geometries that fail to buffer are replaced with an empty polygon (so
    the index stays aligned and they never match) and marked unusable, so their
    crime picture is reported as unknown rather than zero.
    """
    buffers = shapely.buffer(polygons, buffer_m, quad_segs=8)
    broken = shapely.is_missing(buffers) | ~shapely.is_valid(buffers)
    if broken.any():
        print(f"  {int(broken.sum()):,} postcode buffers unusable; left empty")
        buffers[broken] = shapely.from_wkt("POLYGON EMPTY")
    return shapely.STRtree(buffers), ~broken


def _accumulate_counts(
    csvs: list[Path],
    postcodes: np.ndarray,
    tree: shapely.STRtree,
    type_to_idx: dict[str, int],
    year_to_idx: dict[int, int],
    force_to_idx: dict[str, int],
    transformer: Transformer,
    counts: np.ndarray,
    force_votes: np.ndarray,
    records_shard_dir: Path | None,
    records_min_ym: int,
) -> None:
    """Stream the crime CSVs, counting points-in-buffer per (postcode, type, year).

    Also accumulates ``force_votes`` (n_postcodes, n_forces) for the home-force
    election and, when ``records_shard_dir`` is set, writes one parquet shard per
    batch holding every counted (incident, postcode) pair whose month is within
    the records window (month index >= ``records_min_ym``).
    """
    # Type overrides only for the columns we ever read; LSOA is not stored.
    schema = {
        "Longitude": pl.Float64,
        "Latitude": pl.Float64,
        "Month": pl.Utf8,
        "Crime type": pl.Utf8,
        "Location": pl.Utf8,
        "Last outcome category": pl.Utf8,
    }
    years = list(year_to_idx)
    total_points = 0
    total_matches = 0
    total_dropped = 0
    total_records = 0
    unknown_type_counts: dict[str, int] = {}

    for start in range(0, len(csvs), _CSV_BATCH):
        batch = csvs[start : start + _CSV_BATCH]
        path_to_fidx = {}
        batch_max_ym = -1
        for path in batch:
            m = STREET_CSV_NAME_RE.fullmatch(path.name)
            if m is not None:
                ym = m.group(1)
                batch_max_ym = max(batch_max_ym, int(ym[:4]) * 12 + int(ym[5:7]) - 1)
                if m.group(2) in force_to_idx:
                    path_to_fidx[str(path)] = force_to_idx[m.group(2)]
        # The per-incident record strings (Location, outcome) are by far the
        # heaviest columns; read them only for batches that fall inside the
        # records window, so the ~50% of pre-window months cost nothing extra.
        want_records = records_shard_dir is not None and batch_max_ym >= records_min_ym
        record_cols = ["Location", "Last outcome category"] if want_records else []

        frame = (
            pl.scan_csv(
                batch,
                schema_overrides=schema,
                ignore_errors=True,
                include_file_paths="_source_path",
            )
            .select(
                "Longitude", "Latitude", "Month", "Crime type", *record_cols, "_source_path"
            )
            .with_columns(
                pl.col("Month").str.slice(0, 4).cast(pl.Int32, strict=False).alias("year"),
                pl.col("Month").str.slice(5, 2).cast(pl.Int32, strict=False).alias("_mm"),
            )
            .filter(
                pl.col("Longitude").is_not_null()
                & pl.col("Latitude").is_not_null()
                & pl.col("Longitude").is_between(*LON_BOUNDS)
                & pl.col("Latitude").is_between(*LAT_BOUNDS)
                & pl.col("Crime type").is_not_null()
                & (pl.col("Crime type") != "")
                & pl.col("year").is_in(years)
            )
            # year*12 + (month-1): an integer month index for window filtering.
            .with_columns(
                (pl.col("year") * 12 + (pl.col("_mm").fill_null(1) - 1)).alias("month_index")
            )
            .with_columns(pl.col("Crime type").replace(LEGACY_CRIME_TYPE_ALIASES))
            .with_columns(
                pl.col("Crime type")
                .replace_strict(type_to_idx, default=None, return_dtype=pl.Int32)
                .alias("tidx"),
                pl.col("year")
                .replace_strict(year_to_idx, return_dtype=pl.Int32)
                .alias("yidx"),
                pl.col("_source_path")
                .replace_strict(path_to_fidx, default=-1, return_dtype=pl.Int32)
                .alias("fidx"),
            )
            .select(
                "Longitude",
                "Latitude",
                "Crime type",
                *record_cols,
                "month_index",
                "tidx",
                "yidx",
                "fidx",
            )
            .collect(engine="streaming")
        )

        if frame.height == 0:
            continue

        unknown = frame.filter(pl.col("tidx").is_null())
        if unknown.height:
            for name, cnt in unknown.group_by("Crime type").len().iter_rows():
                unknown_type_counts[name] = unknown_type_counts.get(name, 0) + cnt
            frame = frame.filter(pl.col("tidx").is_not_null())
            if frame.height == 0:
                continue

        lon = frame["Longitude"].to_numpy()
        lat = frame["Latitude"].to_numpy()
        tidx = frame["tidx"].to_numpy()
        yidx = frame["yidx"].to_numpy()
        fidx = frame["fidx"].to_numpy()
        month_index = frame["month_index"].to_numpy()

        x, y = transformer.transform(lon, lat)
        finite = np.isfinite(x) & np.isfinite(y)
        total_dropped += int((~finite).sum())
        if not finite.any():
            continue
        x, y = x[finite], y[finite]
        tidx, yidx, fidx = tidx[finite], yidx[finite], fidx[finite]
        total_points += x.size

        points = shapely.points(x, y)
        point_index, postcode_index = tree.query(points, predicate="intersects")
        if point_index.size:
            np.add.at(
                counts,
                (postcode_index, tidx[point_index], yidx[point_index]),
                1,
            )
            matched_fidx = fidx[point_index]
            known_force = matched_fidx >= 0
            if known_force.any():
                np.add.at(
                    force_votes,
                    (postcode_index[known_force], matched_fidx[known_force]),
                    1,
                )
            total_matches += point_index.size

            if want_records:
                total_records += _write_record_shard(
                    records_shard_dir,
                    start,
                    postcodes,
                    point_index,
                    postcode_index,
                    month_index[finite],
                    frame["Crime type"].to_numpy()[finite],
                    frame["Location"].to_numpy()[finite],
                    frame["Last outcome category"].to_numpy()[finite],
                    lon[finite],
                    lat[finite],
                    records_min_ym,
                )

        print(
            f"  files {start + len(batch):,}/{len(csvs):,}: "
            f"{total_points:,} located points, {total_matches:,} postcode matches"
            + (f", {total_records:,} records" if records_shard_dir is not None else "")
        )

    if total_dropped:
        print(f"Dropped {total_dropped:,} points outside the BNG transform domain")
    if unknown_type_counts:
        total_unknown = sum(unknown_type_counts.values())
        listed = ", ".join(
            f"{name!r} ({cnt:,})"
            for name, cnt in sorted(
                unknown_type_counts.items(), key=lambda kv: kv[1], reverse=True
            )
        )
        print(
            f"WARNING: dropped {total_unknown:,} incidents with crime types not in "
            f"ALL_CRIME_TYPES (taxonomy is stale -- update SERIOUS/MINOR_CRIME_TYPES): "
            f"{listed}",
            file=sys.stderr,
        )


def _write_record_shard(
    records_shard_dir: Path,
    start: int,
    postcodes: np.ndarray,
    point_index: np.ndarray,
    postcode_index: np.ndarray,
    month_index: np.ndarray,
    crime_type: np.ndarray,
    location: np.ndarray,
    outcome: np.ndarray,
    lon: np.ndarray,
    lat: np.ndarray,
    records_min_ym: int,
) -> int:
    """Write one parquet shard of (incident, postcode) records for this batch.

    Each matched pair becomes a row -- the same multiplicity as the count -- so a
    postcode's records are exactly the incidents that made up its counts. Only
    incidents within the records window (month index >= ``records_min_ym``) are
    kept. Returns the number of rows written.
    """
    mi = month_index[point_index]
    keep = mi >= records_min_ym
    if not keep.any():
        return 0
    pidx = point_index[keep]
    # Build the string columns from Python lists (.tolist()) rather than numpy
    # object arrays: an all-null Location/outcome slice is an all-None object
    # array that polars cannot cast to String, whereas a list of str|None infers
    # a nullable String column cleanly.
    shard = pl.DataFrame(
        {
            "postcode": postcodes[postcode_index[keep]].astype(str),
            "month_index": mi[keep].astype(np.int32),
            "crime_type": crime_type[pidx].astype(str),
            "location": location[pidx].tolist(),
            "outcome": outcome[pidx].tolist(),
            "lat": lat[pidx].astype(np.float32),
            "lon": lon[pidx].astype(np.float32),
        },
        schema_overrides={
            "postcode": pl.String,
            "crime_type": pl.String,
            "location": pl.String,
            "outcome": pl.String,
        },
    )
    shard.write_parquet(
        records_shard_dir / f"{start:08d}.parquet", compression="zstd"
    )
    return shard.height


def _assign_home_force(
    postcodes: np.ndarray,
    force_votes: np.ndarray,
    forces: list[str],
) -> np.ndarray:
    """Elect each postcode's home (territorial) force by majority incident vote."""
    votes = force_votes.astype(np.int64, copy=True)
    for idx, force in enumerate(forces):
        if force in NON_TERRITORIAL_FORCES:
            votes[:, idx] = 0

    home = votes.argmax(axis=1).astype(np.int32)
    has_vote = votes.max(axis=1) > 0
    home[~has_vote] = -1

    if not has_vote.any():
        raise ValueError("No incidents matched any postcode; cannot assign forces")

    outcodes = np.array([pc.split(" ")[0] for pc in postcodes], dtype=object)
    national_modal = int(np.bincount(home[has_vote], minlength=len(forces)).argmax())
    if (~has_vote).any():
        outcode_modal: dict[str, int] = {}
        voted_outcodes = outcodes[has_vote]
        voted_home = home[has_vote]
        for oc in np.unique(voted_outcodes):
            tally = np.bincount(voted_home[voted_outcodes == oc], minlength=len(forces))
            outcode_modal[oc] = int(tally.argmax())
        fallback = np.array(
            [outcode_modal.get(oc, national_modal) for oc in outcodes[~has_vote]],
            dtype=np.int32,
        )
        home[~has_vote] = fallback
        print(
            f"  {int((~has_vote).sum()):,} postcodes had no territorial incidents; "
            "home force inherited from outcode majority"
        )

    return home


def _window_annualised(
    counts: np.ndarray,
    months_in_year_force: np.ndarray,
    home_fidx: np.ndarray,
    usable: np.ndarray,
    year_mask: np.ndarray,
) -> np.ndarray:
    """Raw annualised incidents/yr per (postcode, type) over a window of years.

    For each force, the count is pooled over the force's covered months that fall
    inside ``year_mask`` and annualised by those covered months. A covered
    postcode with no incidents of a type gets 0; a postcode whose force never
    published in the window, or whose geometry is unusable, gets NaN (unknown).
    """
    n_pc, n_types = counts.shape[0], counts.shape[1]
    avg = np.full((n_pc, n_types), np.nan, dtype=np.float64)
    for f in range(months_in_year_force.shape[0]):
        sel = home_fidx == f
        if not sel.any():
            continue
        cov_months = months_in_year_force[f].astype(np.float64) * year_mask
        denom = cov_months.sum()
        if denom <= 0:
            continue  # force published nothing in this window; stays NaN
        window_years = cov_months > 0
        pooled = counts[sel][:, :, window_years].sum(axis=2, dtype=np.float64)
        avg[sel] = pooled * 12.0 / denom
    avg[~usable] = np.nan
    return avg


def _append_rollups(avg14: np.ndarray) -> np.ndarray:
    """Append Serious/Minor rollup columns (sum of components) -> (n_pc, 16)."""
    serious_idx = [ALL_CRIME_TYPES.index(t) for t in SERIOUS_CRIME_TYPES]
    minor_idx = [ALL_CRIME_TYPES.index(t) for t in MINOR_CRIME_TYPES]
    serious = avg14[:, serious_idx].sum(axis=1)
    minor = avg14[:, minor_idx].sum(axis=1)
    return np.column_stack([avg14, serious, minor])


def _write_crime_table(
    postcodes: np.ndarray,
    raw_by_window: dict[str, np.ndarray],
    output_path: Path,
) -> None:
    """Write the average-annual-count parquet (the filterable crime features).

    ``raw_by_window[label]`` is the (n_postcodes, 16) average annual incident
    count for that window (14 leaf types + 2 rollups). Each value is the raw,
    absolute incidents/yr; a postcode with no usable data for a window keeps NaN
    (written as null).
    """
    data: dict[str, np.ndarray] = {"postcode": postcodes}
    for label, _years in WINDOWS:
        counts = np.round(raw_by_window[label], 1).astype(np.float32)
        for type_idx, name in enumerate(ALL_OUTPUT_TYPES):
            data[crime_column(name, label)] = counts[:, type_idx]

    _write_nan_aware(data, output_path, "postcode crime average annual counts")


def _write_nan_aware(
    data: dict[str, np.ndarray], output_path: Path, label: str
) -> None:
    frame = pl.DataFrame(data)
    value_cols = [c for c in frame.columns if c != "postcode"]
    frame = frame.with_columns(pl.col(c).fill_nan(None) for c in value_cols)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    frame.write_parquet(output_path, compression="zstd")
    print(f"Wrote {label}: {output_path}  {frame.shape}")


def _rollup_long(
    long: pl.DataFrame, types: tuple[str, ...], rollup_name: str
) -> pl.DataFrame:
    """Sum per-year counts across ``types`` into a single rollup."""
    return (
        long.filter(pl.col("Crime type").is_in(list(types)))
        .group_by("postcode", "year")
        .agg(pl.col("count").sum().round(1).alias("count"))
        .with_columns(pl.lit(rollup_name).alias("Crime type"))
        .select("postcode", "Crime type", "year", "count")
    )


def _write_by_year(
    postcodes: np.ndarray,
    counts: np.ndarray,
    years: list[int],
    months_in_year_force: np.ndarray,
    home_fidx: np.ndarray,
    usable: np.ndarray,
    min_bar_months: int,
    output_path: Path,
) -> None:
    """Write nested ``"{type} (by year)"`` raw-count series plus rollups + coverage.

    A bar is only emitted for (postcode, year)s where the postcode's home force
    published at least ``min_bar_months`` months. Bars are the raw annualised
    count for that year (``count * 12 / covered_months``); unlike the headline
    windows there is no per-capita normalisation here -- the chart shows incident
    volume over time. Every postcode gets a ``covered_years`` row so consumers can
    distinguish covered-but-crime-free years from coverage gaps.
    """
    cov_pc_year = months_in_year_force[home_fidx, :]

    annual = np.round(
        counts.astype(np.float64) * 12.0 / np.maximum(cov_pc_year[:, None, :], 1),
        1,
    )
    bar_ok = (
        (counts > 0)
        & (cov_pc_year[:, None, :] >= min_bar_months)
        & usable[:, None, None]
    )

    pc_i, ty_i, yr_i = np.nonzero(bar_ok)

    type_names = np.array(ALL_CRIME_TYPES, dtype=object)
    year_values = np.array(years, dtype=np.int32)
    long = pl.DataFrame(
        {
            "postcode": postcodes[pc_i].astype(str),
            "Crime type": type_names[ty_i].astype(str),
            "year": year_values[yr_i],
            "count": annual[pc_i, ty_i, yr_i].astype(np.float32),
        },
        schema_overrides={"postcode": pl.String, "Crime type": pl.String},
    )

    serious = _rollup_long(long, SERIOUS_CRIME_TYPES, "Serious crime")
    minor = _rollup_long(long, MINOR_CRIME_TYPES, "Minor crime")
    combined = pl.concat([long, serious, minor])

    by_type = (
        combined.sort("year")
        .group_by("postcode", "Crime type")
        .agg(pl.struct("year", "count").alias("series"))
    )
    wide = by_type.pivot(on="Crime type", index="postcode", values="series")
    type_cols = [c for c in wide.columns if c != "postcode"]
    wide = wide.rename({col: f"{col} (by year)" for col in type_cols})

    coverage_per_force: list[list[dict[str, int]]] = []
    for f in range(months_in_year_force.shape[0]):
        coverage_per_force.append(
            [
                {"year": int(years[y]), "months": int(m)}
                for y, m in enumerate(months_in_year_force[f])
                if m >= min_bar_months
            ]
        )
    coverage_frame = pl.DataFrame(
        {
            "_fidx": pl.Series(range(len(coverage_per_force)), dtype=pl.Int32),
            COVERAGE_COLUMN: pl.Series(
                coverage_per_force,
                dtype=pl.List(pl.Struct({"year": pl.Int32, "months": pl.Int32})),
            ),
        }
    )
    base = pl.DataFrame(
        {
            "postcode": postcodes,
            "_fidx": pl.Series(home_fidx, dtype=pl.Int32),
            "_usable": pl.Series(usable),
        }
    )
    dense = (
        base.join(coverage_frame, on="_fidx", how="left")
        .with_columns(
            pl.when(pl.col("_usable"))
            .then(pl.col(COVERAGE_COLUMN))
            .otherwise(pl.col(COVERAGE_COLUMN).list.head(0))
            .alias(COVERAGE_COLUMN)
        )
        .drop("_fidx", "_usable")
    )
    wide = dense.join(wide, on="postcode", how="left")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    wide.write_parquet(output_path, compression="zstd")
    print(f"Wrote postcode crime by-year series: {output_path}  {wide.shape}")


def _finalize_records(
    records_shard_dir: Path, output_path: Path
) -> None:
    """Concatenate the per-batch record shards into one postcode-sorted parquet.

    Sorting by postcode lets the server build a contiguous per-postcode slice
    index. The sort runs on the streaming engine so it spills rather than holding
    all ~40M rows in memory.
    """
    shards = sorted(records_shard_dir.glob("*.parquet"))
    output_path.parent.mkdir(parents=True, exist_ok=True)
    if not shards:
        pl.DataFrame(
            schema={
                "postcode": pl.String,
                "month_index": pl.Int32,
                "crime_type": pl.String,
                "location": pl.String,
                "outcome": pl.String,
                "lat": pl.Float32,
                "lon": pl.Float32,
            }
        ).write_parquet(output_path, compression="zstd")
        print(f"Wrote crime records: {output_path}  (empty)")
        return

    (
        pl.scan_parquet(shards)
        .sort("postcode")
        .sink_parquet(output_path, compression="zstd")
    )
    n = pl.scan_parquet(output_path).select(pl.len()).collect().item()
    print(f"Wrote crime records: {output_path}  ({n:,} rows)")


def transform_crime_spatial(
    crime_dir: Path,
    boundaries_dir: Path,
    output_path: Path,
    by_year_output_path: Path,
    records_output_path: Path,
    buffer_m: float = DEFAULT_BUFFER_M,
    max_postcodes: int | None = None,
    max_files: int | None = None,
    min_bar_months: int = MIN_BAR_MONTHS,
) -> None:
    csvs, ignored_csv_count = find_street_crime_csvs(crime_dir)
    if not csvs:
        raise FileNotFoundError(f"No street crime CSV files found in {crime_dir}")
    if max_files is not None:
        csvs = csvs[:max_files]

    years, forces, months_in_year_force = _force_calendar(csvs)
    latest_year = years[-1]
    # Records cover the longest window's calendar years (January of its earliest
    # year onward), so they reconcile exactly with the calendar-year headline.
    records_min_ym = (latest_year - (RECORDS_WINDOW_YEARS - 1)) * 12
    print(
        f"Found {len(csvs):,} street crime CSVs across {len(forces)} forces "
        f"({years[0]}-{years[-1]})"
        + (f" (ignored {ignored_csv_count} non-street CSVs)" if ignored_csv_count else "")
    )

    postcodes, polygons = load_postcode_polygons(boundaries_dir, max_postcodes)
    postcodes = np.asarray(postcodes)

    print(f"Buffering {len(postcodes):,} postcode polygons by {buffer_m:g}m...")
    tree, usable = _build_tree(polygons, buffer_m)

    type_to_idx = {name: idx for idx, name in enumerate(ALL_CRIME_TYPES)}
    year_to_idx = {year: idx for idx, year in enumerate(years)}
    force_to_idx = {force: idx for idx, force in enumerate(forces)}
    counts = np.zeros((len(postcodes), len(ALL_CRIME_TYPES), len(years)), dtype=np.int32)
    force_votes = np.zeros((len(postcodes), len(forces)), dtype=np.int32)

    transformer = Transformer.from_crs("EPSG:4326", "EPSG:27700", always_xy=True)

    with tempfile.TemporaryDirectory(
        prefix="crime_records_", dir=records_output_path.parent
    ) as shard_dir_str:
        shard_dir = Path(shard_dir_str)
        _accumulate_counts(
            csvs,
            postcodes,
            tree,
            type_to_idx,
            year_to_idx,
            force_to_idx,
            transformer,
            counts,
            force_votes,
            shard_dir,
            records_min_ym,
        )

        home_fidx = _assign_home_force(postcodes, force_votes, forces)

        # Per-window raw annualised averages (14 leaf types + Serious/Minor).
        raw_by_window: dict[str, np.ndarray] = {}
        for label, win_years in WINDOWS:
            year_mask = np.array(
                [1.0 if y > latest_year - win_years else 0.0 for y in years]
            )
            avg14 = _window_annualised(
                counts, months_in_year_force, home_fidx, usable, year_mask
            )
            raw_by_window[label] = _append_rollups(avg14)

        _write_crime_table(postcodes, raw_by_window, output_path)
        _write_by_year(
            postcodes,
            counts,
            years,
            months_in_year_force,
            home_fidx,
            usable,
            min_bar_months,
            by_year_output_path,
        )
        _finalize_records(shard_dir, records_output_path)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Count police.uk crime points near each postcode boundary"
    )
    parser.add_argument(
        "--input",
        type=Path,
        default=Path("property-data/crime"),
        help="Directory containing police.uk street crime CSVs",
    )
    parser.add_argument(
        "--boundaries",
        type=Path,
        default=Path("property-data/postcode_boundaries/units"),
        help="Directory of per-district postcode boundary GeoJSONs",
    )
    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Output parquet: postcode + '{type} (/yr, <window>)' average-annual-count columns",
    )
    parser.add_argument(
        "--output-by-year",
        type=Path,
        required=True,
        help="Output parquet: postcode + nested '{type} (by year)' columns",
    )
    parser.add_argument(
        "--output-records",
        type=Path,
        required=True,
        help="Output parquet: one row per counted incident (last 7 years), postcode-sorted",
    )
    parser.add_argument(
        "--max-postcodes",
        type=int,
        default=None,
        help="Testing only: process the first N postcodes",
    )
    parser.add_argument(
        "--max-files",
        type=int,
        default=None,
        help="Testing only: process the first N monthly CSV files",
    )
    parser.add_argument(
        "--min-bar-months",
        type=int,
        default=MIN_BAR_MONTHS,
        help="Minimum covered months for a year to get a by-year bar",
    )
    args = parser.parse_args()

    transform_crime_spatial(
        crime_dir=args.input,
        boundaries_dir=args.boundaries,
        output_path=args.output,
        by_year_output_path=args.output_by_year,
        records_output_path=args.output_records,
        max_postcodes=args.max_postcodes,
        max_files=args.max_files,
        min_bar_months=args.min_bar_months,
    )


if __name__ == "__main__":
    main()