"""Shared property base: the dwelling universe before any area enrichment.

This is the single source of truth for *which* dwellings exist and their
intrinsic, source-level attributes (price, floor area, type, addresses). Both
``merge`` (which enriches it with postcode/LSOA-keyed area features to build
properties.parquet) and price estimation (which only needs the intrinsic
columns) start from exactly these rows, so estimates computed here line up 1:1
with the final properties by the natural key ``(Postcode, coalesced address)``.

Living in its own module is what lets price estimation be *cached* across
merge changes: ``price_inputs.parquet`` depends only on epc_pp + arcgis + this
file, so adding an area column to merge.py does not invalidate it and the
expensive index/kNN steps are skipped.
"""

import argparse
from pathlib import Path

import polars as pl

from pipeline.utils.postcode_mapping import build_postcode_mapping

MIN_FLOOR_AREA_M2 = 10

# Columns price estimation reads, with the final (properties.parquet) names so
# index.py/estimate.py and the join all speak the same schema. The two address
# columns form the natural join key (Postcode + their coalesce).
PRICE_INPUT_SELECT = [
    pl.col("postcode").alias("Postcode"),
    pl.col("total_floor_area").alias("Total floor area (sqm)"),
    pl.col("latest_price").alias("Last known price"),
    pl.col("date_of_transfer").alias("Date of last transaction"),
    "historical_prices",
    pl.col("pp_address").alias("Address per Property Register"),
    pl.col("epc_address").alias("Address per EPC"),
]


def _active_english_postcode_area(arcgis_raw: pl.LazyFrame) -> pl.LazyFrame:
    """Return the supported postcode universe with geography join keys."""
    return (
        arcgis_raw.filter(pl.col("ctry25cd") == "E92000001")
        .filter(pl.col("doterm").is_null())
        .select(
            pl.col("pcds").alias("postcode"),
            "lat",
            pl.col("long").alias("lon"),
            "ctry25cd",
            pl.col("lsoa21cd").alias("lsoa21"),
            pl.col("oa21cd").alias("oa21"),
            pl.col("pcon24cd").alias("pcon"),
        )
        .drop_nulls(["postcode"])
        .unique(["postcode"])
    )


def _remap_terminated_postcodes(
    wide: pl.LazyFrame, postcode_mapping: pl.LazyFrame
) -> pl.LazyFrame:
    return (
        wide.join(
            postcode_mapping,
            left_on="postcode",
            right_on="old_postcode",
            how="left",
        )
        .with_columns(
            pl.coalesce("new_postcode", "postcode").alias("postcode"),
        )
        .drop("new_postcode")
    )


def _dedupe_collapsed_properties(wide: pl.LazyFrame) -> pl.LazyFrame:
    """Keep one row per (postcode, address): the most-recent transaction.

    The terminated-postcode remap can map two distinct postcodes onto one active
    successor, collapsing the same physical address onto a single
    (postcode, address) key with conflicting sale records. Keep the row with the
    latest date_of_transfer so the headline price/date reflect the most recent
    transaction; genuinely distinct addresses are untouched.

    The dedup key coalesces the price-paid address with the EPC address: EPC-only
    dwellings (never sold) have a null pp_address, so keying on pp_address alone
    would collapse EVERY EPC-only dwelling in a postcode onto one
    (postcode, null) key and silently drop all but one. Each dwelling's coalesced
    address is unique within its postcode (the EPC frame is deduped on
    address+postcode upstream), so the coalesced key keeps them distinct while
    leaving sold-property dedup unchanged, since pp_address wins the coalesce whenever
    a sale exists.
    """
    return (
        wide.with_columns(
            pl.coalesce("pp_address", "epc_address").alias("_dedupe_address")
        )
        .sort("date_of_transfer", descending=True, nulls_last=True)
        .unique(
            subset=["postcode", "_dedupe_address"], keep="first", maintain_order=True
        )
        .drop("_dedupe_address")
    )


def _filter_to_active_english_postcodes(
    wide: pl.LazyFrame, active_postcodes: pl.LazyFrame
) -> pl.LazyFrame:
    return wide.join(active_postcodes, on="postcode", how="semi")


def property_type_expr() -> pl.Expr:
    """Unaliased property-type expression: prefer EPC, fall back to price-paid.

    For Houses, use built_form (e.g. Semi-Detached, Mid-Terrace) for finer
    detail. Depends only on intrinsic base columns (epc_property_type,
    pp_property_type, built_form), so merge and price_inputs derive the same
    value. Callers alias it ("property_type" in merge, "Property type" in
    price_inputs).
    """
    bad_built_form = pl.col("built_form").is_null() | pl.col("built_form").is_in(
        ["NO DATA!", "Not Recorded"]
    )
    has_epc = pl.col("epc_property_type").is_not_null()
    is_house = pl.col("epc_property_type") == "House"
    return (
        pl.when(has_epc & is_house & ~bad_built_form)
        .then(pl.col("built_form"))
        .when(has_epc & is_house)
        .then(pl.col("pp_property_type"))
        .when(has_epc)
        .then(pl.col("epc_property_type"))
        .otherwise(pl.col("pp_property_type"))
        # Unify EPC's "Flat"/"Maisonette" with price-paid's "Flats/Maisonettes",
        # collapse terrace sub-types, and fold rare types into "Other"
        .replace(
            {
                "Flat": "Flats/Maisonettes",
                "Maisonette": "Flats/Maisonettes",
                "End-Terrace": "Terraced",
                "Mid-Terrace": "Terraced",
                "Enclosed End-Terrace": "Terraced",
                "Enclosed Mid-Terrace": "Terraced",
                "Bungalow": "Other",
                "Park home": "Other",
            }
        )
    )


def build_postcode_centroids(arcgis_path: Path) -> pl.LazyFrame:
    """One row per active-English postcode with its lat/lon, from arcgis.

    This is the lat/lon source price estimation needs (index sector centroids,
    kNN). It is the same per-postcode lat/lon merge writes into postcode.parquet
    (both come from arcgis), but built straight from arcgis so the index/estimate
    steps do not depend on the merge output. Adding an area column to merge
    therefore does not invalidate the expensive price index/kNN.
    """
    return _active_english_postcode_area(pl.scan_parquet(arcgis_path)).select(
        pl.col("postcode").alias("Postcode"), "lat", "lon"
    )


def build_property_base(epc_pp_path: Path, arcgis_path: Path) -> pl.LazyFrame:
    """The deduped, active-English dwelling universe from epc_pp + arcgis.

    Floor filter -> terminated-postcode remap -> collapse-dedupe -> restrict to
    the active English postcode universe. Returns a LazyFrame with the original
    epc_pp column names; merge enriches it, the CLI projects it to price_inputs.
    """
    wide = pl.scan_parquet(epc_pp_path).filter(
        pl.col("total_floor_area").is_null()
        | (pl.col("total_floor_area") > MIN_FLOOR_AREA_M2)
    )
    postcode_mapping = build_postcode_mapping(arcgis_path)
    wide = _remap_terminated_postcodes(wide, postcode_mapping.lazy())
    wide = _dedupe_collapsed_properties(wide)
    active_postcodes = (
        _active_english_postcode_area(pl.scan_parquet(arcgis_path))
        .select("postcode")
        .unique()
    )
    return _filter_to_active_english_postcodes(wide, active_postcodes)


def main():
    parser = argparse.ArgumentParser(
        description="Write the slim price-estimation inputs from epc_pp + arcgis"
    )
    parser.add_argument("--epc-pp", type=Path, required=True)
    parser.add_argument("--arcgis", type=Path, required=True)
    parser.add_argument(
        "--output", type=Path, required=True, help="price_inputs.parquet output"
    )
    parser.add_argument(
        "--centroids",
        type=Path,
        required=True,
        help="postcode_centroids.parquet output (Postcode, lat, lon)",
    )
    args = parser.parse_args()

    base = build_property_base(args.epc_pp, args.arcgis)
    price_inputs = base.with_columns(
        property_type_expr().alias("Property type")
    ).select(*PRICE_INPUT_SELECT, "Property type")
    price_inputs.sink_parquet(args.output)
    n = pl.scan_parquet(args.output).select(pl.len()).collect().item()
    print(f"Wrote {args.output} ({args.output.stat().st_size / 1e6:.1f} MB), {n:,} dwellings")

    build_postcode_centroids(args.arcgis).sink_parquet(args.centroids)
    n_pc = pl.scan_parquet(args.centroids).select(pl.len()).collect().item()
    print(f"Wrote {args.centroids} ({args.centroids.stat().st_size / 1e6:.1f} MB), {n_pc:,} postcodes")


if __name__ == "__main__":
    main()