#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 1 — geographic concordance.

Spatially joins every transaction (lat/lng) to the MRNF SDA V2026-07 boundaries
(municipality → MRC → administrative region; metropolitan community; borough),
links the MAMH municipal register (population, designation, constitution date)
and the roll's municipality code, and builds fine-location keys
(assessment neighbourhood unit when present, else 1-km / 500-m metric grid cells).
Writes data/interim/geo_transactions.parquet and results/geo/*.csv.
"""
from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
import geopandas as gpd
import numpy as np
import pandas as pd
import polars as pl

from sharp_hpi import config as C

OUT = C.RESULTS / "geo"
OUT.mkdir(parents=True, exist_ok=True)


def main() -> None:
    tx = pl.read_parquet(C.RAW_TX, columns=["id", "date", "lat", "lng", "city", "id_provinc"])
    um = pl.read_parquet(C.RAW_UNITS_MATCHED, columns=["id_provinc", "code_mun", "municipalite", "unite_voisinage", "arrond"])
    tx = tx.join(um, on="id_provinc", how="left")
    pts = gpd.GeoDataFrame(tx.select(["id", "lat", "lng"]).to_pandas(), geometry=gpd.points_from_xy(tx["lng"], tx["lat"]), crs="EPSG:4269")

    mun = gpd.read_file(C.SDA_GPKG, layer="munic_s")[["MUS_CO_GEO", "MUS_NM_MUN", "MUS_CO_DES", "MUS_DE_IND", "MUS_CO_MRC", "MUS_NM_MRC", "MUS_CO_REG", "MUS_NM_REG", "MUS_DA_CON", "MUS_CO_MOD", "MUS_DA_MOD", "geometry"]]
    mun = mun.to_crs("EPSG:4269")
    j = gpd.sjoin(pts, mun, how="left", predicate="within")
    j = j[~j.index.duplicated(keep="first")]
    j["join_method"] = np.where(j["MUS_CO_GEO"].notna(), "within", "unassigned")
    # nearest within 1 km for the unassigned (coastline / water slivers)
    miss = j["MUS_CO_GEO"].isna()
    if miss.any():
        p_m = pts.loc[miss].to_crs(C.CRS_METRIC)
        near = gpd.sjoin_nearest(p_m, mun.to_crs(C.CRS_METRIC), how="left", max_distance=1000, distance_col="dist_m")
        near = near[~near.index.duplicated(keep="first")]
        for c in ["MUS_CO_GEO", "MUS_NM_MUN", "MUS_CO_DES", "MUS_DE_IND", "MUS_CO_MRC", "MUS_NM_MRC", "MUS_CO_REG", "MUS_NM_REG", "MUS_DA_CON", "MUS_CO_MOD", "MUS_DA_MOD"]:
            j.loc[miss, c] = near[c].values
        j.loc[miss, "join_method"] = np.where(near["MUS_CO_GEO"].notna(), "nearest<=1km", "unassigned")

    comet = gpd.read_file(C.SDA_GPKG, layer="comet_s")[["CMS_CO_COM", "CMS_NM_COM", "geometry"]].to_crs("EPSG:4269")
    jc = gpd.sjoin(pts, comet, how="left", predicate="within")
    jc = jc[~jc.index.duplicated(keep="first")]
    arr = gpd.read_file(C.SDA_GPKG, layer="arron_s")[["ARS_CO_ARR", "ARS_NM_ARR", "ARS_NM_MUN", "geometry"]].to_crs("EPSG:4269")
    ja = gpd.sjoin(pts, arr, how="left", predicate="within")
    ja = ja[~ja.index.duplicated(keep="first")]

    # metric grid cells for fine location
    pm = pts.to_crs(C.CRS_METRIC)
    x, y = pm.geometry.x.values, pm.geometry.y.values
    grid1k = (np.floor(x / 1000).astype("int64").astype(str) + "_" + np.floor(y / 1000).astype("int64").astype(str))
    grid500 = (np.floor(x / 500).astype("int64").astype(str) + "_" + np.floor(y / 500).astype("int64").astype(str))
    grid2k = (np.floor(x / 2000).astype("int64").astype(str) + "_" + np.floor(y / 2000).astype("int64").astype(str))
    grid5k = (np.floor(x / 5000).astype("int64").astype(str) + "_" + np.floor(y / 5000).astype("int64").astype(str))

    geo = pl.from_pandas(pd.DataFrame({
        "id": j["id"].values,
        "mun_code": j["MUS_CO_GEO"].values, "mun_name": j["MUS_NM_MUN"].values, "mun_design": j["MUS_CO_DES"].values,
        "mrc_code": j["MUS_CO_MRC"].values, "mrc_name": j["MUS_NM_MRC"].values,
        "region_code": j["MUS_CO_REG"].values, "region_name": j["MUS_NM_REG"].values,
        "mun_date_constituted": pd.to_datetime(j["MUS_DA_CON"]).dt.strftime("%Y-%m-%d").values,
        "mun_date_modified": pd.to_datetime(j["MUS_DA_MOD"]).dt.strftime("%Y-%m-%d").values,
        "join_method": j["join_method"].values,
        "comet_code": jc["CMS_CO_COM"].values, "comet_name": jc["CMS_NM_COM"].values,
        "borough_code": ja["ARS_CO_ARR"].values, "borough_name": ja["ARS_NM_ARR"].values,
        "x_m": x, "y_m": y, "grid500": grid500, "grid1k": grid1k, "grid2k": grid2k, "grid5k": grid5k,
    }).astype({c: object for c in ["mun_code","mun_name","mun_design","mrc_code","mrc_name","region_code","region_name","mun_date_constituted","mun_date_modified","join_method","comet_code","comet_name","borough_code","borough_name","grid500","grid1k","grid2k","grid5k"]}))
    geo = tx.select(["id", "city", "code_mun", "municipalite", "unite_voisinage", "arrond"]).join(geo, on="id", how="left")
    # MAMH register: population, designation, constitution date, official MRC/region labels
    mamh = pl.read_csv(C.MAMH_MUN_CSV, infer_schema_length=0, encoding="utf8-lossy").select(
        pl.col("mcode").alias("mun_code"), pl.col("munnom").alias("mamh_name"), pl.col("mdes").alias("mamh_design"),
        pl.col("regadm").alias("mamh_region"), pl.col("mrc").alias("mamh_mrc"), pl.col("mpopul").cast(pl.Int64, strict=False).alias("mamh_population"),
        pl.col("mdatcons").alias("mamh_date_constituted"), pl.col("msuperf").cast(pl.Float64, strict=False).alias("mamh_area_km2"),
    )
    geo = geo.join(mamh, on="mun_code", how="left")
    # fine-location keys. The roll's `unite_voisinage` turned out to be a unit-level code (379k distinct values for
    # 449k rows), so fine location is a nested metric grid (500 m ⊂ 1 km ⊂ 2 km ⊂ 5 km ⊂ municipality); the
    # estimation code picks the finest level with enough sales (hedonic.assign_loc_fe).
    geo = geo.with_columns(
        pl.concat_str([pl.lit("G1K:"), pl.col("grid1k")]).alias("loc_fine"),
        pl.concat_str([pl.lit("G1K:"), pl.col("grid1k")]).alias("loc_grid1k"),
        pl.concat_str([pl.lit("G500:"), pl.col("grid500")]).alias("loc_grid500"),
        (pl.col("code_mun") == pl.col("mun_code")).alias("roll_mun_agrees"),
    )
    geo.write_parquet(C.GEO_TX)

    # --- diagnostics ---
    geo.group_by("join_method").agg(pl.len().alias("n")).sort("n", descending=True).write_csv(OUT / "join_methods.csv")
    agree = geo.filter(pl.col("code_mun").is_not_null()).select(
        pl.len().alias("n_with_roll_code"), pl.col("roll_mun_agrees").mean().alias("share_roll_code_agrees"),
    )
    agree.write_csv(OUT / "roll_vs_spatial_agreement.csv")
    geo.filter(pl.col("code_mun").is_not_null() & ~pl.col("roll_mun_agrees")).group_by(["code_mun", "municipalite", "mun_code", "mun_name"]).agg(pl.len().alias("n")).sort("n", descending=True).head(50).write_csv(OUT / "roll_vs_spatial_disagreements_top50.csv")
    geo.group_by(["region_code", "region_name"]).agg(pl.len().alias("n"), pl.col("mun_code").n_unique().alias("municipalities"), pl.col("mrc_code").n_unique().alias("mrcs")).sort("region_code").write_csv(OUT / "region_summary.csv")
    geo.group_by(["comet_code", "comet_name"]).agg(pl.len().alias("n")).sort("n", descending=True).write_csv(OUT / "metropolitan_communities.csv")
    # boundary changes since 2021 (municipalities constituted/modified during the sample)
    mun_df = pl.from_pandas(mun.drop(columns="geometry")).with_columns(pl.col("MUS_DA_CON").cast(pl.Utf8), pl.col("MUS_DA_MOD").cast(pl.Utf8))
    mun_df.filter((pl.col("MUS_DA_CON") >= "2021-01-01") | (pl.col("MUS_DA_MOD") >= "2021-01-01")).write_csv(OUT / "boundary_changes_since_2021.csv")
    mamh.filter(pl.col("mamh_date_constituted") >= "2021-01-01").write_csv(OUT / "mamh_constituted_since_2021.csv")
    pl.DataFrame({
        "metric": ["transactions", "share_rows_with_roll_neighbourhood_code", "distinct_grid1k", "distinct_grid2k", "distinct_grid500", "distinct_municipalities", "distinct_mrc"],
        "value": [float(geo.height), float((geo["unite_voisinage"].is_not_null() & (geo["unite_voisinage"] != "")).mean()), float(geo["grid1k"].n_unique()), float(geo["grid2k"].n_unique()), float(geo["grid500"].n_unique()), float(geo["mun_code"].n_unique()), float(geo["mrc_code"].n_unique())],
    }).write_csv(OUT / "fine_location_summary.csv")
    print(geo.group_by("join_method").agg(pl.len()).to_dicts(), agree.to_dicts())


if __name__ == "__main__":
    main()