#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Phase 1 — registration / publication lag measured from earlier data vintages.

Compares four earlier extractions of the same transaction table (ids + dates)
with the final snapshot. For a vintage extracted on day D, a transaction dated
t that is absent from the vintage but present in the final snapshot arrived
with a lag > D - t. This yields an empirical arrival profile (share of
eventual transactions visible after L days) used to simulate real-time vintages.
Also records deletions and price revisions between vintages.
"""
from __future__ import annotations

import sys
from pathlib import Path

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

from sharp_hpi import config as C

OUT = C.RESULTS / "vintages"
OUT.mkdir(parents=True, exist_ok=True)


def load_vintage(path: Path) -> pl.DataFrame:
    if path.suffix == ".gz":
        import gzip, io
        with gzip.open(path, "rb") as f:
            df = pl.read_csv(io.BytesIO(f.read()), infer_schema_length=0)
    else:
        df = pl.read_parquet(path)
    return df.select(pl.col("id"), pl.col("date").cast(pl.Utf8).str.slice(0, 10).str.to_date("%Y-%m-%d"), pl.col("amount").cast(pl.Float64, strict=False))


def main() -> None:
    final = pl.read_parquet(C.RAW_TX, columns=["id", "date", "amount", "property_type"])
    rows, profile, revisions = [], [], []
    for vdate, path in C.VINTAGES.items():
        v = load_vintage(path)
        vd = pl.lit(vdate).str.to_date("%Y-%m-%d")
        in_final = final.join(v.select("id", pl.col("amount").alias("amount_v")), on="id", how="left")
        # eligible: dated on/before the vintage's last observed date
        last_obs = v["date"].max()
        elig = in_final.filter(pl.col("date") <= last_obs)
        present = elig["amount_v"].is_not_null()
        deleted = v.join(final.select("id"), on="id", how="anti")
        rows.append({
            "vintage": vdate, "vintage_rows": v.height, "vintage_last_date": str(last_obs), "vintage_max_lag_days": (pl.Series([vdate]).str.to_date("%Y-%m-%d") - pl.Series([last_obs]))[0].days,
            "final_rows_dated_le_last": elig.height, "present_in_vintage": int(present.sum()), "share_present": float(present.mean()),
            "deleted_since_vintage": deleted.height,
            "amount_revised": int((elig.filter(present)["amount"] != elig.filter(present)["amount_v"]).sum()),
        })
        # arrival profile by lag (days between transaction date and vintage date)
        prof = elig.with_columns(((vd - pl.col("date")).dt.total_days()).alias("lag_days"), pl.col("amount_v").is_not_null().alias("present")) \
            .with_columns((pl.col("lag_days") // 7 * 7).alias("lag_bin")).group_by("lag_bin").agg(pl.len().alias("n_final"), pl.col("present").mean().alias("share_visible")) \
            .with_columns(pl.lit(vdate).alias("vintage")).sort("lag_bin")
        profile.append(prof)
        # by transaction month
        bym = elig.with_columns(pl.col("date").dt.strftime("%Y-%m").alias("month"), pl.col("amount_v").is_not_null().alias("present")).group_by("month").agg(
            pl.len().alias("n_final"), pl.col("present").sum().alias("n_in_vintage"), pl.col("present").mean().alias("share_visible")).with_columns(pl.lit(vdate).alias("vintage")).sort("month")
        revisions.append(bym)
    pl.DataFrame(rows).write_csv(OUT / "vintage_summary.csv")
    pl.concat(profile).write_csv(OUT / "arrival_profile_by_lag.csv")
    pl.concat(revisions).write_csv(OUT / "arrival_by_month.csv")
    print(pl.DataFrame(rows))


if __name__ == "__main__":
    main()