#!/usr/bin/env python3
# =============================================================================
# QWHPI — Quebec Weekly Housing Price Index
# Author  : Simon-Pierre Boucher
# Contact : contact@spboucher.ai
# File    : engine/scripts/11_monthly_canonical.py
# Purpose : v2.1 canonical dataset — qhpi_monthly.parquet with growth
#           horizons, reliability grades, dollar values, vintages, coverage.
# =============================================================================
"""Canonical monthly dataset build (v2.1).

Reliability grades are anchored on the monthly downsampling experiment
(outputs/tables/downsampling_monthly.csv): smoothed-path RMSE ≈ 1.0–1.2% at
150+ tx/month, ≈1.5% at 75, ≈1.9% at 40, ≥2.6% at 10.
  A: median ≥ 150 tx/month and se_log < 1.5%
  B: median ≥ 75 and se_log < 2.5%
  C: median ≥ 40 (direct estimation floor)
  D: median ≥ 15 (hierarchical, heavy shrinkage)
  E: < 15 (model-implied, transparency only)
"""

from __future__ import annotations

import datetime as dt
import sys
import unicodedata
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import numpy as np
import polars as pl

from qwhpi.clean import CLEAN_PARQUET, research_sample
from qwhpi.config import (
    CITY_TARGETS, INTERIM_DIR, PROCESSED_DIR, PROPERTY_TYPES, TABLES_DIR,
    ensure_dirs,
)

MODEL_VERSION = "2.1.0"
CANONICAL = PROCESSED_DIR / "qhpi_monthly.parquet"
FIRST_RELEASE = PROCESSED_DIR / "first_release_monthly.parquet"
VINTAGE_DIR = PROCESSED_DIR / "vintages_monthly"
KEY = ["geography_id", "property_type", "period"]

GROWTH = {"monthly_pct": 1, "three_month_pct": 3, "six_month_pct": 6, "yoy_pct": 12}


def slug(name: str) -> str:
    s = unicodedata.normalize("NFKD", name).encode("ascii", "ignore").decode()
    return s.lower().replace(" ", "-").replace("'", "")


def city_slug(name: str) -> str:
    return "quebec-city" if name == "Québec" else slug(name)


def main() -> None:
    ensure_dirs()
    data_vintage = dt.date.today().isoformat()
    m = pl.read_parquet(INTERIM_DIR / "monthly_indexes.parquet")
    cl = research_sample(pl.read_parquet(CLEAN_PARQUET)).with_columns(
        pl.col("date").dt.strftime("%Y-%m").alias("period"))

    # ---------------- transactions per published cell ---------------- #
    cell_frames = [
        cl.select(pl.lit("quebec").alias("geography_id"),
                  pl.col("propertyType").alias("property_type"),
                  "period", "log_amount", "amount"),
        cl.filter(pl.col("region_code").is_not_null()).select(
            pl.concat_str([pl.lit("region-"), pl.col("region_code")]).alias("geography_id"),
            pl.col("propertyType").alias("property_type"),
            "period", "log_amount", "amount"),
    ]
    for city in CITY_TARGETS:
        sub = cl.filter(pl.col("municipality") == city)
        if sub.height:
            cell_frames.append(sub.select(
                pl.lit(city_slug(city)).alias("geography_id"),
                pl.col("propertyType").alias("property_type"),
                "period", "log_amount", "amount"))
    tx = pl.concat(cell_frames)
    tx_all = pl.concat([tx, tx.with_columns(pl.lit("all").alias("property_type"))])

    # ---------------- representative dollar value ---------------- #
    basket = (
        tx_all.join(m.select(*KEY, "index_smoothed"), on=KEY, how="inner")
        .group_by(["geography_id", "property_type"])
        .agg((pl.col("log_amount") - (pl.col("index_smoothed") / 100.0).log())
             .mean().exp().round(0).alias("basket_value"))
    )
    canon = m.join(basket, on=["geography_id", "property_type"], how="left") \
        .with_columns((pl.col("basket_value") * pl.col("index_smoothed") / 100.0)
                      .round(0).alias("representative_value"))

    # ---------------- growth horizons ---------------- #
    keys = ["geography_id", "property_type"]
    canon = canon.sort([*keys, "period"]).with_columns([
        ((pl.col("index_smoothed") / pl.col("index_smoothed").shift(h).over(keys) - 1) * 100)
        .round(3).alias(name)
        for name, h in GROWTH.items()
    ])

    # ---------------- reliability grades ---------------- #
    med = canon.group_by(keys).agg(
        pl.col("transactions").fill_nan(None).median().alias("cell_median_tx"))
    canon = canon.join(med, on=keys, how="left")
    base_grade = (
        pl.when((pl.col("cell_median_tx") >= 150) & (pl.col("se_log") < 0.015)).then(pl.lit("A"))
        .when((pl.col("cell_median_tx") >= 75) & (pl.col("se_log") < 0.025)).then(pl.lit("B"))
        .when(pl.col("cell_median_tx") >= 40).then(pl.lit("C"))
        .when(pl.col("cell_median_tx") >= 15).then(pl.lit("D"))
        .otherwise(pl.lit("E"))
    )
    downgrade = {"A": "B", "B": "C", "C": "D", "D": "E", "E": "E"}
    canon = canon.with_columns(
        pl.when(pl.col("transactions").fill_nan(0) == 0)
        .then(base_grade.replace(downgrade)).otherwise(base_grade)
        .alias("reliability_grade"))

    # ---------------- partial month (registration lag) ---------------- #
    max_date = cl["date"].max()
    last_period = max_date.strftime("%Y-%m")
    month_end = (dt.date(max_date.year + (max_date.month == 12),
                         max_date.month % 12 + 1, 1) - dt.timedelta(days=1))
    is_partial_last = max_date < month_end
    canon = canon.with_columns(
        ((pl.col("period") == last_period) & is_partial_last).alias("is_partial_month"))

    # ---------------- effective sample size ---------------- #
    fp = pl.read_parquet(INTERIM_DIR / "rtd_residuals.parquet")
    sigma2 = {t: float(fp.filter(pl.col("propertyType") == t)["resid"].var())
              for t in PROPERTY_TYPES}
    sigma2["all"] = float(np.mean(list(sigma2.values())))
    canon = canon.with_columns(
        pl.col("property_type").replace_strict(sigma2, default=None).alias("_s2")
    ).with_columns(
        (pl.col("_s2") / (pl.col("se_log") ** 2).clip(1e-9, None))
        .round(1).alias("effective_sample_size")
    ).drop("_s2")

    canon = canon.with_columns(
        pl.lit(MODEL_VERSION).alias("model_version"),
        pl.lit(data_vintage).alias("data_vintage"),
        pl.lit("monthly").alias("frequency"),
        pl.col("transactions").fill_nan(0).cast(pl.Int64),
    )
    cols = ["period", "geography_level", "geography_id", "geography_name",
            "property_type", "index", "index_smoothed", "representative_value",
            "transactions", "effective_sample_size",
            "monthly_pct", "three_month_pct", "six_month_pct", "yoy_pct",
            "lower_95", "upper_95", "reliability_grade", "shrinkage_weight",
            "is_partial_month", "model_version", "data_vintage"]
    canonical = canon.select(cols).sort(
        ["geography_level", "geography_id", "property_type", "period"])
    canonical.write_parquet(CANONICAL)
    print(f"Canonical: {canonical.height:,} rows -> {CANONICAL}")

    # ---------------- liquidity + coverage ---------------- #
    liq = (tx_all.group_by([*keys, "period"]).len().rename({"len": "transactions"})
           .sort([*keys, "period"]))
    liq.write_parquet(PROCESSED_DIR / "monthly_liquidity.parquet")

    months_total = cl["period"].n_unique()
    cov = (
        cl.group_by(["municipality", "geo_code", "region", "propertyType", "period"]).len()
        .group_by(["municipality", "geo_code", "region", "propertyType"])
        .agg(pl.col("len").median().alias("median_monthly_tx"),
             pl.col("len").sum().alias("total_tx"),
             (pl.len() / months_total * 100).round(1).alias("pct_months_active"))
        .with_columns(
            pl.when(pl.col("municipality").is_in(list(CITY_TARGETS)))
            .then(pl.lit("published"))
            .when((pl.col("median_monthly_tx") >= 15) & (pl.col("pct_months_active") >= 85))
            .then(pl.lit("conditional"))
            .otherwise(pl.lit("not_published"))
            .alias("coverage_status"))
        .sort("total_tx", descending=True)
    )
    cov.write_csv(TABLES_DIR / "coverage_monthly.csv")
    cov.write_parquet(PROCESSED_DIR / "coverage_matrix_monthly.parquet")
    print("Coverage:", dict(cov.group_by("coverage_status").len().iter_rows()))

    # ---------------- assessment gap (monthly) ---------------- #
    gap = cl.filter(pl.col("totalArValue") > 0).with_columns(
        (pl.col("amount") / pl.col("totalArValue")).alias("ratio"))
    gap_cells = pl.concat([
        gap.select(pl.lit("quebec").alias("geography_id"),
                   pl.col("propertyType").alias("property_type"), "period", "ratio"),
        gap.select(pl.concat_str([pl.lit("region-"), pl.col("region_code")])
                   .alias("geography_id"),
                   pl.col("propertyType").alias("property_type"), "period", "ratio"),
    ])
    (gap_cells.group_by([*keys, "period"] if False else ["geography_id", "property_type", "period"])
     .agg(pl.col("ratio").median().round(4).alias("median_ratio"), pl.len().alias("n"))
     .sort(["geography_id", "property_type", "period"])
     .write_parquet(PROCESSED_DIR / "assessment_gap_monthly.parquet"))

    # ---------------- vintages (monthly store) ---------------- #
    vd = VINTAGE_DIR / f"data_vintage={data_vintage}"
    vd.mkdir(parents=True, exist_ok=True)
    canonical.write_parquet(vd / "qhpi_monthly.parquet")
    fresh = canonical.select(
        *KEY,
        pl.col("index").alias("first_release_index"),
        pl.col("index_smoothed").alias("first_release_index_smoothed"),
        pl.lit(data_vintage).alias("first_release_vintage"))
    if FIRST_RELEASE.exists():
        existing = pl.read_parquet(FIRST_RELEASE)
        fresh = pl.concat([existing, fresh.join(existing.select(KEY), on=KEY, how="anti")])
    fresh.write_parquet(FIRST_RELEASE)

    latest_period = canonical.filter(~pl.col("is_partial_month"))["period"].max()
    canonical.filter(pl.col("period") == latest_period) \
        .write_csv(TABLES_DIR / "index_latest_monthly.csv")
    print(f"Latest complete month: {latest_period}")


if __name__ == "__main__":
    main()