#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 6 — revision analysis of the real-time vintage database (output of 40_vintages.py).
For each cell and month: 1st, 2nd, 3rd release vs the final retrospective panel (smoothed, full sample) and vs the latest
real-time vintage; revision magnitudes of monthly and 12-month changes; timeliness (first-release nowcast error of the newest
month); publication policies: full re-estimation (revising), movement splice (never revise history), revision-limited (freeze
after 3 releases). Writes results/vintages/revisions_*.csv."""
from __future__ import annotations

import sys
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

OUT = C.RESULTS / "vintages"


def main() -> None:
    v = pl.read_parquet(C.PROCESSED / "index_vintages.parquet")
    modes = v["mode"].unique().to_list() if "mode" in v.columns else ["filtered"]
    if "mode" not in v.columns:
        v = v.with_columns(pl.lit("filtered").alias("mode"))
    for mode in modes:
        _analyse(v.filter(pl.col("mode") == mode), panel_path(), mode)


def panel_path():
    return pl.read_parquet(C.PROCESSED / "index_panel.parquet").select(
        pl.col("property_type").alias("class"), pl.col("geography_level").alias("level"), pl.col("geography_id").alias("geo_id"), "month", (pl.col("index_value") / 100).log().alias("li_final"))


def _analyse(v: pl.DataFrame, panel: pl.DataFrame, mode: str) -> None:
    sfx = "" if mode == "smoothed" else "_" + mode
    # release number of month m in vintage vm: k = months between m and vm (+1)
    def mi(s):
        return (s.str.slice(0, 4).cast(pl.Int32) * 12 + s.str.slice(5, 2).cast(pl.Int32))
    v = v.with_columns((mi(pl.col("vintage")) - mi(pl.col("month")) + 1).alias("release"))
    # monthly change in each vintage (levels are only identified up to a constant per vintage → compare changes)
    v = v.sort(["class", "level", "geo_id", "vintage", "month"]).with_columns(
        (pl.col("log_index") - pl.col("log_index").shift(1)).over(["class", "level", "geo_id", "vintage"]).alias("d1"),
        (pl.col("log_index") - pl.col("log_index").shift(12)).over(["class", "level", "geo_id", "vintage"]).alias("d12"))
    final = panel.sort(["class", "level", "geo_id", "month"]).with_columns(
        (pl.col("li_final") - pl.col("li_final").shift(1)).over(["class", "level", "geo_id"]).alias("d1_final"),
        (pl.col("li_final") - pl.col("li_final").shift(12)).over(["class", "level", "geo_id"]).alias("d12_final"))
    last_v = v.filter(pl.col("vintage") == v["vintage"].max()).select(["class", "level", "geo_id", "month", pl.col("d1").alias("d1_last"), pl.col("d12").alias("d12_last")])
    j = v.join(final, on=["class", "level", "geo_id", "month"], how="left").join(last_v, on=["class", "level", "geo_id", "month"], how="left")
    j = j.with_columns((pl.col("d1") - pl.col("d1_final")).alias("rev1_final"), (pl.col("d12") - pl.col("d12_final")).alias("rev12_final"),
                       (pl.col("d1") - pl.col("d1_last")).alias("rev1_last"), (pl.col("d12") - pl.col("d12_last")).alias("rev12_last"))
    summ = j.filter(pl.col("release") <= 12).group_by(["class", "level", "release"]).agg(
        pl.len().alias("obs"), pl.col("rev1_final").abs().mean().alias("mae_rev_m1_vs_final"), pl.col("rev1_final").mean().alias("bias_rev_m1_vs_final"),
        pl.col("rev12_final").abs().mean().alias("mae_rev_12m_vs_final"), pl.col("rev12_final").mean().alias("bias_rev_12m_vs_final"),
        pl.col("rev1_last").abs().mean().alias("mae_rev_m1_vs_last_rt"), pl.col("rev12_last").abs().mean().alias("mae_rev_12m_vs_last_rt"),
        (pl.col("rev1_final").abs() > 0.01).mean().alias("share_m1_rev_gt_1pct")).sort(["class", "level", "release"])
    summ.write_csv(OUT / f"revisions_by_release{sfx}.csv")
    # policies on the provincial monthly change: (a) full re-estimation = release k value; (b) movement splice = first release of each
    # month chained; (c) revision-limited = value at release 3
    pol = j.filter(pl.col("level") == "province")
    first = pol.filter(pl.col("release") == 1).select(["class", "month", pl.col("d1").alias("d1_first"), pl.col("d1_final")])
    third = pol.filter(pl.col("release") == 3).select(["class", "month", pl.col("d1").alias("d1_third")])
    p = first.join(third, on=["class", "month"], how="left")
    pol_summ = p.group_by("class").agg(pl.len().alias("months"), (pl.col("d1_first") - pl.col("d1_final")).abs().mean().alias("mae_first_release_vs_final"),
                                       (pl.col("d1_first") - pl.col("d1_final")).mean().alias("bias_first_release_vs_final"),
                                       (pl.col("d1_third") - pl.col("d1_final")).abs().mean().alias("mae_third_release_vs_final"),
                                       ((pl.col("d1_first") - pl.col("d1_final")) ** 2).mean().sqrt().alias("rmse_first"), ((pl.col("d1_third") - pl.col("d1_final")) ** 2).mean().sqrt().alias("rmse_third"),
                                       pl.col("d1_final").abs().mean().alias("mean_abs_true_change"))
    pol_summ.write_csv(OUT / f"revision_policies_province{sfx}.csv")
    j.filter(pl.col("release") <= 6).select(["class", "level", "geo_id", "month", "vintage", "release", "n", "log_index", "se", "d1", "d12", "rev1_final", "rev12_final"]).write_parquet(OUT / f"revisions_long{sfx}.parquet")
    print(mode); print(summ); print(pol_summ)


if __name__ == "__main__":
    main()