#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Estimations complémentaires pour la révision JHE :
   - OLS A/B avec SE HC3 ET clusterisées par municipalité
   - ablation in-sample du modèle A (base / +texte / +texte+vision)
   - effets fixes de ville extrêmes
   - CV 5-fold OLS et RF pour les panels A (texte) et B (texte+vision)
   Résultats -> revision_extras.json
"""
import json, warnings
import numpy as np, pandas as pd
import statsmodels.formula.api as smf
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor, HistGradientBoostingRegressor
from sklearn.model_selection import KFold, cross_val_predict
from sklearn.metrics import r2_score
warnings.filterwarnings("ignore")

# pipeline v5 jusqu'avant la CV (construit df, AMEN, rv, gv, vv, NUM_v4, NUM_v5, MISS, cats, y, kf)
exec(open("scripts/09_v5_ml_text_vision.py").read().split("kf=KFold")[0])
amen_ = list(AMEN.keys())
out = {}

# ---------- OLS A/B : HC3 vs cluster ----------
df["ln_price"] = np.log(df["price"])
df["ln_sqft"] = np.log(df["sqft_final"].fillna(df.groupby("ptype")["sqft_final"].transform("median")).fillna(df["sqft_final"].median()))
df["ln_lot"] = np.log(df["lot_sqft"].fillna(df["lot_sqft"].median())); df["has_lot"] = df["lot_sqft"].notna().astype(int)
df["has_age"] = df["age"].notna().astype(int)
for c in ["beds","baths","storeys","age","age2"]:
    df[c] = df[c].fillna(df.groupby("ptype")[c].transform("median")).fillna(df[c].median())
df["parking"] = df["parking"].fillna(0)

form = ("ln_price ~ ln_sqft+ln_lot+has_lot+beds+baths+parking+storeys+age+age2+has_age+C(source)+C(ptype_fe)+"
        + "+".join(amen_) + "+n_images_c+virtual_tour+" + "+".join(rv) + "+" + "+".join(gv) + "+" + "+".join(vv) + "+C(city_fe)")
dfA = df[df["has_sqft"] == 1]
mA  = smf.ols(form, data=dfA).fit(cov_type="HC3")
mB  = smf.ols(form.replace("ln_sqft+",""), data=df).fit(cov_type="HC3")
mAc = smf.ols(form, data=dfA).fit(cov_type="cluster", cov_kwds={"groups": dfA["city_fe"]})
mBc = smf.ols(form.replace("ln_sqft+",""), data=df).fit(cov_type="cluster", cov_kwds={"groups": df["city_fe"]})

keys = ["ln_sqft","ln_lot","beds","baths","parking","age","age2",
        "central_air","transit","waterfront",
        "g_cond","g_renovated_kitchen","g_hardwood","g_new_construction","g_income_property",
        "v_curb","v_reno","v_luxury","v_garage_vis",
        "v_mat_brique","v_mat_pierre","v_mat_vinyle","v_set_rural","v_land","v_cond"]
def grab(m, mc):
    o = {}
    for k in keys:
        if k in m.params.index:
            o[k] = {"b": float(m.params[k]), "se_hc3": float(m.bse[k]), "p_hc3": float(m.pvalues[k]),
                    "se_cl": float(mc.bse[k]), "p_cl": float(mc.pvalues[k])}
    return o
out["modelA"] = {"r2": mA.rsquared, "adj": mA.rsquared_adj, "n": int(mA.nobs), "coef": grab(mA, mAc)}
out["modelB"] = {"r2": mB.rsquared, "adj": mB.rsquared_adj, "n": int(mB.nobs), "coef": grab(mB, mBc)}
print("A R2=%.4f N=%d | B R2=%.4f N=%d" % (mA.rsquared, mA.nobs, mB.rsquared, mB.nobs))

# ---------- ablation in-sample modèle A ----------
base_form = ("ln_price ~ ln_sqft+ln_lot+has_lot+beds+baths+parking+storeys+age+age2+has_age+C(source)+C(ptype_fe)+"
             + "+".join(amen_) + "+n_images_c+virtual_tour+" + "+".join(rv) + "+C(city_fe)")
mA0 = smf.ols(base_form, data=dfA).fit()
mA1 = smf.ols(base_form + "+" + "+".join(gv), data=dfA).fit()
mA2 = smf.ols(base_form + "+" + "+".join(gv) + "+" + "+".join(vv), data=dfA).fit()
out["ablationA_insample"] = {"base": mA0.rsquared, "plus_text": mA1.rsquared, "plus_text_vision": mA2.rsquared}
print("ablation A in-sample:", out["ablationA_insample"])

# ---------- effets fixes ville extrêmes ----------
fes = {k.split("T.")[1].rstrip("]"): float(v) for k, v in mA.params.items() if k.startswith("C(city_fe)")}
top = sorted(fes.items(), key=lambda kv: -kv[1])[:12]
bot = sorted(fes.items(), key=lambda kv: kv[1])[:8]
out["city_fe_top"] = [(k, (np.exp(v) - 1) * 100) for k, v in top]
out["city_fe_bot"] = [(k, (np.exp(v) - 1) * 100) for k, v in bot]
print("top villes:", out["city_fe_top"][:4])

# ---------- CV : OLS et RF, panels texte / texte+vision ----------
kf = KFold(n_splits=5, shuffle=True, random_state=42)
Xcat = pd.get_dummies(df[cats], prefix=cats, dtype=float)
dfi = df.copy()
def stats(pred):
    e = np.abs(np.exp(pred - y) - 1)
    return {"r2": float(r2_score(y, pred)), "mdae": float(np.median(e) * 100),
            "w15": float(np.mean(e <= .15) * 100), "w25": float(np.mean(e <= .25) * 100)}
def cv_ols(numcols):
    X = pd.concat([dfi[numcols].reset_index(drop=True), Xcat.reset_index(drop=True)], axis=1)
    X = X.fillna(X.median())
    return stats(cross_val_predict(LinearRegression(), X, y, cv=kf, n_jobs=-1))
def cv_rf(numcols):
    X = pd.concat([df[numcols].reset_index(drop=True), Xcat.reset_index(drop=True)], axis=1)
    m = RandomForestRegressor(n_estimators=400, max_features=0.5, min_samples_leaf=3, n_jobs=-1, random_state=42)
    return stats(cross_val_predict(m, X, y, cv=kf, n_jobs=-1))
def cv_gbm(numcols):
    X = pd.concat([df[numcols].reset_index(drop=True), Xcat.reset_index(drop=True)], axis=1)
    m = HistGradientBoostingRegressor(max_iter=600, learning_rate=0.05, max_leaf_nodes=63, l2_regularization=1.0, random_state=42)
    return stats(cross_val_predict(m, X, y, cv=kf, n_jobs=-1))

# matrice OLS imputée médiane par type (comme hedonic_ml)
for c in ["sqft_final","lot_sqft","beds","baths","storeys","age","age2"]:
    dfi[c] = dfi.groupby("ptype")[c].transform(lambda s: s.fillna(s.median()))
    dfi[c] = dfi[c].fillna(dfi[c].median())
dfi["parking"] = dfi["parking"].fillna(0)

print(">>> CV OLS v4/v5")
out["cv"] = {"ols_v4": cv_ols(NUM_v4), "ols_v5": cv_ols(NUM_v5)}
print("OLS:", out["cv"]["ols_v4"], out["cv"]["ols_v5"])
print(">>> CV GBM v4/v5 (contrôle)")
out["cv"]["gbm_v4"] = cv_gbm(NUM_v4); out["cv"]["gbm_v5"] = cv_gbm(NUM_v5)
print("GBM:", out["cv"]["gbm_v4"], out["cv"]["gbm_v5"])
print(">>> CV RF v4/v5 (plus long)")
out["cv"]["rf_v4"] = cv_rf(NUM_v4); out["cv"]["rf_v5"] = cv_rf(NUM_v5)
print("RF:", out["cv"]["rf_v4"], out["cv"]["rf_v5"])

json.dump(out, open("results/revision_extras.json", "w"), indent=1, default=float)
print("-> revision_extras.json")