import polars as pl

from pipeline.transform.area_crime_averages import (
    NATIONAL_AREA,
    SCOPE_NATIONAL,
    SCOPE_OUTCODE,
    SCOPE_SECTOR,
    compute_area_crime_averages,
)

_BURGLARY = "Burglary (/yr, 7y)"
_ROBBERY = "Robbery (/yr, 7y)"


def _postcodes() -> pl.LazyFrame:
    return pl.LazyFrame(
        {
            "Postcode": ["E14 2DG", "E14 2AB", "E14 9XY", "M1 1AE", "E14 2ZZ"],
            # E14 9XY has no usable crime data; E14 2AB lacks robbery; E14 2ZZ has
            # crime but (below) no properties, so it must not weight any average.
            _BURGLARY: [10.0, 20.0, None, 5.0, 100.0],
            _ROBBERY: [2.0, None, None, 1.0, 50.0],
            # An unrelated column proves only the crime columns are averaged.
            "Median age": [40.0, 41.0, 42.0, 30.0, 99.0],
        }
    )


def _properties() -> pl.LazyFrame:
    # Property rows per postcode become the weights (3 / 1 / 2 / 4). E14 2ZZ has
    # none, so it is excluded entirely.
    postcodes = ["E14 2DG"] * 3 + ["E14 2AB"] + ["E14 9XY"] * 2 + ["M1 1AE"] * 4
    return pl.LazyFrame({"Postcode": postcodes})


def _row(df: pl.DataFrame, scope: str, area: str) -> dict:
    matched = df.filter((pl.col("scope") == scope) & (pl.col("area") == area))
    assert matched.height == 1, f"expected one {scope} row for {area!r}"
    return matched.to_dicts()[0]


def test_property_weighted_means_skip_nulls():
    result = compute_area_crime_averages(_postcodes(), _properties())

    national = _row(result, SCOPE_NATIONAL, NATIONAL_AREA)
    # Burglary: (10*3 + 20*1 + 5*4) / (3+1+4) = 70/8; E14 9XY null dilutes nothing.
    assert national[_BURGLARY] == 8.75
    # Robbery: (2*3 + 1*4) / (3+4) = 10/7; both null postcodes are excluded from
    # the numerator AND the denominator.
    assert national[_ROBBERY] == pl.Series([10.0 / 7.0]).cast(pl.Float32).item()

    outcode = _row(result, SCOPE_OUTCODE, "E14")
    assert outcode[_BURGLARY] == 12.5  # (10*3 + 20*1) / 4
    assert outcode[_ROBBERY] == 2.0  # only E14 2DG has robbery (2 * 3 / 3)


def test_sector_aggregation_and_all_null_rows_dropped():
    result = compute_area_crime_averages(_postcodes(), _properties())

    sector = _row(result, SCOPE_SECTOR, "E14 2")
    assert sector[_BURGLARY] == 12.5
    assert sector[_ROBBERY] == 2.0

    # E14 9XY has properties but no crime data at all, so its sector "E14 9" is
    # all-null and must be dropped rather than reported as a known area.
    assert result.filter(pl.col("area") == "E14 9").height == 0


def test_postcodes_without_properties_are_excluded():
    result = compute_area_crime_averages(_postcodes(), _properties())

    # E14 2ZZ carries crime values but no properties; including it would pull the
    # E14 outcode burglary mean toward its 100.0. It must contribute nothing.
    outcode = _row(result, SCOPE_OUTCODE, "E14")
    assert outcode[_BURGLARY] == 12.5


def test_only_crime_columns_are_emitted():
    result = compute_area_crime_averages(_postcodes(), _properties())
    assert set(result.columns) == {"scope", "area", _BURGLARY, _ROBBERY}
    assert result.schema[_BURGLARY] == pl.Float32