"""Download Census 2021 household tenure (TS054) by LSOA.

Downloads the household-tenure breakdown (TS054, classification C2021_TENURE_9)
from the NOMIS API at LSOA 2021 granularity and folds the 8 detailed leaf
categories into our 3 output buckets (Owner occupied / Social rent /
Private rent), emitting one row per LSOA with the percentage of households in
each. The three buckets sum to 100%, so downstream they render as a
composition/ratio (like the ethnicity and qualifications stacked bars) AND each
percentage is independently filterable.

The 3-bucket grouping follows ONS's standard tenure summary:
  * Owner occupied = Owns outright + Owns with a mortgage/loan + Shared ownership
  * Social rent    = Rents from council/LA + Other social rented
  * Private rent    = Private landlord/letting agency + Other private + Lives rent free
(Shared ownership is part-owned and rolls into owner-occupied; "lives rent free"
rolls into private rent, mirroring ONS's "Private rented or lives rent free".)

NOTE this table counts HOUSEHOLDS (not usual residents). The join key downstream
(merge.py) is `lsoa21`, the same key used for ethnicity, qualifications,
median age, and IoD.

Source: NOMIS (ONS Census 2021, TS054 dataset, NM_2072_1)
License: Open Government Licence v3.0
"""

import argparse
from pathlib import Path

import polars as pl

from pipeline.utils import ENGLAND_LSOA_COUNT_2021, download_nomis_csv

pl.Config.set_tbl_cols(-1)

# NOMIS API: Census 2021 TS054 (household tenure) by LSOA 2021 (TYPE151).
# c2021_tenure_9=1..8 selects the 8 substantive leaf categories (excluding the
# aggregates 0/1001-1004/9996/9997, whose components we sum ourselves).
# measures=20100 selects the count.
BASE_URL = (
    "https://www.nomisweb.co.uk/api/v01/dataset/NM_2072_1.data.csv"
    "?date=latest"
    "&geography=TYPE151"
    "&c2021_tenure_9=1,2,3,4,5,6,7,8"
    "&measures=20100"
    "&select=GEOGRAPHY_CODE,C2021_TENURE_9_NAME,OBS_VALUE"
)

# Map each canonical NOMIS C2021_TENURE_9 leaf to our 3 output buckets. Keyed on
# the exact NOMIS label so a relabelled or missing leaf fails loudly in
# validation instead of silently dropping households from the denominator.
TENURE_MAP = {
    "Owned: Owns outright": "Owner occupied",
    "Owned: Owns with a mortgage or loan": "Owner occupied",
    "Shared ownership: Shared ownership": "Owner occupied",
    "Social rented: Rents from council or Local Authority": "Social rent",
    "Social rented: Other social rented": "Social rent",
    "Private rented: Private landlord or letting agency": "Private rent",
    "Private rented: Other private rented": "Private rent",
    "Lives rent free": "Private rent",
}

# Output buckets in a fixed order (own -> social rent -> private rent), so the
# stacked composition reads left-to-right and the largest-remainder rounding
# below is deterministic regardless of pivot column ordering.
OUTPUT_BUCKETS = ["Owner occupied", "Social rent", "Private rent"]
assert set(TENURE_MAP.values()) == set(OUTPUT_BUCKETS), (
    "TENURE_MAP values must be exactly the OUTPUT_BUCKETS"
)


def _tenure_percentages(df: pl.DataFrame) -> pl.DataFrame:
    """Fold the 8 NOMIS tenure leaves into 3-bucket percentages per LSOA.

    `df` is the long-format NOMIS download with columns GEOGRAPHY_CODE,
    C2021_TENURE_9_NAME (the leaf label) and OBS_VALUE (a household count). A
    missing, extra, or relabelled leaf would silently change the denominator (the
    sum over leaves) so we validate the leaf set against TENURE_MAP first and fail
    otherwise. Returns one row per LSOA with `lsoa21` and `% <bucket>` for each
    bucket.
    """
    found = set(df["C2021_TENURE_9_NAME"].unique().to_list())
    expected = set(TENURE_MAP)
    if found != expected:
        missing = sorted(expected - found)
        unexpected = sorted(found - expected)
        raise ValueError(
            "Census tenure categories do not match the expected NOMIS "
            "TS054 C2021_TENURE_9 leaf set.\n"
            f"  expected {len(expected)} categories, found {len(found)}\n"
            f"  missing:    {missing}\n"
            f"  unexpected: {unexpected}\n"
            "Refusing to compute percentages against an unrecognised breakdown."
        )

    # Map each leaf to its output bucket and sum counts per (LSOA, bucket).
    # Summing counts (not rounded percentages) keeps the denominator exact.
    grouped = (
        df.with_columns(
            pl.col("C2021_TENURE_9_NAME").replace_strict(TENURE_MAP).alias("bucket"),
            pl.col("OBS_VALUE").cast(pl.Float64, strict=False).alias("_count"),
        )
        .group_by("GEOGRAPHY_CODE", "bucket")
        .agg(pl.col("_count").sum())
    )
    wide = grouped.pivot(on="bucket", index="GEOGRAPHY_CODE", values="_count").rename(
        {"GEOGRAPHY_CODE": "lsoa21"}
    )

    # A bucket with no households in an LSOA is absent from the long rows, so the
    # pivot leaves a null; treat it as 0 before normalising.
    wide = wide.with_columns(pl.col(OUTPUT_BUCKETS).fill_null(0.0))

    # Normalize so each row sums to exactly 100%, then round with the
    # largest-remainder method to preserve the sum (independent rounding of 3
    # values can drift +/-0.1).
    row_total = sum(pl.col(c) for c in OUTPUT_BUCKETS)
    wide = wide.with_columns(
        [(pl.col(c) / row_total * 100.0).alias(c) for c in OUTPUT_BUCKETS]
    )
    wide = wide.with_columns([pl.col(c).round(1).alias(c) for c in OUTPUT_BUCKETS])
    rounded_sum = sum(pl.col(c) for c in OUTPUT_BUCKETS)
    residual = (100.0 - rounded_sum).round(1)
    largest_col = pl.concat_list(OUTPUT_BUCKETS).list.arg_max()
    wide = wide.with_columns(
        [
            pl.when(largest_col == i)
            .then(pl.col(c) + residual)
            .otherwise(pl.col(c))
            .alias(c)
            for i, c in enumerate(OUTPUT_BUCKETS)
        ]
    )

    rename_map = {col: f"% {col}" for col in OUTPUT_BUCKETS}
    return wide.rename(rename_map).select(
        "lsoa21", *[f"% {c}" for c in OUTPUT_BUCKETS]
    )


def download_and_convert(output_path: Path) -> None:
    print("Downloading Census 2021 household tenure (TS054) by LSOA from NOMIS...")
    df = download_nomis_csv(BASE_URL)
    print(f"Total rows: {df.height}")

    # Filter to England only (E-prefixed LSOA codes); the merge joins on the
    # English postcode universe and the LSOA coverage check is England-wide.
    df = df.filter(pl.col("GEOGRAPHY_CODE").str.starts_with("E"))

    wide = _tenure_percentages(df)

    print(f"England LSOAs: {wide.height}")
    if wide.height != ENGLAND_LSOA_COUNT_2021:
        raise ValueError(
            f"Expected {ENGLAND_LSOA_COUNT_2021} England LSOAs, "
            f"got {wide.height}: truncated NOMIS download?"
        )
    print(f"Columns: {wide.columns}")
    for bucket in OUTPUT_BUCKETS:
        col = wide[f"% {bucket}"]
        print(f"% {bucket}: {col.min()}% - {col.max()}% (mean {col.mean():.1f}%)")

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


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Download Census 2021 household tenure (TS054) by LSOA"
    )
    parser.add_argument(
        "--output", type=Path, required=True, help="Output parquet file path"
    )
    args = parser.parse_args()
    download_and_convert(args.output)


if __name__ == "__main__":
    main()