#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Ablation CV complète (4 couches x OLS/GBM) pour cohérence du tableau ablation.tex."""
import json, warnings
import numpy as np, pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.model_selection import KFold, cross_val_predict
from sklearn.metrics import r2_score
warnings.filterwarnings("ignore")

exec(open("scripts/09_v5_ml_text_vision.py").read().split("kf=KFold")[0])
kf = KFold(n_splits=5, shuffle=True, random_state=42)
amen_ = list(AMEN.keys())

L1 = ["sqft_final","has_sqft","lot_sqft","beds","baths","parking","storeys","age","age2"] + amen_
L2 = L1 + ["n_images_c","virtual_tour"] + rv
L3 = L2 + gv
L4 = L3 + vv
print("feats:", len(L1), len(L2), len(L3), len(L4))

Xcat = pd.get_dummies(df[cats], prefix=cats, dtype=float)
dfi = df.copy()
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)

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_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))

out = {}
for name, cols in [("L1", L1), ("L2", L2), ("L3", L3), ("L4", L4)]:
    out[name] = {"nfeat": len(cols), "ols": cv_ols(cols), "gbm": cv_gbm(cols)}
    print(name, out[name])
json.dump(out, open("results/ablation_extras.json", "w"), indent=1)
print("-> ablation_extras.json")