"""Model historical school catchment areas and count them per postcode.

No national dataset of school catchment areas exists for England: catchments
are set per admission authority, only a handful of councils publish polygons,
and the pupil-residence data behind commercial "heatmap" catchments lives in
the restricted National Pupil Database. This module therefore COMPILES one
from open data, estimating each school's admission cutoff distance ("last
distance offered") — the radius within which an applicant would plausibly be
offered a place.

Model: English state admissions are run as deferred acceptance with distance
tie-breaks, which in a continuum economy is equivalent to finding
market-clearing cutoff distances (Azevedo & Leshno 2016). Per phase
(primary/secondary):

1. Demand — Census 2021 children per LSOA (TS007A age bands, prorated to the
   phase's cohort ages) split evenly across the LSOA's live postcodes.
2. Supply — every open, non-selective state-funded school (GIAS), with a fill
   target of max(capacity, headcount) prorated to the phase's cohorts
   (sixth-form and nursery years carry reduced weight, since their class
   sizes differ and they are not allocated by the same admissions round).
3. Preferences — children prefer nearby schools, trading distance against
   Ofsted grade: a school's effective distance is its real distance minus a
   grade bonus (Outstanding > Good > ungraded > below-Good). Because real
   first preferences are heterogeneous, each postcode's children split
   across nearby feasible schools with logit weights over effective
   distance rather than all picking the same one.
4. Equilibrium — cutoffs start unbounded and tighten monotonically: each
   round, children apply to their preferred feasible school(s), and
   oversubscribed schools tighten their cutoff to the distance of their
   marginal admitted child. Converges to the deferred-acceptance outcome.
5. Schools that never fill have no binding cutoff — anyone who applies gets
   in — so their feasibility radius is the distance within which the local
   child population would cover their fill target, capped.

The free parameters (preference bonuses, demand scale, choice temperature,
residual calibration factors) are CALIBRATED against published "last
distance offered" figures scraped from nine local authorities' allocation
reports — see check_school_cutoffs.py and the constants below.

A postcode is "inside the catchment" of every school whose cutoff radius
covers it. The output counts those schools per postcode for the four
good+/outstanding x primary/secondary categories (Ofsted-classified, same
rules as the previous proximity metric). Selective (grammar) schools are
excluded throughout: their intakes are test-based and region-wide, so a
distance model would fabricate a catchment that does not exist.

Known limitations: faith oversubscription criteria are not modelled (whether
a faith school's catchment is open to a given family depends on the family),
and Census 2021 child counts lag current rolls slightly. Cutoffs are
straight-line distances, the modal LA tie-break criterion.
"""

import argparse
from pathlib import Path

import numpy as np
import polars as pl
from scipy.spatial import cKDTree

from pipeline.utils.poi_counts import _project_lat_lng_km, valid_uk_coords_mask

SCHOOL_GROUPS = {
    "good_primary": ["good_primary", "outstanding_primary"],
    "good_secondary": ["good_secondary", "outstanding_secondary"],
    "outstanding_primary": ["outstanding_primary"],
    "outstanding_secondary": ["outstanding_secondary"],
}

# Age thresholds for deciding which phase(s) a school serves. A school serves
# PRIMARY-age children if its statutory lowest age is <= 10, and SECONDARY-age
# children if its statutory highest age is >= 12. All-through (e.g. 3-18) and
# middle-deemed-secondary (e.g. 9-13) schools satisfy BOTH and so are counted in
# both the primary and the secondary metrics — Ofsted's coarse "Ofsted phase"
# labels such schools as just "Secondary", which previously hid them from every
# postcode's primary-school count.
PRIMARY_MAX_AGE = 10
SECONDARY_MIN_AGE = 12

# Cohort ages (inclusive) each phase competes for: Reception-Y6 and Y7-Y11.
PRIMARY_AGES = (4, 10)
SECONDARY_AGES = (11, 15)

# Cohort weights for prorating a school's headcount/capacity across the ages
# it teaches. Nursery classes are typically part-time and small; sixth forms
# run at roughly 60% of a school's Y7-Y11 cohort size. A flat proration
# undersupplied secondary places by ~8%.
NURSERY_COHORT_WEIGHT = 0.5  # ages < 4
SIXTH_FORM_COHORT_WEIGHT = 0.6  # ages >= 16

# Assumed bounds for the one-sided age-range shapes GIAS emits when a
# statutory age is missing: "up to {high}" starts at the earliest nursery
# intake, "{low}+" runs to the end of sixth form.
EARLIEST_INTAKE_AGE = 2
DEFAULT_LEAVING_AGE = 19

# Only schools that admit (mostly) by geography take part in the assignment.
# Independent, special and Welsh schools and post-16 colleges either don't
# admit by distance or fall outside the England postcode universe; selective
# (grammar) schools admit by test from a wide region.
STATE_SCHOOL_TYPE_GROUPS = [
    "Academies",
    "Local authority maintained schools",
    "Free Schools",
]

# Preference bonuses (km of extra travel a family accepts for a better
# school), applied as a discount on effective distance when children choose.
# Grade 3/4 schools repel by the same magnitudes.
PREF_BONUS_OUTSTANDING_KM = 0.6
PREF_BONUS_GOOD_KM = 0.3

# Share of resident children who actually compete for state places. Census
# 2021 counts overstate current entry cohorts (birth rates fell ~10% between
# 2016 and 2021, which is exactly the gap between the census stock and the
# children reaching Reception by mid-decade) and independent/home-educated
# children (~7%) never enter the allocation at all. Without this, modelled
# cutoffs run systematically tight and undersubscribed schools look full.
DEMAND_SCALE = 0.8

# Logit choice temperature (km). With deterministic choice every child at a
# postcode ranks the same school first, so popular schools fill entirely from
# their nearest band and the marginal admitted child sits unrealistically
# close. Real first preferences are heterogeneous; a school draws only a
# distance-decaying share of nearby families. Children therefore split across
# nearby feasible schools with weights softmax(-effective_distance / tau):
# higher tau = more smearing = wider cutoffs. tau -> 0 recovers the
# deterministic model (used by the unit tests). Calibrated 2026-06 against
# 240 published binding cutoffs from 9 LAs (check_school_cutoffs.py): 0.3 km
# maximises rank correlation and within-2x share; beyond ~0.6 the smearing
# erases school-to-school differentiation (Spearman 0.24 -> 0.01).
CHOICE_TEMPERATURE_KM = 0.3

# Residual calibration from the same ground truth: after the equilibrium
# solve, modelled cutoffs still ran systematically tight (median log2 bias
# -0.53 primary / -0.36 secondary at the settings above — published "last
# distance offered" reflects offer-day frictions, waiting-list churn and
# furthest-applicant noise that no clean equilibrium reproduces). Radii are
# multiplied by 2^-bias so the modelled median matches the published median;
# rank ordering is unaffected.
CUTOFF_CALIBRATION_FACTOR = {"primary": 1.44, "secondary": 1.28}

# Each demand postcode considers this many nearest schools; beyond ~16
# candidates assignment shares are negligible.
NEAREST_SCHOOL_CANDIDATES = 16

# Radius guard rails: the floor absorbs postcode-centroid noise around tiny
# urban catchments; the cap bounds feasibility radii for schools the model
# never fills (mostly rural).
MIN_RADIUS_KM = 0.3
MAX_RADIUS_KM = 25.0

EQUILIBRIUM_MAX_ITER = 100


def classify_good_plus_schools(
    ofsted: pl.DataFrame, open_urns: set[int] | None = None
) -> pl.DataFrame:
    """Label good+/outstanding primary & secondary schools for catchment counts.

    Derives a grade ("1" = outstanding, "2" = good) and one or two
    ``category`` rows per school, returning a ``(urn, category)`` frame.

    Schools with a recent GRADED inspection carry a 1-4 grade in "Latest OEIF
    overall effectiveness" (OEIF = the previous Ofsted Education Inspection
    Framework). A large and growing share of schools were last inspected under an
    UNGRADED (Section 8) inspection or the post-2024 report-card framework, so
    that column is null/"Not judged" for them even when they are demonstrably
    good — their status lives in "Ungraded inspection overall outcome" ("School
    remains Good"/"School remains Outstanding"). Filtering on the graded column
    alone dropped ~7,000 genuinely good/outstanding schools. We fall back to the
    ungraded outcome, but ONLY when there is no usable graded result
    (null/"Not judged"), so a genuine grade 3/4 is never overridden.

    Outcomes flagged "(Concerns)" are NOT treated as good+: a "remains Good
    (Concerns)" outcome signals inspectors found issues warranting an earlier
    graded re-inspection, so marketing it as a good+ school is misleading.

    Phase assignment uses the statutory age range when available (so all-through
    and middle schools count toward BOTH primary and secondary), falling back to
    the coarse "Ofsted phase" label when age columns are absent. When
    ``open_urns`` is given, schools whose URN is not in the current GIAS open
    register are dropped so closed/merged schools are not counted.
    """
    graded = _with_derived_grade(ofsted).filter(
        pl.col("Ofsted phase").is_in(["Primary", "Secondary"])
        & pl.col("_ofsted_grade").is_in(["1", "2"])
    )

    # Drop schools no longer open (closed/merged) when the GIAS open register is
    # provided, so stale Ofsted "latest inspection" rows are not counted.
    if open_urns is not None and "URN" in graded.columns:
        graded = graded.filter(pl.col("URN").is_in(list(open_urns)))

    # Decide which phase(s) each school serves.
    if {"Statutory lowest age", "Statutory highest age"} <= set(graded.columns):
        low = pl.col("Statutory lowest age").cast(pl.Int64, strict=False)
        high = pl.col("Statutory highest age").cast(pl.Int64, strict=False)
        serves_primary = (
            pl.when(low.is_not_null())
            .then(low <= PRIMARY_MAX_AGE)
            .otherwise(pl.col("Ofsted phase") == "Primary")
        )
        serves_secondary = (
            pl.when(high.is_not_null())
            .then(high >= SECONDARY_MIN_AGE)
            .otherwise(pl.col("Ofsted phase") == "Secondary")
        )
    else:
        serves_primary = pl.col("Ofsted phase") == "Primary"
        serves_secondary = pl.col("Ofsted phase") == "Secondary"

    graded = graded.with_columns(
        serves_primary.alias("_serves_primary"),
        serves_secondary.alias("_serves_secondary"),
    )

    # Good+ groups include both grade variants; outstanding groups count grade 1.
    # A school can yield up to two rows (primary and secondary).
    primary = graded.filter(pl.col("_serves_primary")).with_columns(
        pl.when(pl.col("_ofsted_grade") == "1")
        .then(pl.lit("outstanding_primary"))
        .otherwise(pl.lit("good_primary"))
        .alias("category")
    )
    secondary = graded.filter(pl.col("_serves_secondary")).with_columns(
        pl.when(pl.col("_ofsted_grade") == "1")
        .then(pl.lit("outstanding_secondary"))
        .otherwise(pl.lit("good_secondary"))
        .alias("category")
    )
    return pl.concat([primary, secondary]).select(
        pl.col("URN").cast(pl.Int64).alias("urn"),
        "category",
    )


def _with_derived_grade(ofsted: pl.DataFrame) -> pl.DataFrame:
    """Attach ``_ofsted_grade`` ("1"-"4" or null): graded OEIF result first,
    falling back to ungraded "School remains Good/Outstanding" outcomes (minus
    "(Concerns)") only when there is no usable graded result."""
    # Cast to Utf8 so the string predicates below are well-defined even if a
    # column happens to be entirely null (read back as a Null dtype).
    oeif = pl.col("Latest OEIF overall effectiveness").cast(pl.Utf8, strict=False)
    ungraded = pl.col("Ungraded inspection overall outcome").cast(pl.Utf8, strict=False)
    no_usable_grade = oeif.is_null() | (oeif == "Not judged")
    has_concern = ungraded.str.contains(r"\(Concerns\)")
    remains_outstanding = (
        ungraded.str.starts_with("School remains Outstanding") & ~has_concern
    )
    remains_good = ungraded.str.starts_with("School remains Good") & ~has_concern
    return ofsted.with_columns(
        pl.when(oeif.is_in(["1", "2", "3", "4"]))
        .then(oeif)
        .when(no_usable_grade & remains_outstanding)
        .then(pl.lit("1"))
        .when(no_usable_grade & remains_good)
        .then(pl.lit("2"))
        .otherwise(None)
        .alias("_ofsted_grade")
    )


def school_preference_bonuses(
    ofsted: pl.DataFrame,
    bonus_outstanding_km: float = PREF_BONUS_OUTSTANDING_KM,
    bonus_good_km: float = PREF_BONUS_GOOD_KM,
) -> pl.DataFrame:
    """Per-school preference bonus in km, from the derived Ofsted grade.

    Outstanding/Good schools attract demand from further away; grade 3/4
    schools repel it symmetrically. Ungraded (typically new) schools are
    neutral. Returns ``(urn, bonus_km)`` with one row per URN.
    """
    bonus = {
        "1": bonus_outstanding_km,
        "2": bonus_good_km,
        "3": -bonus_good_km,
        "4": -bonus_outstanding_km,
    }
    return (
        _with_derived_grade(ofsted)
        .filter(pl.col("URN").is_not_null())
        .select(
            pl.col("URN").cast(pl.Int64).alias("urn"),
            pl.col("_ofsted_grade")
            .replace_strict(bonus, default=0.0, return_dtype=pl.Float64)
            .alias("bonus_km"),
        )
        .unique(subset="urn", keep="first")
    )


def phase_intakes(gias: pl.DataFrame) -> pl.DataFrame:
    """Per-school phase-prorated fill targets for the admissions model.

    Returns one row per open, non-selective state-funded school with valid
    coordinates: ``(urn, lat, lng, primary_intake, secondary_intake)``. The
    fill target — max(capacity, headcount), so over-full schools keep their
    demonstrated size and under-full schools can admit up to capacity — is
    spread over the cohort ages the school teaches (parsed from ``age_range``,
    e.g. "3–11" = ages 3..10) with nursery and sixth-form ages down-weighted,
    and each phase receives the share of cohort weight in its age band.
    """
    # gias._format_age_range emits three shapes: "{low}–{high}", "up to {high}"
    # (StatutoryLowAge missing) and "{low}+" (StatutoryHighAge missing). Parse
    # all three — the one-sided shapes previously fell through the two-number
    # parse and silently dropped the school from the catchment supply.
    age = pl.col("age_range")
    leading = age.str.extract(r"^\s*(\d+)", 1).cast(pl.Int64, strict=False)
    trailing = age.str.extract(r"(\d+)\s*$", 1).cast(pl.Int64, strict=False)
    low = (
        pl.when(age.str.starts_with("up to"))
        .then(pl.lit(EARLIEST_INTAKE_AGE, dtype=pl.Int64))
        .otherwise(leading)
    )
    # The leaving age is exclusive as a cohort: a "3-11" school teaches
    # children aged 3 through 10. "{low}+" schools get the end of sixth form
    # as their assumed leaving age (post-19 institutions then carry no
    # primary/secondary cohort weight and drop out naturally).
    high = (
        pl.when(age.str.ends_with("+"))
        .then(pl.lit(DEFAULT_LEAVING_AGE, dtype=pl.Int64))
        .otherwise(trailing)
        - 1
    )

    schools = (
        gias.filter(
            pl.col("type_group").is_in(STATE_SCHOOL_TYPE_GROUPS)
            & (
                pl.col("admissions_policy").is_null()
                | (pl.col("admissions_policy") != "Selective")
            )
            & pl.col("lat").is_not_null()
            & pl.col("lng").is_not_null()
        )
        .with_columns(low.alias("_low"), high.alias("_high"))
        .filter(pl.col("_low").is_not_null() & (pl.col("_high") >= pl.col("_low")))
        .with_columns(
            pl.max_horizontal(
                pl.col("pupils").fill_null(0), pl.col("capacity").fill_null(0)
            )
            .cast(pl.Float64)
            .alias("_fill_target"),
        )
        .filter(pl.col("_fill_target") > 0)
    )

    def weighted_overlap(lo: int, hi: int, weight: float = 1.0) -> pl.Expr:
        """Cohort weight contributed by ages [lo, hi] within [_low, _high]."""
        return (
            weight
            * (
                pl.min_horizontal(pl.col("_high"), hi)
                - pl.max_horizontal(pl.col("_low"), lo)
                + 1
            ).clip(lower_bound=0)
        ).cast(pl.Float64)

    total_weight = (
        weighted_overlap(0, 3, NURSERY_COHORT_WEIGHT)
        + weighted_overlap(4, 15)
        + weighted_overlap(16, 30, SIXTH_FORM_COHORT_WEIGHT)
    )
    return schools.select(
        pl.col("urn").cast(pl.Int64),
        "lat",
        "lng",
        (pl.col("_fill_target") * weighted_overlap(*PRIMARY_AGES) / total_weight).alias(
            "primary_intake"
        ),
        (
            pl.col("_fill_target") * weighted_overlap(*SECONDARY_AGES) / total_weight
        ).alias("secondary_intake"),
    )


def children_per_postcode(
    postcodes: pl.DataFrame, lsoa_children: pl.DataFrame
) -> pl.DataFrame:
    """Estimate phase-age children living at each live postcode.

    Census age bands don't align with school phases, so phase totals take
    fractional shares of bands (one fifth per single year of age): primary
    (4-10) = age 4 + ages 5-9 + age 10, secondary (11-15) = ages 11-14 +
    age 15. LSOA totals are then split evenly across the LSOA's postcodes.
    """
    lsoa = lsoa_children.select(
        "lsoa21",
        (
            0.2 * pl.col("aged_0_4") + pl.col("aged_5_9") + 0.2 * pl.col("aged_10_14")
        ).alias("_lsoa_primary"),
        (0.8 * pl.col("aged_10_14") + 0.2 * pl.col("aged_15_19")).alias(
            "_lsoa_secondary"
        ),
    )
    return (
        postcodes.join(lsoa, left_on="lsoa21cd", right_on="lsoa21", how="inner")
        .with_columns(pl.len().over("lsoa21cd").alias("_lsoa_postcodes"))
        .select(
            "postcode",
            "lat",
            "lng",
            (pl.col("_lsoa_primary") / pl.col("_lsoa_postcodes")).alias(
                "primary_children"
            ),
            (pl.col("_lsoa_secondary") / pl.col("_lsoa_postcodes")).alias(
                "secondary_children"
            ),
        )
    )


def equilibrium_cutoffs(
    school_xy: np.ndarray,
    fill_target: np.ndarray,
    bonus_km: np.ndarray,
    pc_xy: np.ndarray,
    pc_children: np.ndarray,
    k: int = NEAREST_SCHOOL_CANDIDATES,
    max_iter: int = EQUILIBRIUM_MAX_ITER,
    tau_km: float = CHOICE_TEMPERATURE_KM,
) -> np.ndarray:
    """Market-clearing admission cutoff distance (km) per school.

    Deferred acceptance with distance priority, solved as cutoff dynamics
    (Azevedo & Leshno): cutoffs start unbounded; each round every child unit
    applies to its preferred feasible school(s) — a logit split over
    effective distance (distance - school bonus) among schools whose cutoff
    covers it, collapsing to the single best school when ``tau_km`` is 0 —
    and each oversubscribed school tightens its cutoff to its marginal
    admitted child's distance. Cutoffs only ever tighten, so the iteration
    converges.

    Returns np.inf for schools that never fill (no binding cutoff).
    """
    n_schools = len(school_xy)
    k = min(k, n_schools)
    demand = np.flatnonzero(pc_children > 0)
    weights = pc_children[demand]
    tree = cKDTree(school_xy)
    dist, cand = tree.query(pc_xy[demand], k=k, workers=-1)
    if k == 1:
        dist = dist[:, None]
        cand = cand[:, None]
    eff = dist - bonus_km[cand]

    rows = np.arange(len(demand))
    cutoff = np.full(n_schools, np.inf)
    for _ in range(max_iter):
        eff_feasible = np.where(dist <= cutoff[cand], eff, np.inf)
        if tau_km <= 0:
            choice = np.argmin(eff_feasible, axis=1)
            valid = np.isfinite(eff_feasible[rows, choice])
            chosen_school = cand[rows[valid], choice[valid]]
            chosen_dist = dist[rows[valid], choice[valid]]
            chosen_mass = weights[valid]
        else:
            z = -eff_feasible / tau_km
            z_max = z.max(axis=1, keepdims=True)
            share = np.exp(z - np.where(np.isfinite(z_max), z_max, 0.0))
            share[~np.isfinite(eff_feasible)] = 0.0
            total = share.sum(axis=1, keepdims=True)
            mass = weights[:, None] * share / np.where(total > 0, total, 1.0)
            # Sub-thousandth-of-a-child applications only slow the sort down.
            keep = mass > 1e-3
            chosen_school = cand[keep]
            chosen_dist = dist[keep]
            chosen_mass = mass[keep]

        order = np.lexsort((chosen_dist, chosen_school))
        s_sorted = chosen_school[order]
        d_sorted = chosen_dist[order]
        m_cum = np.cumsum(chosen_mass[order])
        boundaries = np.flatnonzero(np.diff(s_sorted)) + 1
        starts = np.concatenate(([0], boundaries))
        ends = np.concatenate((boundaries, [len(s_sorted)]))

        changed = False
        for start, end in zip(starts, ends):
            school = s_sorted[start]
            seg_cum = m_cum[start:end] - (m_cum[start - 1] if start else 0.0)
            if seg_cum[-1] <= fill_target[school]:
                continue
            marginal = d_sorted[start + np.searchsorted(seg_cum, fill_target[school])]
            if marginal < cutoff[school]:
                cutoff[school] = marginal
                changed = True
        if not changed:
            break

    return cutoff


def capacity_fill_radii(
    school_xy: np.ndarray,
    fill_target: np.ndarray,
    pc_xy: np.ndarray,
    pc_children: np.ndarray,
    max_radius_km: float = MAX_RADIUS_KM,
) -> np.ndarray:
    """Feasibility radius for schools without a binding cutoff.

    An undersubscribed school admits anyone who applies, so its catchment is
    bounded by plausibility rather than competition: the distance within
    which the local child population would cover its fill target. Capped at
    ``max_radius_km``.
    """
    demand = np.flatnonzero(pc_children > 0)
    tree = cKDTree(pc_xy[demand])
    radii = np.full(len(school_xy), max_radius_km)
    k = min(4096, len(demand))
    for i in range(len(school_xy)):
        dists, idx = tree.query(
            school_xy[i], k=k, distance_upper_bound=max_radius_km
        )
        found = np.isfinite(dists)
        cum = np.cumsum(pc_children[demand[idx[found]]])
        if len(cum) and cum[-1] >= fill_target[i]:
            radii[i] = dists[found][np.searchsorted(cum, fill_target[i])]
    return radii


def count_covering_catchments(
    pc_xy: np.ndarray,
    pc_valid: np.ndarray,
    school_xy: np.ndarray,
    school_radii: np.ndarray,
    n_postcodes: int,
) -> np.ndarray:
    """Count, per postcode, how many schools' catchment radii cover it."""
    counts = np.zeros(n_postcodes, dtype=np.int32)
    if len(school_xy) == 0:
        return counts
    valid_indices = np.flatnonzero(pc_valid)
    tree = cKDTree(pc_xy[valid_indices])
    covered = np.zeros(len(valid_indices), dtype=np.int32)
    for indices in tree.query_ball_point(school_xy, school_radii, workers=-1):
        covered[indices] += 1
    counts[valid_indices] = covered
    return counts


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Model school admission cutoff radii and count good+/outstanding "
            "primary/secondary catchments covering each postcode"
        )
    )
    parser.add_argument(
        "--ofsted", type=Path, required=True, help="Ofsted inspection parquet"
    )
    parser.add_argument(
        "--gias", type=Path, required=True, help="GIAS open-school parquet"
    )
    parser.add_argument(
        "--arcgis", type=Path, required=True, help="ArcGIS postcode parquet"
    )
    parser.add_argument(
        "--lsoa-children",
        type=Path,
        required=True,
        help="Census 2021 children by LSOA parquet",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Per-postcode counts parquet; omit for calibration runs that only "
        "need --schools-output",
    )
    parser.add_argument(
        "--schools-output",
        type=Path,
        default=None,
        help="Optional per-school catchment radii parquet (for calibration/debugging)",
    )
    parser.add_argument(
        "--bonus-outstanding-km",
        type=float,
        default=PREF_BONUS_OUTSTANDING_KM,
        help="Preference bonus for Outstanding schools (calibration sweeps)",
    )
    parser.add_argument(
        "--bonus-good-km",
        type=float,
        default=PREF_BONUS_GOOD_KM,
        help="Preference bonus for Good schools (calibration sweeps)",
    )
    parser.add_argument(
        "--demand-scale",
        type=float,
        default=DEMAND_SCALE,
        help="Share of resident children competing for state places",
    )
    parser.add_argument(
        "--choice-temperature-km",
        type=float,
        default=CHOICE_TEMPERATURE_KM,
        help="Logit choice temperature over effective distance",
    )
    args = parser.parse_args()

    gias = pl.read_parquet(args.gias)
    open_urns = set(
        gias.select(pl.col("urn").cast(pl.Int64, strict=False))
        .to_series()
        .drop_nulls()
        .to_list()
    )
    print(f"GIAS open register: {len(open_urns):,} open school URNs")

    ofsted = pl.read_parquet(args.ofsted)
    rated = classify_good_plus_schools(ofsted, open_urns=open_urns)
    if rated.is_empty():
        raise ValueError("No good+ primary/secondary Ofsted schools found")
    print(f"Good+ school/phase rows: {len(rated):,}")

    supply = phase_intakes(gias).join(
        school_preference_bonuses(
            ofsted,
            bonus_outstanding_km=args.bonus_outstanding_km,
            bonus_good_km=args.bonus_good_km,
        ),
        on="urn",
        how="left",
    ).with_columns(pl.col("bonus_km").fill_null(0.0))
    print(f"State schools in admissions model: {len(supply):,}")

    arcgis = pl.read_parquet(args.arcgis).select(
        pl.col("pcds").alias("postcode"),
        "lat",
        pl.col("long").alias("lng"),
        "lsoa21cd",
        "doterm",
    )
    live = arcgis.filter(
        pl.col("doterm").is_null() & pl.col("lsoa21cd").str.starts_with("E")
    )
    demand = children_per_postcode(live, pl.read_parquet(args.lsoa_children))
    print(
        f"Demand postcodes: {len(demand):,} "
        f"({demand['primary_children'].sum():,.0f} primary-age, "
        f"{demand['secondary_children'].sum():,.0f} secondary-age children)"
    )

    # Shared local-km projection so assignment and coverage use one metric.
    pc_lats = arcgis["lat"].to_numpy()
    pc_lngs = arcgis["lng"].to_numpy()
    pc_valid = valid_uk_coords_mask(pc_lats, pc_lngs)
    origin_lat = float(np.mean(pc_lats[pc_valid]))
    pc_xy = _project_lat_lng_km(pc_lats, pc_lngs, origin_lat)

    demand_lats = demand["lat"].to_numpy()
    demand_lngs = demand["lng"].to_numpy()
    demand_valid = valid_uk_coords_mask(demand_lats, demand_lngs)
    demand_xy = _project_lat_lng_km(demand_lats, demand_lngs, origin_lat)

    school_xy = _project_lat_lng_km(
        supply["lat"].to_numpy(), supply["lng"].to_numpy(), origin_lat
    )

    radii = {}
    for phase in ("primary", "secondary"):
        in_phase = supply[f"{phase}_intake"].to_numpy() > 0
        targets = supply[f"{phase}_intake"].to_numpy()[in_phase]
        xy = school_xy[in_phase]
        children = np.where(
            demand_valid,
            demand[f"{phase}_children"].to_numpy() * args.demand_scale,
            0.0,
        )
        print(f"Solving {phase} admissions for {in_phase.sum():,} schools...")
        cutoffs = equilibrium_cutoffs(
            xy,
            targets,
            supply["bonus_km"].to_numpy()[in_phase],
            demand_xy,
            children,
            tau_km=args.choice_temperature_km,
        )
        filled = np.isfinite(cutoffs)
        print(
            f"  {filled.sum():,} schools have binding cutoffs "
            f"(median {np.median(cutoffs[filled]):.2f} km); "
            f"{(~filled).sum():,} undersubscribed"
        )
        fallback = capacity_fill_radii(
            xy[~filled], targets[~filled], demand_xy, children
        )
        raw = cutoffs.copy()
        raw[~filled] = fallback
        radii[phase] = pl.DataFrame(
            {
                "urn": supply["urn"].to_numpy()[in_phase],
                "phase": phase,
                "cutoff_km": raw,
                "filled": filled,
                "radius_km": np.clip(
                    raw * CUTOFF_CALIBRATION_FACTOR[phase],
                    MIN_RADIUS_KM,
                    MAX_RADIUS_KM,
                ),
            }
        )
        print(
            f"  radius km: median {radii[phase]['radius_km'].median():.2f}, "
            f"p90 {radii[phase]['radius_km'].quantile(0.9):.2f}"
        )

    # Attach each rated school's phase radius; rated schools outside the
    # admissions model (special schools, selective schools, missing
    # headcounts) cannot be given a defensible radius and are dropped.
    rated = rated.with_columns(
        pl.col("category").str.split("_").list.get(1).alias("phase")
    )
    rated_with_radius = rated.join(
        pl.concat(list(radii.values())), on=["urn", "phase"], how="inner"
    ).join(supply.select("urn", "lat", "lng"), on="urn", how="inner")
    dropped = len(rated) - len(rated_with_radius)
    print(
        f"Rated school/phase rows with radii: {len(rated_with_radius):,} "
        f"(dropped {dropped:,}, incl. selective schools)"
    )

    if args.output is None and args.schools_output is None:
        raise SystemExit("Provide --output and/or --schools-output")

    if args.output is not None:
        category_counts = {}
        for category in set(c for cats in SCHOOL_GROUPS.values() for c in cats):
            cat = rated_with_radius.filter(pl.col("category") == category)
            cat_xy = _project_lat_lng_km(
                cat["lat"].to_numpy(), cat["lng"].to_numpy(), origin_lat
            )
            category_counts[category] = count_covering_catchments(
                pc_xy, pc_valid, cat_xy, cat["radius_km"].to_numpy(), len(arcgis)
            )
            print(f"  {category}: {len(cat):,} schools")

        result = pl.DataFrame(
            {
                "postcode": arcgis["postcode"],
                **{
                    f"{group}_catchments": sum(category_counts[c] for c in categories)
                    for group, categories in SCHOOL_GROUPS.items()
                },
            }
        )
        for group in SCHOOL_GROUPS:
            col = result[f"{group}_catchments"]
            print(f"  {group}_catchments: mean {col.mean():.2f}, max {col.max()}")

        args.output.parent.mkdir(parents=True, exist_ok=True)
        result.write_parquet(args.output)
        size_mb = args.output.stat().st_size / (1024 * 1024)
        print(f"Wrote {args.output} ({size_mb:.1f} MB)")

    if args.schools_output is not None:
        schools_out = rated_with_radius.select(
            "urn", "category", "phase", "cutoff_km", "filled", "radius_km", "lat", "lng"
        )
        args.schools_output.parent.mkdir(parents=True, exist_ok=True)
        schools_out.write_parquet(args.schools_output)
        print(f"Wrote {args.schools_output}")


if __name__ == "__main__":
    main()