# -*- coding: utf-8 -*-
"""03 — Modèles d'ajustement et modèles de référence, ré-estimés chaque trimestre avec les seules ventes
antérieures au trimestre (as-of 2024-01, 2024-04, …, 2026-07) — comme le ferait un AVM ré-entraîné régulièrement :
  • grille hédonique globale (coefficients) et régionale rétrécie   → ajustements B / C
  • LightGBM (objectif L1 = médiane) sur le prix ajusté dans le temps (avec et sans rôle) → ajustements C-GBM, M11
  • hédonique globale directe avec effets de municipalité           → M10
Sorties : models_<période>.json, gbm_pred_<période>.npy (prédiction f pour TOUTES les ventes)."""
import pickle
from common import *
import lightgbm as lgb

df = load_parquet("sales.parquet")
meta = load_json("index_meta.json"); regions = meta["regions"]; M_EVAL0 = meta["M_EVAL0"]
rid = {r: i for i, r in enumerate(regions)}
df["rid"] = df["region"].map(rid).astype(int)
idx = np.load(OUT / "index_asof.npy")

FEATS = ["ln_floor", "ln_land", "age", "age2", "lien1", "lien2", "genre1", "genre2", "genre3", "etages", "neuf"]
def feats(d):
    return pd.DataFrame({
        "ln_floor": d["ln_floor"], "ln_land": d["ln_land"], "age": d["age"], "age2": d["age"] ** 2 / 100,
        "lien1": (d["lien"] == 1).astype(float), "lien2": (d["lien"] >= 2).astype(float),
        "genre1": (d["genre"] == 1).astype(float), "genre2": (d["genre"] == 2).astype(float), "genre3": (d["genre"] == 3).astype(float),
        "etages": d["etages"].astype(float), "neuf": d["neuf"].astype(float)}, index=d.index)

def time_adjusted(d, m_asof):
    """ln P ramené au niveau du mois m_asof-1 avec l'indice connu au mois m_asof."""
    lvl = idx[m_asof - M_EVAL0]
    return d["ln_price"].to_numpy() - lvl[d["rid"].to_numpy(), d["month"].to_numpy()] + lvl[d["rid"].to_numpy(), m_asof - 1]

def ols(X, y):
    A = np.column_stack([np.ones(len(X)), X])
    b, *_ = np.linalg.lstsq(A, y, rcond=None)
    return b

ASOF = list(range(M_EVAL0, int(df["month"].max()) + 1, 3))   # ré-estimation trimestrielle (2024-01, 2024-04, …)
save_json({"asof": ASOF}, "asof.json")
TODO = [int(a) for a in os.environ["CV_ASOF"].split(",")] if os.environ.get("CV_ASOF") else ASOF   # sous-ensemble (répartition sur plusieurs nœuds)
for m_asof in TODO:
    period = str(m_asof); start = f"{2021 + m_asof // 12}-{m_asof % 12 + 1:02d}-01"
    tr = df[df["month"] < m_asof].copy()
    y = time_adjusted(tr, m_asof)
    F = feats(tr)
    log(f"[{period}] entraînement sur {len(tr):,} ventes < {start} (indice connu au mois {m_asof})")
    # --- grille hédonique globale : estimateur « within » (effets fixes de municipalité), pour que les β
    #     mesurent des différences entre biens voisins et non entre marchés
    def within(frame, cols, by):
        return frame[cols] - frame.groupby(by)[cols].transform("mean")
    # effets mensuels résiduels (base = dernier mois connu m_asof-1) : l'indice régional ne retire pas toute la
    # tendance commune ; sans eux, les modèles globaux sous-estiment les ventes récentes
    months_tr = sorted(tr["month"].unique()); MCOLS = [f"m{m}" for m in months_tr if m != m_asof - 1]
    Fm = pd.concat([F] + [pd.Series((tr["month"] == m).astype(float), index=tr.index, name=f"m{m}") for m in months_tr if m != m_asof - 1], axis=1)
    Fw = within(pd.concat([Fm, tr[["code_mun"]]], axis=1), FEATS + MCOLS, "code_mun")
    yw = y - pd.Series(y, index=tr.index).groupby(tr["code_mun"]).transform("mean").to_numpy()
    bw_all, *_ = np.linalg.lstsq(Fw.to_numpy(), yw, rcond=None)
    bw = bw_all[:len(FEATS)]; bm = bw_all[len(FEATS):]
    beta_global = dict(zip(FEATS, bw.tolist()))
    xb = F.to_numpy() @ bw + Fm[MCOLS].to_numpy() @ bm
    resid = y - xb
    r2w = 1 - np.var(yw - Fw.to_numpy() @ bw_all) / np.var(yw)
    Fw = Fw[FEATS]
    log(f"  hédonique within-municipalité : R² within = {r2w:.3f} ; β = { {k: round(v, 3) for k, v in beta_global.items()} }")
    # --- grille régionale (within-municipalité) rétrécie vers la globale (n/(n+300))
    beta_region = {}
    for r in regions:
        sel = (tr["region"] == r).to_numpy()
        if sel.sum() < 100: beta_region[r] = beta_global; continue
        b, *_ = np.linalg.lstsq(Fw.to_numpy()[sel], yw[sel], rcond=None)
        w = sel.sum() / (sel.sum() + 300.0)
        beta_region[r] = {k: float(w * b[j] + (1 - w) * beta_global[k]) for j, k in enumerate(FEATS)}
    # --- M10 : hédonique directe = Xβ + effet de région + MRC + municipalité + unité de voisinage (rétrécis en cascade)
    reg_eff = pd.Series(resid, index=tr.index).groupby(tr["region"]).mean().to_dict()
    tr["res_g"] = resid - tr["region"].map(reg_eff).to_numpy()
    mrc_eff = tr.groupby("mrc")["res_g"].agg(["sum", "count"])
    mrc_eff = (mrc_eff["sum"] / (mrc_eff["count"] + 30)).to_dict()
    tr["res_m"] = tr["res_g"] - tr["mrc"].map(mrc_eff).fillna(0)
    mun_eff = tr.groupby("code_mun")["res_m"].agg(["sum", "count"])
    mun_eff = (mun_eff["sum"] / (mun_eff["count"] + 20)).to_dict()
    tr["res_u"] = tr["res_m"] - tr["code_mun"].map(mun_eff).fillna(0)
    # correction médiane : la MCO vise la moyenne du log ; la distribution est asymétrique (queue basse), on recentre sur la médiane
    uv_tmp = tr.groupby(["code_mun", "unite_voisinage"])["res_u"].transform(lambda v: v.sum() / (len(v) + 10))
    median_shift = float(np.nanmedian(tr["res_u"] - uv_tmp.fillna(0)))
    uv = tr.groupby(["code_mun", "unite_voisinage"])["res_u"].agg(["sum", "count"])
    uv_eff = {f"{a}|{b}": v for (a, b), v in (uv["sum"] / (uv["count"] + 10)).items()}
    # --- LightGBM (prix ajusté dans le temps) avec rôle et sans rôle
    def gbm_frame(d):
        G = feats(d)
        G["lat"] = d["lat"]; G["lng"] = d["lng"]; G["x_km"] = d["x_km"]; G["y_km"] = d["y_km"]
        G["rid"] = d["rid"]; G["btype"] = d["btype"]
        G["ln_role"] = np.log(d["role_vente"]); G["part_terrain"] = d["part_terrain"]; G["role_vintage"] = d["role_vintage"]
        G["month"] = np.minimum(d["month"], m_asof - 1).astype(float)   # tendance résiduelle ; hors échantillon = dernier mois connu
        return G
    G = gbm_frame(tr)
    params = dict(objective="regression_l1", learning_rate=0.05, num_leaves=127, min_data_in_leaf=40, feature_fraction=0.8,
                  bagging_fraction=0.8, bagging_freq=1, lambda_l2=1.0, verbose=-1, num_threads=int(os.environ.get("CV_THREADS", 20)))
    models = {}
    for name, cols in [("gbm", [c for c in G.columns if c not in ("ln_role", "part_terrain", "role_vintage")]),
                       ("gbm_role", list(G.columns))]:
        ds = lgb.Dataset(G[cols], y, categorical_feature=["rid", "btype"] if "btype" in cols else ["rid"])
        # validation interne : derniers 10 % chronologiques pour fixer le nombre d'arbres
        n = len(G); cut = int(n * 0.9)
        dtr = lgb.Dataset(G[cols].iloc[:cut], y[:cut], categorical_feature=["rid", "btype"])
        dva = lgb.Dataset(G[cols].iloc[cut:], y[cut:], reference=dtr)
        bst = lgb.train(params, dtr, num_boost_round=3000, valid_sets=[dva], callbacks=[lgb.early_stopping(100, verbose=False)])
        best = max(bst.best_iteration, 300)   # plancher : l'arrêt précoce sur la tranche finale (mois partiels) peut être trop court
        bst = lgb.train(params, ds, num_boost_round=best)
        models[name] = bst
        Gall = gbm_frame(df); Gall["month"] = float(m_asof - 1)   # f évalué au niveau du dernier mois connu pour toutes les ventes
        pred_all = bst.predict(Gall[cols])
        np.save(OUT / f"{name}_pred_{period}.npy", pred_all.astype(np.float32))
        rr = y - bst.predict(G[cols]); r2 = 1 - np.var(rr) / np.var(y)
        log(f"  {name}: {best} arbres, R² (in-sample) = {r2:.3f}, résidu médian = {np.median(rr) * 100:.2f} %")
        bst.save_model(str(OUT / f"{name}_{period}.txt"))
    save_json({"period": period, "m_asof": m_asof, "start": start, "n_train": len(tr), "feats": FEATS,
               "beta_global": beta_global, "intercept_global": 0.0, "region_fe": reg_eff,
               "beta_region": beta_region, "mrc_eff": mrc_eff, "mun_eff": mun_eff, "uv_eff": uv_eff,
               "r2_within": float(r2w), "sd_resid_u": float(tr["res_u"].std()), "median_shift": median_shift},
              f"models_{period}.json")
log("modèles terminés")