"""Precompute per-outcode / per-postcode-sector / national mean headline crime
counts for the right pane's area comparison.

The right pane shows each crime metric next to its area context: the mean
average-annual count (``"X (/yr, 7y)"``) across the selection's postcode sector (e.g.
``"E14 2"``), its outcode (e.g. ``"E14"``), and the nation. Crime is constant
within a postcode (the merge keys it on the postcode), so each postcode
contributes its single value weighted by how many properties sit in it — keeping
every scope on the same property-weighted basis as the per-selection mean, so the
four numbers (this selection / sector / outcode / nation) are directly
comparable. The national figure here is an EXACT property-weighted mean, which is
why it overrides the upward-biased histogram-bin national average for crime.

This used to be recomputed inside the server on every boot from the loaded
property matrix. It is a pure function of the two merge outputs, so it belongs in
the data build; the server now just loads the parquet this writes. Reading the
crime values from ``postcode.parquet`` and the per-postcode property weights from
``properties.parquet`` mirrors exactly the two inputs the server loads, so the
result matches what the server used to compute (minus its u16 quantization loss).

Output schema — one row per area:

    scope            : ``"national"`` | ``"outcode"`` | ``"sector"``
    area             : the outcode (``"E14"``) / sector (``"E14 2"``);
                       ``""`` for the single national row
    ``<type> (/yr, 7y|2y)`` : Float32 property-weighted mean crime count per year
                          (null = the scope has no data for that type)
"""

import argparse
from pathlib import Path

import polars as pl

# Filterable crime columns are the average-annual incident counts and carry this
# marker in postcode.parquet (e.g. "Burglary (/yr, 7y)"). We average those. The
# full column name is kept; the server discovers and keys area averages by the
# same names.
COUNT_MARKER = " (/yr, "

# `scope` discriminator values. The server's loader dispatches on these.
SCOPE_NATIONAL = "national"
SCOPE_OUTCODE = "outcode"
SCOPE_SECTOR = "sector"

# Area label on the national row — it spans the whole country, so it has no code.
NATIONAL_AREA = ""

# Both merge outputs key on the canonical NSPL `pcds` postcode (spaced, e.g.
# "E14 2DG").
POSTCODE_COLUMN = "Postcode"

# Internal weight / split columns dropped before write.
_WEIGHT_COLUMN = "_weight"
_OUTCODE_COLUMN = "_outcode"
_SECTOR_COLUMN = "_sector"


def _crime_columns(columns: list[str]) -> list[str]:
    crime_cols = [name for name in columns if COUNT_MARKER in name]
    if not crime_cols:
        raise ValueError(
            f"postcode parquet has no '*{COUNT_MARKER}*' crime count columns to average"
        )
    return crime_cols


def _weighted_mean(column: str) -> pl.Expr:
    """Property-weighted mean of ``column`` that excludes nulls from BOTH the
    value sum and the weight.

    A null crime value is a genuine gap (the postcode's police force published no
    usable data), not zero crime, so it must dilute neither the numerator nor the
    denominator — exactly as the server's former estimator skipped NaN values.
    Yields null when no postcode in the group has data for this type.
    """
    weight = pl.col(_WEIGHT_COLUMN)
    numerator = (pl.col(column) * weight).sum()
    denominator = weight.filter(pl.col(column).is_not_null()).sum()
    return (
        pl.when(denominator > 0)
        .then(numerator / denominator)
        .otherwise(None)
        .cast(pl.Float32)
        .alias(column)
    )


def compute_area_crime_averages(
    postcodes: pl.LazyFrame, properties: pl.LazyFrame
) -> pl.DataFrame:
    """Build the national / per-outcode / per-sector crime-average table.

    ``postcodes`` is the merge's postcode output (one row per active postcode,
    carrying the ``"* (/yr, *)"`` crime count columns); ``properties`` is the merge's
    per-property output, used only to weight each postcode by its property count.
    """
    crime_cols = _crime_columns(postcodes.collect_schema().names())

    # Property weight per postcode = how many property rows the server indexes
    # under it. The inner join keeps only postcodes that actually carry
    # properties, matching the server's per-postcode row index (a postcode with
    # no properties never contributed to any average).
    weights = properties.group_by(POSTCODE_COLUMN).agg(pl.len().alias(_WEIGHT_COLUMN))

    # Outcode / sector of the spaced `pcds` postcode, matching the server's
    # postcode_outcode / postcode_sector (split on the single space; sector =
    # outcode + space + first inward character). Null where the form has no
    # inward code, so such rows drop out of the per-area groups.
    parts = pl.col(POSTCODE_COLUMN).str.splitn(" ", 2).struct
    outward = parts.field("field_0")
    inward = parts.field("field_1")
    base = (
        postcodes.select(POSTCODE_COLUMN, *crime_cols)
        .join(weights, on=POSTCODE_COLUMN, how="inner")
        .with_columns(
            pl.when(inward.is_not_null())
            .then(outward)
            .otherwise(None)
            .alias(_OUTCODE_COLUMN),
            pl.when(inward.str.len_chars() >= 1)
            .then(outward + pl.lit(" ") + inward.str.slice(0, 1))
            .otherwise(None)
            .alias(_SECTOR_COLUMN),
        )
    )

    mean_exprs = [_weighted_mean(column) for column in crime_cols]

    national = base.select(
        pl.lit(SCOPE_NATIONAL).alias("scope"),
        pl.lit(NATIONAL_AREA).alias("area"),
        *mean_exprs,
    )
    by_outcode = (
        base.drop_nulls(_OUTCODE_COLUMN)
        .group_by(_OUTCODE_COLUMN)
        .agg(mean_exprs)
        .select(
            pl.lit(SCOPE_OUTCODE).alias("scope"),
            pl.col(_OUTCODE_COLUMN).alias("area"),
            *crime_cols,
        )
    )
    by_sector = (
        base.drop_nulls(_SECTOR_COLUMN)
        .group_by(_SECTOR_COLUMN)
        .agg(mean_exprs)
        .select(
            pl.lit(SCOPE_SECTOR).alias("scope"),
            pl.col(_SECTOR_COLUMN).alias("area"),
            *crime_cols,
        )
    )

    result = pl.concat([national, by_outcode, by_sector], how="vertical").collect(
        engine="streaming"
    )

    # Drop per-area rows where every crime type is null: the server only created
    # a map entry once a scope had at least one finite value, so an all-null
    # outcode/sector reported no code at all. The national row is always kept (it
    # always has data, and is emitted even for areas absent from both maps).
    has_any = pl.any_horizontal(pl.col(column).is_not_null() for column in crime_cols)
    return result.filter((pl.col("scope") == SCOPE_NATIONAL) | has_any)


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Precompute national / per-outcode / per-sector mean headline crime "
            "counts from the merge outputs"
        )
    )
    parser.add_argument(
        "--postcodes",
        type=Path,
        required=True,
        help="postcode.parquet (area features, incl. the '* (/yr, *)' crime columns)",
    )
    parser.add_argument(
        "--properties",
        type=Path,
        required=True,
        help="properties.parquet (per-property rows; supplies postcode property weights)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Output area_crime_averages.parquet path",
    )
    args = parser.parse_args()

    result = compute_area_crime_averages(
        pl.scan_parquet(args.postcodes), pl.scan_parquet(args.properties)
    )
    outcodes = result.filter(pl.col("scope") == SCOPE_OUTCODE).height
    sectors = result.filter(pl.col("scope") == SCOPE_SECTOR).height
    print(
        f"Area crime averages: {result.height} rows "
        f"({outcodes} outcodes, {sectors} sectors, "
        f"{len(_crime_columns(result.columns))} crime types)"
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    result.write_parquet(args.output, compression="zstd")
    print(f"Saved to {args.output}")


if __name__ == "__main__":
    main()