# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Synthetic Quebec-like housing markets with known quality-constant latent price paths (protocol §12).

A market has R regions × M MRCs × municipalities of heterogeneous size (expected monthly sales from a log-normal
mimicking results/profile/municipality_size_bands.csv). True log index of municipality u at month t:
    L_{u,t} = μ_t + δ^R_{r(u),t} + δ^M_{m(u),t} + δ^U_{u,t}
with μ a random walk with drift and optional structural break / turning point, and nested random-walk deviations
whose innovation sd differs across municipalities. Prices:
    log P_i = L_{u(i),t(i)} + f_t(X_i) + α_cell(i) + ε_i,  ε ~ Student-t(4) scaled (+ contamination)
f is nonlinear in floor area with an optional drifting elasticity; α_cell is spatially correlated.
Composition shifts: the distribution of sold X and of sold cells can trend over time (quality-up / quality-down,
location drift). Sales counts are Poisson with June–July seasonality. Floor area missing MCAR or time-increasing.
Repeat sales: a fraction of parcels resell. Registration lags: exponential days (for real-time experiments).
The output has the same columns as the real analytic layer so every estimator runs unchanged.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
import polars as pl


@dataclass
class Scenario:
    name: str = "baseline"
    T: int = 60
    n_regions: int = 4
    mrc_per_region: int = 3
    mun_per_mrc: int = 5
    size_logmean: float = 1.2      # log expected monthly sales per municipality (median ≈ 3.3)
    size_logsd: float = 1.3
    drift: float = 0.004           # provincial monthly drift
    tau_prov: float = 0.012
    tau_region: float = 0.006
    tau_mrc: float = 0.004
    tau_mun: float = 0.006
    tau_mun_hetero: bool = True    # half the municipalities have tau_mun/4
    break_month: int | None = None # turning point: drift sign flips
    break_drift: float = -0.008
    sigma_eps: float = 0.16
    t_df: float = 4.0
    contamination: float = 0.01
    beta_fa: float = 0.55
    beta_fa_drift: float = 0.0     # change in floor-area elasticity over the sample
    quality_trend: float = 0.0     # monthly drift of mean log floor area of SOLD homes (composition)
    location_trend: float = 0.0    # monthly drift of mean cell quality of SOLD homes
    fa_missing: float = 0.03
    fa_missing_trend: float = 0.0  # additional missingness at the end of the sample
    resale_prob: float = 0.12
    seasonal_amp: float = 0.6
    reg_lag_mean_days: float = 10.0
    seed: int = 0
    scale_sales: float = 1.0       # multiply all expected sales (sample-size experiments)


@dataclass
class SimOutput:
    df: pl.DataFrame
    truth_mun: np.ndarray      # (n_mun, T) true log index
    truth_region: np.ndarray   # (n_regions, T) stock-weighted
    truth_mrc: np.ndarray
    truth_prov: np.ndarray     # (T,)
    mun_codes: list[str]
    mun_region: np.ndarray
    mun_mrc: np.ndarray
    stock: np.ndarray          # housing stock weights per municipality
    months: list[str]
    scenario: Scenario


def _months(T: int, start: str = "2021-01") -> list[str]:
    y, m = int(start[:4]), int(start[5:7]); out = []
    for _ in range(T):
        out.append(f"{y:04d}-{m:02d}"); m += 1
        if m == 13: y, m = y + 1, 1
    return out


def simulate(sc: Scenario) -> SimOutput:
    rng = np.random.default_rng(sc.seed)
    T = sc.T; months = _months(T)
    n_mrc = sc.n_regions * sc.mrc_per_region; n_mun = n_mrc * sc.mun_per_mrc
    mun_mrc = np.repeat(np.arange(n_mrc), sc.mun_per_mrc); mun_region = mun_mrc // sc.mrc_per_region
    lam = np.exp(rng.normal(sc.size_logmean, sc.size_logsd, n_mun)) * sc.scale_sales
    stock = lam * rng.uniform(80, 160, n_mun)  # stock ∝ turnover with noise
    # latent paths
    drift = np.full(T, sc.drift)
    if sc.break_month is not None:
        drift[sc.break_month:] = sc.break_drift
    mu = np.cumsum(drift + rng.normal(0, sc.tau_prov, T)); mu -= mu[0]
    dR = np.cumsum(rng.normal(0, sc.tau_region, (sc.n_regions, T)), axis=1); dR -= dR[:, [0]]
    dM = np.cumsum(rng.normal(0, sc.tau_mrc, (n_mrc, T)), axis=1); dM -= dM[:, [0]]
    tau_u = np.full(n_mun, sc.tau_mun)
    if sc.tau_mun_hetero:
        tau_u[rng.random(n_mun) < 0.5] = sc.tau_mun / 4
    dU = np.cumsum(rng.normal(0, 1, (n_mun, T)) * tau_u[:, None], axis=1); dU -= dU[:, [0]]
    L = mu[None, :] + dR[mun_region] + dM[mun_mrc] + dU
    # geography: municipality centres clustered by region, cells inside
    reg_c = rng.uniform(0, 400_000, (sc.n_regions, 2))
    mrc_c = reg_c[np.arange(n_mrc) // sc.mrc_per_region] + rng.normal(0, 40_000, (n_mrc, 2))
    mun_c = mrc_c[mun_mrc] + rng.normal(0, 12_000, (n_mun, 2))
    mun_level = rng.normal(0, 0.25, n_mun)  # price level of the municipality (absorbed by location FE)
    # sales
    season = 1 + sc.seasonal_amp * np.cos(2 * np.pi * (np.arange(T) % 12 - 5) / 12)
    rows = []
    parcel_counter = 0
    for u in range(n_mun):
        for t in range(T):
            n = rng.poisson(lam[u] * season[t])
            if n == 0:
                continue
            # composition drift of sold homes
            lfa = rng.normal(np.log(120) + 0.15 * mun_level[u] + sc.quality_trend * t, 0.35, n)
            age = np.clip(rng.gamma(2.0, 20, n), 0, 150)
            lla = rng.normal(np.log(600), 0.6, n)
            storeys = rng.choice([1, 2], n, p=[0.6, 0.4])
            # within-municipality location: cell offset with spatially correlated quality; location drift = mean cell quality trends
            theta = rng.uniform(0, 2 * np.pi, n); rad = rng.gamma(2, 1200, n)
            x = mun_c[u, 0] + rad * np.cos(theta); y = mun_c[u, 1] + rad * np.sin(theta)
            cell_q = 0.08 * np.sin(x / 1500) * np.cos(y / 1500) + 0.06 * (rad < 1500) + sc.location_trend * t * rng.uniform(0, 1, n)
            beta_fa_t = sc.beta_fa + sc.beta_fa_drift * t / max(T - 1, 1)
            fX = beta_fa_t * (lfa - np.log(120)) + 0.08 * (lfa - np.log(120)) ** 2 - 0.004 * age + 0.00002 * age ** 2 + 0.09 * (lla - np.log(600)) + 0.05 * (storeys - 1)
            eps = rng.standard_t(sc.t_df, n) * sc.sigma_eps / np.sqrt(sc.t_df / (sc.t_df - 2))
            cont = rng.random(n) < sc.contamination
            eps[cont] += rng.choice([-1.0, 1.0], cont.sum()) * rng.uniform(0.6, 1.2, cont.sum())
            lp = 12.6 + mun_level[u] + L[u, t] + fX + cell_q + eps
            pid = parcel_counter + np.arange(n); parcel_counter += n
            rows.append(pl.DataFrame({
                "mun_i": u, "t": t, "log_price": lp, "lfa_true": lfa, "age": age, "lla_true": lla, "storeys": storeys, "x_m": x, "y_m": y, "parcel_i": pid,
                "year_built_offset": age,
            }))
    df = pl.concat(rows)
    # repeat sales: a share of sales are resales of an earlier parcel (same X, same cell, new time)
    n_rows = df.height
    is_resale = rng.random(n_rows) < sc.resale_prob
    df = df.with_row_index("_i")
    # pick an earlier parcel in the same municipality for resales
    first = df.filter(~pl.Series(is_resale)).select(["_i", "mun_i", "t", "parcel_i", "lfa_true", "age", "lla_true", "storeys", "x_m", "y_m"])
    res = df.filter(pl.Series(is_resale)).select(["_i", "mun_i", "t", "log_price"])
    cand = first.join(res, on="mun_i", how="inner", suffix="_r").filter(pl.col("t_r") - pl.col("t") >= 6)
    if cand.height:  # one random earlier parcel per resale (not always the earliest)
        cand = cand.with_columns(pl.Series("_rk", rng.random(cand.height))).sort("_rk").group_by("_i_r").agg(pl.all().first()).drop("_rk")
    if cand.height:
        # replace resale rows' parcel/X by the matched earlier sale (price keeps the new period's latent level)
        rep = cand.select(pl.col("_i_r").alias("_i"), pl.col("parcel_i").alias("parcel_new"), pl.col("lfa_true").alias("lfa_new"), pl.col("age").alias("age_new"),
                          pl.col("lla_true").alias("lla_new"), pl.col("storeys").alias("st_new"), pl.col("x_m").alias("x_new"), pl.col("y_m").alias("y_new"))
        df = df.join(rep, on="_i", how="left")
        # price adjustment so the resale keeps its own latent level: recompute hedonic part difference approximately via lfa change
        beta = sc.beta_fa
        df = df.with_columns(
            pl.when(pl.col("parcel_new").is_not_null()).then(pl.col("log_price") + beta * (pl.col("lfa_new") - pl.col("lfa_true"))).otherwise(pl.col("log_price")).alias("log_price"),
            pl.coalesce(pl.col("parcel_new"), pl.col("parcel_i")).alias("parcel_i"),
            pl.coalesce(pl.col("lfa_new"), pl.col("lfa_true")).alias("lfa_true"), pl.coalesce(pl.col("age_new"), pl.col("age")).alias("age"),
            pl.coalesce(pl.col("lla_new"), pl.col("lla_true")).alias("lla_true"), pl.coalesce(pl.col("st_new"), pl.col("storeys")).alias("storeys"),
            pl.coalesce(pl.col("x_new"), pl.col("x_m")).alias("x_m"), pl.coalesce(pl.col("y_new"), pl.col("y_m")).alias("y_m"),
        ).drop(["parcel_new", "lfa_new", "age_new", "lla_new", "st_new", "x_new", "y_new"])
    # missing floor area (MCAR + trend)
    p_miss = sc.fa_missing + sc.fa_missing_trend * (df["t"].to_numpy() / max(T - 1, 1))
    miss = rng.random(df.height) < p_miss
    mun_codes = [f"M{u:03d}" for u in range(n_mun)]
    mrc_codes = [f"R{m:02d}" for m in range(n_mrc)]
    reg_codes = [f"{r:02d}" for r in range(sc.n_regions)]
    lag_days = rng.exponential(sc.reg_lag_mean_days, df.height)
    df = df.with_columns(
        pl.Series("month", [months[t] for t in df["t"].to_list()]),
        pl.Series("mun_code", [mun_codes[u] for u in df["mun_i"].to_list()]),
        pl.Series("mrc_code", [mrc_codes[mun_mrc[u]] for u in df["mun_i"].to_list()]),
        pl.Series("region_code", [reg_codes[mun_region[u]] for u in df["mun_i"].to_list()]),
        pl.Series("floor_area_m2", np.where(miss, np.nan, np.exp(df["lfa_true"].to_numpy()))),
        pl.col("lla_true").exp().alias("land_area_m2"),
        pl.Series("reg_lag_days", lag_days),
    ).with_columns(
        pl.col("month").str.to_date("%Y-%m-%d", strict=False).alias("_d"),
    )
    # dates: uniform day within month
    day = rng.integers(1, 28, df.height)
    df = df.with_columns(pl.Series("_day", day)).with_columns(
        (pl.col("month") + "-" + pl.col("_day").cast(pl.Utf8).str.zfill(2)).str.to_date("%Y-%m-%d").alias("date"),
    ).with_columns(
        pl.col("date").dt.year().alias("year"), pl.col("log_price").exp().alias("price"),
        pl.lit("house").alias("property_class"), pl.lit("detached").alias("subtype"), pl.lit("De plain-pied").alias("genre"),
        pl.lit(1).alias("n_units"), pl.col("storeys").cast(pl.Float64).alias("storeys"),
        pl.col("parcel_i").cast(pl.Utf8).alias("parcel_id"), pl.col("_i").cast(pl.Utf8).alias("tx_id"),
        pl.when(pl.Series(miss)).then(None).otherwise(pl.col("floor_area_m2")).alias("floor_area_m2"),
        (pl.col("date") + pl.duration(days=pl.col("reg_lag_days").cast(pl.Int64))).alias("available_date"),
    )
    for lvl, size in (("grid500", 500), ("grid1k", 1000), ("grid2k", 2000), ("grid5k", 5000)):
        df = df.with_columns((pl.col("x_m") // size).cast(pl.Int64).cast(pl.Utf8).add("_").add((pl.col("y_m") // size).cast(pl.Int64).cast(pl.Utf8)).alias(lvl))
    df = df.drop(["_i", "_d", "_day", "mun_i", "t", "parcel_i", "lfa_true", "lla_true", "year_built_offset"])
    # truths: stock-weighted aggregates
    w = stock / stock.sum()
    truth_prov = (w[:, None] * L).sum(axis=0)
    truth_region = np.vstack([(w[mun_region == r][:, None] * L[mun_region == r]).sum(axis=0) / w[mun_region == r].sum() for r in range(sc.n_regions)])
    truth_mrc = np.vstack([(w[mun_mrc == m][:, None] * L[mun_mrc == m]).sum(axis=0) / w[mun_mrc == m].sum() for m in range(n_mrc)])
    return SimOutput(df, L, truth_region, truth_mrc, truth_prov, mun_codes, mun_region, mun_mrc, stock, months, sc)


SCENARIOS: dict[str, Scenario] = {
    "baseline": Scenario(name="baseline"),
    "turning_point": Scenario(name="turning_point", break_month=36, drift=0.006, break_drift=-0.008),
    "flat_quality_up": Scenario(name="flat_quality_up", drift=0.0, tau_prov=0.006, quality_trend=0.004, location_trend=0.0008),
    "rising_quality_down": Scenario(name="rising_quality_down", drift=0.008, quality_trend=-0.004, location_trend=-0.0008),
    "drifting_hedonics": Scenario(name="drifting_hedonics", beta_fa_drift=0.15, fa_missing_trend=0.08),
    "heavy_tails": Scenario(name="heavy_tails", t_df=2.5, contamination=0.04, sigma_eps=0.2),
    "sparse": Scenario(name="sparse", scale_sales=0.35),
    "no_local_dynamics": Scenario(name="no_local_dynamics", tau_mun=0.0005, tau_mrc=0.0005, tau_region=0.001),  # counter-example: pooling should win big
    "idiosyncratic_towns": Scenario(name="idiosyncratic_towns", tau_mun=0.02, tau_mun_hetero=False, tau_region=0.002),  # counter-example: pooling hurts
}