"""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()