#!/usr/bin/env python3
# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Forecasting Monte Carlo on synthetic markets (protocol §14). Markets are generated with the WP22 simulator (known latent
quality-constant paths per municipality, nested hierarchy, heterogeneous sizes) under scenarios that differ in *forecastability*:
  rw          pure random-walk latent paths (no predictable component: the no-change forecast is optimal)
  momentum    latent drift follows an AR(1) (persistent momentum)
  regimes     drift switches between two regimes by a Markov chain (regime adaptation should help)
  spillover   each municipality's drift partly follows its region's lagged drift (spatial information is predictive)
  sparse      momentum with 35 % of the sales (real-time estimates noisier)
The "real-time" index at an origin is the WP22 SHARP filtered estimate on the synthetic sales (so estimation noise and
pooling are realistic); the target is the true latent change. Each forecaster is evaluated across origins and horizons.
Usage: python3 scripts/40_simulation.py --reps 20 --family classic|torch|trees"""
from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

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

from wp23 import config as C
from wp23.data import month_add

sys.path.insert(0, str(C.WP22_SRC))
from sharp_hpi.hedonic import build_features  # noqa: E402
from sharp_hpi.sharp import cell_log_index, run_sharp  # noqa: E402
from sharp_hpi.simulate import Scenario, simulate  # noqa: E402

OUT = C.RESULTS / "simulation"; OUT.mkdir(parents=True, exist_ok=True)
T = 72; LAGS = 12; HORIZONS = (1, 3, 6, 12)


def latent_paths(sc: Scenario, kind: str, rng: np.random.Generator, n_regions: int, n_mun: int, mun_region: np.ndarray):
    """Override the WP22 random-walk latent paths with forecastable dynamics. Returns (mu, dR, dU) arrays."""
    T_ = sc.T
    if kind == "rw":
        drift_p = np.zeros(T_) + 0.003; drift_r = np.zeros((n_regions, T_)); drift_u = np.zeros((n_mun, T_))
    elif kind in ("momentum", "sparse"):
        drift_p = np.zeros(T_); drift_r = np.zeros((n_regions, T_)); drift_u = np.zeros((n_mun, T_))
        for t in range(1, T_):
            drift_p[t] = 0.85 * drift_p[t - 1] + rng.normal(0, 0.003)
            drift_r[:, t] = 0.8 * drift_r[:, t - 1] + rng.normal(0, 0.002, n_regions)
            drift_u[:, t] = 0.7 * drift_u[:, t - 1] + rng.normal(0, 0.003, n_mun)
    elif kind == "regimes":
        P = np.array([[0.93, 0.07], [0.10, 0.90]]); s = 0; drift_p = np.zeros(T_); states = np.zeros(T_, int)
        for t in range(T_):
            s = rng.choice(2, p=P[s]); states[t] = s; drift_p[t] = (0.012 if s == 0 else -0.006) + rng.normal(0, 0.002)
        drift_r = np.zeros((n_regions, T_)); drift_u = np.zeros((n_mun, T_))
        for t in range(1, T_):
            drift_r[:, t] = 0.6 * drift_r[:, t - 1] + rng.normal(0, 0.002, n_regions); drift_u[:, t] = 0.5 * drift_u[:, t - 1] + rng.normal(0, 0.003, n_mun)
    elif kind == "spillover":
        drift_p = np.full(T_, 0.004); drift_r = np.zeros((n_regions, T_)); drift_u = np.zeros((n_mun, T_))
        for t in range(1, T_):
            drift_r[:, t] = 0.8 * drift_r[:, t - 1] + rng.normal(0, 0.003, n_regions)
            drift_u[:, t] = 0.3 * drift_u[:, t - 1] + 0.6 * drift_r[mun_region, t - 1] + rng.normal(0, 0.002, n_mun)  # municipalities follow their region with a lag
    mu = np.cumsum(drift_p + rng.normal(0, 0.004, T_)); mu -= mu[0]
    dR = np.cumsum(drift_r + rng.normal(0, 0.002, (n_regions, T_)), axis=1); dR -= dR[:, [0]]
    dU = np.cumsum(drift_u + rng.normal(0, 0.003, (n_mun, T_)), axis=1); dU -= dU[:, [0]]
    return mu, dR, dU


def one_market(kind: str, rep: int):
    rng = np.random.default_rng(5000 + rep)
    sc = Scenario(name=kind, T=T, seed=5000 + rep, scale_sales=0.35 if kind == "sparse" else 1.0, tau_prov=1e-4, tau_region=1e-4, tau_mrc=1e-4, tau_mun=1e-4, tau_mun_hetero=False)
    sim = simulate(sc)
    n_mun = len(sim.mun_codes); n_regions = sim.truth_region.shape[0]
    mu, dR, dU = latent_paths(sc, kind, rng, n_regions, n_mun, sim.mun_region)
    L_new = mu[None, :] + dR[sim.mun_region] + dU
    # re-price the sales: replace the simulator's latent component by the new one (log_price += L_new − L_old)
    old = sim.truth_mun; mi = {c: i for i, c in enumerate(sim.mun_codes)}; months = sim.months
    df = sim.df.with_columns(pl.col("mun_code").replace_strict(mi, default=None).alias("_u"), pl.col("month").replace_strict({m: i for i, m in enumerate(months)}, default=None).alias("_t"))
    u = df["_u"].to_numpy(); t = df["_t"].to_numpy()
    df = df.with_columns((pl.col("log_price") + pl.Series(L_new[u, t] - old[u, t])).alias("log_price")).with_columns(pl.col("log_price").exp().alias("price")).drop(["_u", "_t"])
    w = sim.stock / sim.stock.sum(); truth_reg = np.vstack([(w[sim.mun_region == r][:, None] * L_new[sim.mun_region == r]).sum(0) / w[sim.mun_region == r].sum() for r in range(n_regions)])
    truth_prov = (w[:, None] * L_new).sum(0)
    return df, L_new, truth_reg, truth_prov, sim


def build_sim_panel(kind: str, rep: int) -> pl.DataFrame:
    df, L, truth_reg, truth_prov, sim = one_market(kind, rep)
    f = build_features(df); months = sim.months
    truth = {("municipality", c): L[i] for i, c in enumerate(sim.mun_codes)}; truth.update({("region", f"{r:02d}"): truth_reg[r] for r in range(truth_reg.shape[0])}); truth[("province", "QC")] = truth_prov
    rows = []
    origins = [months[i] for i in range(24, T - 1)]
    for o in origins:
        oi = months.index(o); avail = f.filter(pl.col("month") <= o)
        r = run_sharp(avail, months[: oi + 1], realtime_parent=True)
        cells = [("province", "QC", r.stage1.mu)] + [(lvl, gid, cell_log_index(r, lvl, gid, "filtered")[0]) for (lvl, gid) in r.cells if lvl in ("region", "municipality")]
        for lvl, gid, path in cells:
            if len(path) < LAGS + 1 or (lvl, gid) not in truth:
                continue
            H = path[-(LAGS + 1):]; n = np.array([avail.filter((pl.col("month") == m) & ((pl.col("mun_code") == gid) if lvl == "municipality" else (pl.col("region_code") == gid) if lvl == "region" else True)).height for m in months[oi - LAGS: oi + 1]], float)
            se = np.zeros(LAGS + 1) + 0.01
            feats = {"rt_level": H[-1], "d1": H[-1] - H[-2], "d3": H[-1] - H[-4], "d6": H[-1] - H[-7], "d12": H[-1] - H[0], "vol6": np.std(np.diff(H[-7:])), "vol12": np.std(np.diff(H)),
                     "n_last": n[-1], "n_mean12": n.mean(), "n_ratio": (n[-3:].mean() + 1) / (n.mean() + 1), "se_last": 0.01, "accel": (H[-1] - H[-4]) - (H[-4] - H[-7])}
            tr = truth[(lvl, gid)]
            for h in HORIZONS:
                if oi + h >= T:
                    continue
                rows.append({"class": "house", "level": lvl, "geo_id": gid, "origin": o, "h": h, "month_of_year": int(o[5:7]), "origin_t": oi, "hist": list(H), **feats,
                             "target": float(tr[oi + h] - tr[oi]), "y_final_o": float(tr[oi]), "region_code": gid if lvl == "region" else (f"{sim.mun_region[sim.mun_codes.index(gid)]:02d}" if lvl == "municipality" else None),
                             "cx": float(df.filter(pl.col("mun_code") == gid)["x_m"].mean()) if lvl == "municipality" else 0.0, "cy": float(df.filter(pl.col("mun_code") == gid)["y_m"].mean()) if lvl == "municipality" else 0.0,
                             "grade_target": "A", "n_target": 0, "scenario": kind, "rep": rep})
    p = pl.DataFrame(rows, infer_schema_length=None)
    reg = p.filter(pl.col("level") == "region").select(["origin", "h", pl.col("geo_id").alias("region_code"), pl.col("d3").alias("reg_d3"), pl.col("d12").alias("reg_d12"), pl.col("vol12").alias("reg_vol12")]).unique()
    prov = p.filter(pl.col("level") == "province").select(["origin", "h", pl.col("d3").alias("prov_d3"), pl.col("d12").alias("prov_d12"), pl.col("vol12").alias("prov_vol12")]).unique()
    return p.join(reg, on=["origin", "h", "region_code"], how="left").join(prov, on=["origin", "h"], how="left")


def main() -> None:
    ap = argparse.ArgumentParser(); ap.add_argument("--reps", type=int, default=20); ap.add_argument("--family", type=str, default="classic"); ap.add_argument("--build", action="store_true"); ap.add_argument("--scenarios", type=str, default="rw,momentum,regimes,spillover,sparse"); ap.add_argument("--horizons", type=str, default="1,3,6,12")
    a = ap.parse_args()
    if a.build:
        for kind in a.scenarios.split(","):
            for rep in range(a.reps):
                t0 = time.time(); p = build_sim_panel(kind, rep); p.write_parquet(OUT / f"panel_{kind}_{rep}.parquet"); print("built", kind, rep, p.shape, f"{time.time() - t0:.0f}s", flush=True)
        return
    if a.family == "torch":
        from wp23.models_b_torch import registry_torch_fast as registry
    else:
        from wp23.models_b import registry
    rows = []
    for kind in a.scenarios.split(","):
        for rep in range(a.reps):
            f = OUT / f"panel_{kind}_{rep}.parquet"
            if not f.exists():
                continue
            p = pl.read_parquet(f)
            for h in [int(x) for x in a.horizons.split(",")]:
                ph = p.filter(pl.col("h") == h); origins = sorted(ph["origin"].unique().to_list())
                for k in range(0, len(origins), 6):
                    o = origins[k]; test_o = origins[k:k + 6]
                    tr = ph.filter(pl.col("origin").map_elements(lambda x: month_add(x, h), return_dtype=pl.Utf8) <= o)
                    if tr["origin"].n_unique() < 8:
                        continue
                    te = ph.filter(pl.col("origin").is_in(test_o))
                    for m in registry(a.family, fast=True):
                        try:
                            m.fit(tr); pr = m.predict(te); Q = m.predict_quantiles(te)
                        except Exception as ex:
                            print("FAILED", kind, rep, h, m.name, repr(ex), flush=True); continue
                        y = te["target"].to_numpy()
                        rows.append(te.select(["level", "geo_id", "origin", "h", "target"]).with_columns(pl.lit(kind).alias("scenario"), pl.lit(rep).alias("rep"), pl.lit(m.name).alias("model"), pl.Series("pred", np.asarray(pr, dtype=np.float64)), pl.Series("q05", np.asarray(Q[:, 0], dtype=np.float64)), pl.Series("q95", np.asarray(Q[:, 4], dtype=np.float64))))
            print(kind, rep, "done", flush=True)
            pl.concat(rows).write_parquet(OUT / f"predictions_{a.family}.parquet")
    print("done")


if __name__ == "__main__":
    main()