# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Build the analytic layer from the WP 22 clean-transactions snapshot (real data) or from the synthetic table.

Real-data rules (``configs/data.yaml`` → ``sample``):
* keep ``in_headline_sample`` (WP 22 cleaning: residential, roll-matched, non-duplicate, ratio not an outlier, …);
* ``property_class`` ∈ {house, condo, plex}; non-null coordinates and ``parcel_id``; price in [price_min, price_max];
* roll attributes are mapped from ``roll_*`` to the contract names; the **assessment in force at the sale**
  (``assessment_at_sale``, recorded by the provider at the time of the sale) replaces the current roll total;
  the land/building split of the current roll is rescaled to that total (the split of the earlier roll is not in
  the snapshot); the roll reference date is the one of the roll in force at the sale, inferred from the triennial
  cycle (reference date + 18 months = entry into force);
* ``urban_class`` from the MAMH municipal population (≥ 100 000 metro, ≥ 10 000 suburban, else rural).

Everything downstream (features, agents, evaluation) consumes only the resulting table.
"""
from __future__ import annotations

import json
import time
from datetime import date
from pathlib import Path
from typing import Any

import polars as pl

from ..config import PROCESSED, RAW, RESULTS, SYNTHETIC, load_config
from .analytic_schema import ANALYTIC_COLUMNS, add_derived_columns, coerce_analytic, validate_analytic

URBAN_METRO_POP = 100_000
URBAN_SUBURBAN_POP = 10_000
ROLL_ENTRY_LAG_MONTHS = 18   # Quebec triennial rolls: market reference date = 1 July, entry into force 18 months later
ROLL_CYCLE_YEARS = 3


def urban_class_expr(pop_col: str = "mamh_population") -> pl.Expr:
    return (
        pl.when(pl.col(pop_col) >= URBAN_METRO_POP).then(pl.lit("metro"))
        .when(pl.col(pop_col) >= URBAN_SUBURBAN_POP).then(pl.lit("suburban"))
        .otherwise(pl.lit("rural"))
    )


def roll_reference_at_sale_expr(ref_col: str = "roll_reference_date", date_col: str = "date") -> pl.Expr:
    """Reference date of the roll in force at the sale: step back whole triennial cycles from the snapshot roll's
    reference date until reference + 18 months ≤ sale date. Null when the reference date is unknown."""
    ref = pl.col(ref_col).str.to_date("%Y-%m-%d", strict=False)
    entry = ref.dt.offset_by(f"{ROLL_ENTRY_LAG_MONTHS}mo")
    # number of 3-year cycles to step back (0 when the snapshot roll was already in force at the sale)
    months_short = (entry.dt.year() * 12 + entry.dt.month()) - (pl.col(date_col).dt.year() * 12 + pl.col(date_col).dt.month())
    # entry > sale when months_short > 0, or same month but later day
    later = (months_short > 0) | ((months_short == 0) & (entry.dt.day() > pl.col(date_col).dt.day()))
    k = pl.when(later).then(((months_short - 1) // (ROLL_CYCLE_YEARS * 12)) + 1).otherwise(0).clip(0, None)
    out = ref.dt.offset_by(pl.format("-{}y", k * ROLL_CYCLE_YEARS))
    return out.dt.strftime("%Y-%m-%d")


def _funnel(raw: pl.DataFrame, cfg: dict[str, Any]) -> tuple[pl.DataFrame, list[dict[str, Any]]]:
    s = cfg["sample"]
    steps: list[tuple[str, pl.Expr]] = [
        ("in_headline_sample", pl.col("in_headline_sample").fill_null(False)),
        ("property_class_in_scope", pl.col("property_class").is_in(s["property_classes"])),
        ("coordinates_present", pl.col("lat").is_not_null() & pl.col("lng").is_not_null()),
        ("parcel_id_present", pl.col("parcel_id").is_not_null()),
        ("price_in_range", (pl.col("price") >= s["price_min"]) & (pl.col("price") <= s["price_max"])),
    ]
    if s.get("require_roll_match", True) and "roll_matched" in raw.columns:
        steps.append(("roll_matched", pl.col("roll_matched").fill_null(False)))
    funnel: list[dict[str, Any]] = [{"step": "raw", "rows": raw.height, "dropped": 0}]
    df = raw
    for name, expr in steps:
        before = df.height
        df = df.filter(expr)
        funnel.append({"step": name, "rows": df.height, "dropped": before - df.height})
    return df, funnel


def build_from_real(raw: pl.DataFrame, cfg: dict[str, Any] | None = None) -> tuple[pl.DataFrame, dict[str, Any]]:
    """WP 22 clean transactions → analytic table + build summary."""
    cfg = cfg or load_config("data")
    df, funnel = _funnel(raw, cfg)
    scale = pl.when(pl.col("roll_total_value") > 0).then(pl.col("assessment_at_sale") / pl.col("roll_total_value")).otherwise(None)
    df = df.with_columns(
        pl.col("roll_floor_area_m2").alias("floor_area_m2"),
        pl.col("roll_land_area_m2").alias("land_area_m2"),
        pl.col("roll_year_built").cast(pl.Int64).alias("year_built"),
        pl.col("roll_storeys").alias("storeys"),
        pl.col("roll_n_units").cast(pl.Int64).alias("n_units"),
        (pl.col("roll_land_value") * scale).round(0).alias("roll_land_value"),
        (pl.col("roll_building_value") * scale).round(0).alias("roll_building_value"),
        pl.col("roll_total_value").alias("_snapshot_roll_total"),
        pl.coalesce(pl.col("assessment_at_sale"), pl.col("roll_total_value")).alias("roll_total_value"),
        pl.col("roll_reference_date").alias("_snapshot_roll_ref"),
        roll_reference_at_sale_expr().alias("roll_reference_date"),
        urban_class_expr().alias("urban_class"),
        pl.col("subtype").fill_null("unknown"),
    )
    # diagnostic: does the inferred roll cycle agree with the recorded assessment? (k=0 ⇔ assessment == snapshot roll)
    diag_df = df.select(
        (pl.col("roll_reference_date") == pl.col("_snapshot_roll_ref")).alias("k0"),
        (pl.col("roll_total_value") == pl.col("_snapshot_roll_total")).alias("eq_cur"),
        (pl.col("roll_total_value") == pl.col("roll_previous_value")).alias("eq_prev"),
    )
    roll_diag = {
        "share_inferred_current_roll": float(diag_df["k0"].mean()),
        "share_assessment_equals_snapshot_roll": float(diag_df["eq_cur"].mean()),
        "share_assessment_equals_previous_roll": float(diag_df["eq_prev"].mean()),
        "agreement_k0_vs_eq_cur": float((diag_df["k0"] == diag_df["eq_cur"]).mean()),
        "note": "roll_total_value = assessment recorded at the sale (roll in force); reference date inferred from the triennial cycle",
    }
    df = add_derived_columns(df, recompute_xy=False)
    out = coerce_analytic(df)
    validate_analytic(out)
    summary = summarize(out)
    summary["funnel"] = funnel
    summary["roll_as_of_diagnostic"] = roll_diag
    summary["source"] = cfg["snapshot"]["transactions_parquet"]
    return out, summary


def build_from_synthetic(syn: pl.DataFrame) -> tuple[pl.DataFrame, dict[str, Any]]:
    """Synthetic generator output (24 columns) → analytic table (x/y projected with pyproj)."""
    df = add_derived_columns(syn, recompute_xy=True)
    out = coerce_analytic(df)
    validate_analytic(out)
    summary = summarize(out)
    summary["source"] = "data/synthetic/synthetic_transactions.parquet (SYNTHETIC)"
    return out, summary


def summarize(df: pl.DataFrame) -> dict[str, Any]:
    """Row counts by year / class / region group, missingness, tercile cut points, date range."""
    def counts(col: pl.Expr, name: str) -> dict[str, int]:
        t = df.group_by(col.alias(name)).len().sort(name)
        return {str(k): int(v) for k, v in zip(t[name].to_list(), t["len"].to_list())}

    terciles = (
        df.group_by("property_class").agg(
            pl.col("price").quantile(1 / 3, interpolation="linear").alias("q33"),
            pl.col("price").quantile(2 / 3, interpolation="linear").alias("q67"),
        ).sort("property_class")
    )
    year_class = df.group_by(pl.col("date").dt.year().alias("year"), "property_class").len().sort(["year", "property_class"])
    return {
        "rows": df.height,
        "parcels": int(df["parcel_id"].n_unique()),
        "date_min": df["date"].min().isoformat(),
        "date_max": df["date"].max().isoformat(),
        "by_year": counts(pl.col("date").dt.year(), "year"),
        "by_property_class": counts(pl.col("property_class"), "property_class"),
        "by_region_group": counts(pl.col("region_group"), "region_group"),
        "by_region_name": counts(pl.col("region_name"), "region_name"),
        "by_urban_class": counts(pl.col("urban_class"), "urban_class"),
        "by_year_class": [{"year": int(y), "property_class": c, "rows": int(n)} for y, c, n in year_class.iter_rows()],
        "missingness_pct": {c: round(100.0 * df[c].null_count() / max(df.height, 1), 3) for c in ANALYTIC_COLUMNS},
        "price_tercile_cuts": {r["property_class"]: [r["q33"], r["q67"]] for r in terciles.iter_rows(named=True)},
        "rows_with_prior_sale": int(df["prior_sale_price"].is_not_null().sum()),
        "price_median": float(df["price"].median()),
    }


def run(synthetic: bool = False, raw_path: Path | None = None) -> dict[str, Any]:
    """Script entry point: build, write parquet + summary JSON, return summary with timing."""
    cfg = load_config("data")
    t0 = time.perf_counter()
    if synthetic:
        src = raw_path or SYNTHETIC / "synthetic_transactions.parquet"
        if not src.exists():
            from .synthetic import write_default

            write_default()
        df, summary = build_from_synthetic(pl.read_parquet(src))
        out = SYNTHETIC / "analytic.parquet"
        summary_path = RESULTS / "data" / "analytic_summary_synthetic.json"
    else:
        src = raw_path or RAW / cfg["snapshot"]["transactions_parquet"].split("data/raw/")[-1]
        raw = pl.read_parquet(src)
        df, summary = build_from_real(raw, cfg)
        out = PROCESSED / "analytic.parquet"
        summary_path = RESULTS / "data" / "analytic_summary.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    summary_path.parent.mkdir(parents=True, exist_ok=True)
    df.write_parquet(out, compression="zstd")
    summary["output"] = str(out.relative_to(out.parents[2]))
    summary["built_at"] = date.today().isoformat()
    summary["seconds"] = round(time.perf_counter() - t0, 2)
    summary_path.write_text(json.dumps(summary, indent=1, ensure_ascii=False))
    return summary