"""Download Census 2021 highest level of qualification (TS067) by LSOA.

Downloads the 8-category "highest level of qualification" breakdown (TS067,
classification C2021_HIQUAL_8) from the NOMIS API at LSOA 2021 granularity and
emits one row per LSOA with the percentage of usual residents aged 16+ in each
qualification band. The bands sum to 100%, so downstream they render as a
composition/ratio (like the ethnicity and political-vote-share stacked bars)
rather than a single headline number.

We give the ONS bands colloquial labels (the census's "Level 1/2/3/4+" jargon
means little to a homebuyer): No qualifications / Some GCSEs / Good GCSEs /
Apprenticeship / A-levels / Degree or higher / Other qualifications. NOTE the
census does NOT split undergraduate from postgraduate. "Level 4 and above" is a
single bucket ("Degree or higher").

The join key downstream (merge.py) is `lsoa21`, the same key used for ethnicity,
median age, and IoD.

Source: NOMIS (ONS Census 2021, TS067 dataset, NM_2084_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 TS067 (highest level of qualification) by LSOA 2021
# (TYPE151). c2021_hiqual_8=1..7 selects the 7 substantive bands (excluding
# 0 = Total, which we re-derive by summing). measures=20100 selects the count.
BASE_URL = (
    "https://www.nomisweb.co.uk/api/v01/dataset/NM_2084_1.data.csv"
    "?date=latest"
    "&geography=TYPE151"
    "&c2021_hiqual_8=1,2,3,4,5,6,7"
    "&measures=20100"
    "&select=GEOGRAPHY_CODE,C2021_HIQUAL_8_NAME,OBS_VALUE"
)

# Map each canonical NOMIS C2021_HIQUAL_8 band to our colloquial output bucket.
# 1:1 mapping (no folding): keyed on the exact NOMIS label so a relabelled or
# missing band fails loudly in validation instead of silently dropping people.
BAND_MAP = {
    "No qualifications": "No qualifications",
    "Level 1 and entry level qualifications": "Some GCSEs",
    "Level 2 qualifications": "Good GCSEs",
    "Apprenticeship": "Apprenticeship",
    "Level 3 qualifications": "A-levels",
    "Level 4 qualifications or above": "Degree or higher",
    "Other qualifications": "Other qualifications",
}

# Output buckets in ascending-attainment order, so the stacked composition reads
# left-to-right from least to most qualified.
OUTPUT_BUCKETS = [
    "No qualifications",
    "Some GCSEs",
    "Good GCSEs",
    "Apprenticeship",
    "A-levels",
    "Degree or higher",
    "Other qualifications",
]
assert set(BAND_MAP.values()) == set(OUTPUT_BUCKETS), (
    "BAND_MAP values must be exactly the OUTPUT_BUCKETS"
)


def _qualification_percentages(df: pl.DataFrame) -> pl.DataFrame:
    """Fold the 7 NOMIS bands into percentage buckets per LSOA (summing to 100).

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

    grouped = (
        df.with_columns(
            pl.col("C2021_HIQUAL_8_NAME").replace_strict(BAND_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 people 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 7
    # values can drift +/-0.3).
    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 highest qualification (TS067) 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 = _qualification_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}")
    deg = wide["% Degree or higher"]
    print(f"% Degree or higher: {deg.min()}% - {deg.max()}% (mean {deg.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 highest level of qualification (TS067) 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()