"""Shared data preparation for the TechnicalGIS series "Spatial regression in Python".

Builds data/west_yorkshire_msoa.gpkg: the 2021 MSOAs of West Yorkshire with
  - the ONS median price paid for all dwellings, year ending March 2026 (formerly HPSSA dataset 2),
  - four Census 2021 variables from the Nomis bulk tables (TS067, TS044, TS045, TS006),
  - the ONS Open Geography 2021 MSOA boundaries (generalised, clipped) in British National Grid.

Run it on its own with `python wy_data.py`, or import load_west_yorkshire() from either part of the series.
Data: Office for National Statistics and HM Land Registry, Census 2021 via Nomis, ONS Open Geography Portal
(all Open Government Licence v3.0; boundaries contain OS data (c) Crown copyright and database right).
"""
import zipfile
from pathlib import Path

import numpy as np
import pandas as pd
import geopandas as gpd
import requests

DATA = Path("data")
CRS = "EPSG:27700"
GPKG = DATA / "west_yorkshire_msoa.gpkg"

# West Yorkshire's five metropolitan districts
LADS = {"E08000032": "Bradford", "E08000033": "Calderdale", "E08000034": "Kirklees",
        "E08000035": "Leeds", "E08000036": "Wakefield"}

PRICE_URL = ("https://www.ons.gov.uk/file?uri=/peoplepopulationandcommunity/housing/datasets/"
             "medianhousepricesbymiddlelayersuperoutputarea/yearendingmarch2026/medianpricepaidformsoa.xlsx")
PRICE_COL = "Year ending Mar 2026"
NOMIS = "https://www.nomisweb.co.uk/output/census/2021/census2021-{t}.zip"
MSOA_URL = ("https://services1.arcgis.com/ESMARspQHYMw9BZ9/arcgis/rest/services/"
            "Middle_layer_Super_Output_Areas_December_2021_Boundaries_EW_BGC_V3/FeatureServer/0/query")


def download(url, path, **params):
    if not path.exists():
        r = requests.get(url, params=params or None, timeout=900, headers={"User-Agent": "Mozilla/5.0 (TechnicalGIS tutorial)"})
        r.raise_for_status()
        path.write_bytes(r.content)
    return path


def nomis_msoa(table):
    """One Census 2021 bulk table at MSOA level, indexed by the 2021 MSOA code."""
    z = download(NOMIS.format(t=table), DATA / f"census2021-{table}.zip")
    with zipfile.ZipFile(z) as zf:
        df = pd.read_csv(zf.open(f"census2021-{table}-msoa.csv"))
    return df.set_index("geography code")


def build():
    DATA.mkdir(exist_ok=True)

    # 1. median price paid, all dwellings, by 2021 MSOA (table 1a); "[x]" = fewer than 5 sales
    xlsx = download(PRICE_URL, DATA / "medianpricepaidformsoa_ye_mar2026.xlsx")
    price = pd.read_excel(xlsx, sheet_name="1a", header=2, usecols=["Local authority code", "Local authority name",
                                                                    "MSOA code", "MSOA name", PRICE_COL])
    price = price[price["Local authority code"].isin(LADS)].rename(columns={
        "Local authority code": "lad_code", "Local authority name": "lad", "MSOA code": "msoa", "MSOA name": "name",
        PRICE_COL: "price"})
    price["price"] = pd.to_numeric(price.price, errors="coerce")

    # 2. Census 2021 explanatory variables (shares in per cent)
    q = nomis_msoa("ts067")
    acc = nomis_msoa("ts044")
    car = nomis_msoa("ts045")
    dens = nomis_msoa("ts006")
    census = pd.DataFrame({
        "level4": 100 * q["Highest level of qualification: Level 4 qualifications and above"]
        / q["Highest level of qualification: Total: All usual residents aged 16 years and over"],
        "detached": 100 * acc["Accommodation type: Detached"] / acc["Accommodation type: Total: All households"],
        "no_car": 100 * car["Number of cars or vans: No cars or vans in household"]
        / car["Number of cars or vans: Total: All households"],
        "density": dens["Population Density: Persons per square kilometre; measures: Value"],
    })

    # 3. 2021 MSOA boundaries (generalised to 20 m, clipped to the coastline), in British National Grid
    where = " OR ".join(f"MSOA21NM LIKE '{n} %'" for n in LADS.values())
    geo = gpd.read_file(download(MSOA_URL, DATA / "msoa2021_west_yorkshire.geojson", where=where,
                                 outFields="MSOA21CD,MSOA21NM", outSR=27700, f="geojson"))
    geo = geo.set_crs(CRS, allow_override=True).rename(columns={"MSOA21CD": "msoa", "MSOA21NM": "msoa_nm"})

    gdf = geo.merge(price, on="msoa", how="inner").merge(census, left_on="msoa", right_index=True, how="left")
    gdf["ln_price"] = np.log(gdf.price)
    gdf["ln_density"] = np.log(gdf.density)
    gdf = gdf.sort_values("msoa").reset_index(drop=True)
    gdf[["msoa", "msoa_nm", "name", "lad_code", "lad", "price", "ln_price", "level4", "detached", "no_car",
         "density", "ln_density", "geometry"]].to_file(GPKG, layer="msoa", driver="GPKG")
    return gdf


def load_west_yorkshire(drop_missing=True):
    """Return the West Yorkshire MSOA GeoDataFrame, building it on the first call."""
    gdf = gpd.read_file(GPKG, layer="msoa") if GPKG.exists() else build()
    if drop_missing:
        gdf = gdf.dropna(subset=["price", "level4", "detached", "no_car", "density"]).reset_index(drop=True)
    return gdf


if __name__ == "__main__":
    g = load_west_yorkshire(drop_missing=False)
    print(f"{len(g)} MSOAs in West Yorkshire; {g.price.isna().sum()} without a median price; written to {GPKG}")