#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 6 — real-time vintages. For each monthly publication date v (10 days after month end, the measured publication
delay), re-estimate SHARP-HPI using only transactions dated ≤ v − 10 days (filtered mode: one-sided state-space deviations,
stage 1 re-estimated on the available sample). Records every vintage of every cell path in two modes — `filtered` (one-sided: the value of month m never
changes after its first release except through stage-1 re-estimation; the no-revision policy) and `smoothed` (two-sided given the data
available at the vintage: full re-estimation policy) — then the revision statistics
(1st/2nd/3rd release vs final), timeliness and three publication policies: (a) full re-estimation, (b) movement splice of
the newest month onto the frozen history, (c) revision-limited (history frozen after 3 releases).
Writes data/processed/index_vintages.parquet and results/vintages/*.csv."""
from __future__ import annotations

import sys
import time
from datetime import date, timedelta
from pathlib import Path

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

from sharp_hpi import config as C
from sharp_hpi.estimators import month_grid
from sharp_hpi.hedonic import build_features
from sharp_hpi.sharp import cell_log_index, run_sharp

OUT = C.RESULTS / "vintages"
OUT.mkdir(parents=True, exist_ok=True)
PUB_LAG_DAYS = 10
FIRST_VINTAGE = "2022-07"   # 18 months of history before the first simulated release
LEVELS_KEEP = {"province", "region", "municipality"}
MIN_MUN_N = 1500            # municipalities tracked in the vintage database (cost)


def month_end(m: str) -> date:
    y, mm = int(m[:4]), int(m[5:7])
    return (date(y + (mm == 12), (mm % 12) + 1, 1) - timedelta(days=1))


def main() -> None:
    a = pl.read_parquet(C.ANALYTIC_TX)
    f = build_features(a)
    big_mun = set(f.group_by("mun_code").agg(pl.len().alias("n")).filter(pl.col("n") >= MIN_MUN_N)["mun_code"].to_list())
    rows = []
    months_all = month_grid("2021-01", "2026-10")
    vintages = [m for m in months_all if m >= FIRST_VINTAGE]
    for cls in C.PROPERTY_TYPES:
        d = f.filter(pl.col("property_class") == cls)
        for vm in vintages:
            pub = month_end(vm) + timedelta(days=PUB_LAG_DAYS)
            cutoff = pub - timedelta(days=PUB_LAG_DAYS)  # transactions visible at publication
            if cutoff > date.fromisoformat(C.SAMPLE_END):
                cutoff = date.fromisoformat(C.SAMPLE_END)
            avail = d.filter(pl.col("date") <= cutoff)
            grid = month_grid("2021-01", vm)
            t0 = time.time()
            r = run_sharp(avail, grid, realtime_parent=True)
            n_prov = np.array([avail.filter(pl.col("month") == m).height for m in grid], float)
            for mode in ("filtered", "smoothed"):
                rows.append(pl.DataFrame({"class": cls, "vintage": vm, "mode": mode, "level": "province", "geo_id": "QC", "month": grid, "log_index": r.stage1.mu, "se": r.stage1.mu_se, "n": n_prov}))
                for (lvl, gid), c in r.cells.items():
                    if lvl not in LEVELS_KEEP or (lvl == "municipality" and gid not in big_mun):
                        continue
                    p, v = cell_log_index(r, lvl, gid, mode)
                    rows.append(pl.DataFrame({"class": cls, "vintage": vm, "mode": mode, "level": lvl, "geo_id": gid, "month": grid, "log_index": p, "se": np.sqrt(v), "n": c.n}))
            print(cls, vm, f"{avail.height} tx, {time.time() - t0:.1f}s", flush=True)
        pl.concat(rows).write_parquet(C.PROCESSED / "index_vintages.parquet")
    print("done")


if __name__ == "__main__":
    main()