# -*- coding: utf-8 -*-
"""02 — Indice temporel régional « point dans le temps ».
Pour chaque mois d'évaluation m, on ré-estime un indice hédonique mensuel par région administrative
avec toutes les ventes antérieures au mois m (fenêtre expansive) (jamais le mois m lui-même). Niveau de référence = m-1,
dernier mois entièrement connu. Sortie : index_asof.npy [m_eval, région, mois] (log), regions.json,
et un indice « plein échantillon » descriptif (index_full.csv)."""
from common import *

df = load_parquet("sales.parquet")
regions = sorted(df["region"].unique().tolist())
rid = {r: i for i, r in enumerate(regions)}
df["rid"] = df["region"].map(rid).astype(int)
M = int(df["month"].max()) + 1
M_EVAL0 = int(month_index(pd.Series([VAL_START]))[0])
WINDOW = 999  # fenêtre « expansive » : toutes les ventes antérieures au mois m
SHRINK = 40.0  # n/(n+SHRINK) : rétrécissement de l'effet mensuel régional vers l'effet provincial

X_cols = ["ln_floor", "ln_land", "age"]
def design(d):
    X = np.column_stack([np.ones(len(d)), d["ln_floor"], d["ln_land"], d["age"], d["age"] ** 2 / 100,
                         (d["lien"] == 1), (d["lien"] >= 2), (d["genre"] == 1), (d["genre"] == 2), (d["genre"] == 3),
                         d["etages"], d["neuf"]]).astype(float)
    return X

def fit_month_effects(d, months):
    """OLS à effets fixes de municipalité (estimateur within) : ln P ~ X + effets mensuels (base = dernier mois),
    toutes les variables étant centrées par municipalité pour que l'indice ne soit pas contaminé par les
    déplacements du mélange des ventes entre municipalités chères et bon marché. Renvoie effets et n par mois."""
    X = design(d)[:, 1:]
    mm = d["month"].to_numpy()
    base = months[-1]
    dums = np.column_stack([(mm == m).astype(float) for m in months if m != base])
    A = np.hstack([X, dums]); y = d["ln_price"].to_numpy(float)
    codes, inv = np.unique(d["code_mun"].to_numpy(), return_inverse=True)
    cnt = np.bincount(inv).astype(float)
    mA = np.zeros((len(codes), A.shape[1])); np.add.at(mA, inv, A); A = A - mA[inv] / cnt[inv, None]
    my = np.bincount(inv, y) / cnt; y = y - my[inv]
    beta, *_ = np.linalg.lstsq(A, y, rcond=None)
    eff = {base: 0.0}
    k = X.shape[1]
    for j, m in enumerate([m for m in months if m != base]):
        eff[m] = float(beta[k + j])
    n = {m: int((mm == m).sum()) for m in months}
    return eff, n

idx = np.full((M - M_EVAL0, len(regions), M), np.nan, dtype=np.float32)
log(f"indice point-dans-le-temps : mois d'évaluation {M_EVAL0}..{M-1}, {len(regions)} régions")
for m in range(M_EVAL0, M):
    months = list(range(max(0, m - WINDOW), m))
    win = df[(df["month"] >= months[0]) & (df["month"] < m)]
    prov_eff, _ = fit_month_effects(win, months)
    for r in regions:
        dr = win[win["region"] == r]
        if len(dr) < 200:
            for mo in months: idx[m - M_EVAL0, rid[r], mo] = prov_eff[mo]
            continue
        eff, n = fit_month_effects(dr, months)
        for mo in months:
            w = n[mo] / (n[mo] + SHRINK)
            idx[m - M_EVAL0, rid[r], mo] = w * eff[mo] + (1 - w) * prov_eff[mo]
    if m % 6 == 0: log(f"  mois {m} ({len(win):,} ventes dans la fenêtre)")
np.save(OUT / "index_asof.npy", idx)
save_json({"regions": regions, "M_EVAL0": M_EVAL0, "M": M, "window": WINDOW, "shrink": SHRINK}, "index_meta.json")

# indice descriptif plein échantillon (pour le rapport seulement — pas utilisé dans la rétro-évaluation)
rows = []
months = list(range(0, M))
eff_p, n_p = fit_month_effects(df, months)
for mo in months:
    rows.append({"month": mo, "ym": f"{2021 + mo // 12}-{mo % 12 + 1:02d}", "region": "QC", "log_idx": eff_p[mo] - eff_p[months[0]], "n": n_p[mo]})
for r in regions:
    dr = df[df["region"] == r]
    if len(dr) < 2000: continue
    eff, n = fit_month_effects(dr, months)
    for mo in months:
        rows.append({"month": mo, "ym": f"{2021 + mo // 12}-{mo % 12 + 1:02d}", "region": r, "log_idx": eff[mo] - eff[months[0]], "n": n[mo]})
pd.DataFrame(rows).to_csv(OUT / "index_full.csv", index=False)
log("indice terminé")