#!/usr/bin/env python3
# =============================================================================
# QWHPI — Quebec Weekly Housing Price Index
# Author  : Simon-Pierre Boucher
# Contact : contact@spboucher.ai
# File    : engine/scripts/07_canonical.py
# Purpose : Step 7 — assemble data/processed/qwhpi_weekly.parquet (canonical
#           dataset), coverage matrix, assessment gap, vintages, liquidity.
# =============================================================================
"""Canonical dataset build (Execution Order step 7, schema §4.9)."""

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,
)
from qwhpi.index import (
    add_growth, add_reliability, add_representative_value, compute_basket,
)
from qwhpi.nowcast import partial_week_flags
from qwhpi.seasonal import seasonality_test
from qwhpi.vintages import snapshot, update_first_release

MODEL_VERSION = "0.1.0"
CANONICAL = PROCESSED_DIR / "qwhpi_weekly.parquet"


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()
    hier = pl.read_parquet(INTERIM_DIR / "hierarchical_indexes.parquet")
    cl = research_sample(pl.read_parquet(CLEAN_PARQUET))
    cl = cl.with_columns(pl.col("week").dt.strftime("%Y-%m-%d").alias("week_str"))

    # ------------------------------------------------------------------ #
    # Map transactions to geography_ids (for baskets & liquidity)
    # ------------------------------------------------------------------ #
    cell_frames = []
    for gid_expr, filt in [
        (pl.lit("quebec"), pl.lit(True)),
        (pl.concat_str([pl.lit("region-"), pl.col("region_code")]),
         pl.col("region_code").is_not_null()),
    ]:
        cell_frames.append(
            cl.filter(filt).select(
                gid_expr.alias("geography_id"),
                pl.col("propertyType").alias("property_type"),
                pl.col("week_str").alias("week"),
                "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"),
                pl.col("week_str").alias("week"),
                "log_amount", "amount",
            ))
    tx_cells = pl.concat(cell_frames)
    tx_cells_all = pl.concat([
        tx_cells,
        tx_cells.with_columns(pl.lit("all").alias("property_type")),
    ])

    # ------------------------------------------------------------------ #
    # Canonical columns
    # ------------------------------------------------------------------ #
    basket = compute_basket(tx_cells_all, hier)
    canon = add_representative_value(hier, basket)
    canon = add_growth(canon)
    canon = add_reliability(canon)

    prov_vol = (
        cl.group_by("week_str").len()
        .rename({"week_str": "week", "len": "volume"})
    )
    canon = canon.join(partial_week_flags(prov_vol), on="week", how="left") \
        .with_columns(pl.col("is_partial_week").fill_null(False))

    # Effective sample size: information units of the filtered estimate.
    sigma2_map = {
        t: float(pl.read_parquet(INTERIM_DIR / "stage1_residuals.parquet")
                 .filter(pl.col("propertyType") == t)["resid"].var())
        for t in PROPERTY_TYPES
    }
    sigma2_map["all"] = float(np.mean(list(sigma2_map.values())))
    canon = canon.with_columns(
        pl.col("property_type").replace_strict(sigma2_map, default=None)
        .alias("_sigma2")
    ).with_columns(
        (pl.col("_sigma2") / (pl.col("se_log") ** 2).clip(1e-9, None))
        .round(1).alias("effective_sample_size")
    ).drop("_sigma2")

    canon = canon.with_columns(
        pl.lit(MODEL_VERSION).alias("model_version"),
        pl.lit(data_vintage).alias("data_vintage"),
        pl.col("transactions").fill_nan(0).cast(pl.Int64),
    )

    cols = ["week", "geography_level", "geography_id", "geography_name",
            "property_type", "index", "index_smoothed", "representative_value",
            "transactions", "effective_sample_size",
            "weekly_pct", "four_week_pct", "thirteen_week_pct",
            "twenty_six_week_pct", "yoy_pct",
            "lower_95", "upper_95", "reliability_grade", "shrinkage_weight",
            "is_partial_week", "model_version", "data_vintage"]
    canonical = canon.select(cols).sort(
        ["geography_level", "geography_id", "property_type", "week"])
    canonical.write_parquet(CANONICAL)
    print(f"Canonical table: {canonical.height:,} rows -> {CANONICAL}")

    # ------------------------------------------------------------------ #
    # Coverage matrix (all municipalities × types)
    # ------------------------------------------------------------------ #
    weeks_total = cl["week_str"].n_unique()
    liq = (
        cl.group_by(["municipality", "geo_code", "region", "propertyType", "week_str"])
        .len()
        .group_by(["municipality", "geo_code", "region", "propertyType"])
        .agg(pl.col("len").median().alias("median_weekly_tx"),
             pl.col("len").sum().alias("total_tx"),
             (pl.len() / weeks_total * 100).round(1).alias("pct_weeks_active"))
    )
    published_cities = {city_slug(c) for c in CITY_TARGETS}
    liq = liq.with_columns(
        pl.when(pl.col("municipality").is_in(list(CITY_TARGETS)))
        .then(pl.lit("published"))
        .when((pl.col("median_weekly_tx") >= 5) & (pl.col("pct_weeks_active") >= 80))
        .then(pl.lit("conditional"))
        .otherwise(pl.lit("not_published"))
        .alias("coverage_status")
    ).sort("total_tx", descending=True)
    liq.write_csv(TABLES_DIR / "coverage.csv")
    liq.write_parquet(PROCESSED_DIR / "coverage_matrix.parquet")
    counts = liq.group_by("coverage_status").len()
    print("Coverage matrix:", dict(counts.iter_rows()))

    # Weekly liquidity per published cell.
    weekly_liq = (
        tx_cells_all.group_by(["geography_id", "property_type", "week"]).len()
        .rename({"len": "transactions"}).sort(["geography_id", "property_type", "week"])
    )
    weekly_liq.write_parquet(PROCESSED_DIR / "weekly_liquidity.parquet")
    (
        weekly_liq.group_by(["geography_id", "property_type"])
        .agg(pl.col("transactions").median().alias("median"),
             pl.col("transactions").mean().round(1).alias("mean"))
        .sort(["geography_id", "property_type"])
        .write_csv(TABLES_DIR / "weekly_liquidity.csv")
    )

    # ------------------------------------------------------------------ #
    # Assessment Gap Index (secondary module — never mixed with QWHPI)
    # ------------------------------------------------------------------ #
    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"),
                   pl.col("week_str").alias("week"), "ratio"),
        gap.select(pl.concat_str([pl.lit("region-"), pl.col("region_code")])
                   .alias("geography_id"),
                   pl.col("propertyType").alias("property_type"),
                   pl.col("week_str").alias("week"), "ratio"),
    ])
    assessment = (
        gap_cells.group_by(["geography_id", "property_type", "week"])
        .agg(pl.col("ratio").median().round(4).alias("median_ratio"),
             pl.len().alias("n"))
        .sort(["geography_id", "property_type", "week"])
    )
    assessment.write_parquet(PROCESSED_DIR / "assessment_gap.parquet")
    print(f"Assessment gap: {assessment.height:,} rows")

    # ------------------------------------------------------------------ #
    # Seasonality check (documents the NSA decision)
    # ------------------------------------------------------------------ #
    seas_rows = []
    for gid, ptype in [("quebec", "all"), ("quebec", "unifamilial"),
                       ("montreal", "condo")]:
        s = canonical.filter((pl.col("geography_id") == gid)
                             & (pl.col("property_type") == ptype)).sort("week")
        res = seasonality_test(np.log(s["index_smoothed"].to_numpy()))
        seas_rows.append({"cell": f"{gid}/{ptype}", **res})
    seas = pl.DataFrame(seas_rows)
    seas.write_csv(TABLES_DIR / "seasonality_tests.csv")
    print(seas)

    # ------------------------------------------------------------------ #
    # Vintages + latest snapshot table
    # ------------------------------------------------------------------ #
    snapshot(canonical, data_vintage)
    update_first_release(canonical, data_vintage)

    latest_week = canonical.filter(~pl.col("is_partial_week"))["week"].max()
    latest = canonical.filter(pl.col("week") == latest_week)
    latest.write_csv(TABLES_DIR / "index_latest.csv")
    print(f"Latest complete week: {latest_week} | {latest.height} series")


if __name__ == "__main__":
    main()