#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Figures + tableau d'ablation détaillé pour le papier v5."""
import numpy as np, pandas as pd, warnings, json
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
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
from sklearn.inspection import permutation_importance
warnings.filterwarnings("ignore")
plt.rcParams.update({"font.family":"serif","font.size":10,"axes.grid":True,
    "grid.alpha":0.25,"axes.axisbelow":True,"figure.dpi":150,"savefig.bbox":"tight"})
C_OLS,C_GBM,C_BL="#3B6EA5","#C4462E","#4A8C5B"

# build df via v5_ml prefix
exec(open("v5_ml.py").read().split("kf=KFold")[0])
kf=KFold(n_splits=5,shuffle=True,random_state=42)
onehot=pd.get_dummies(df[cats],dtype=float).reset_index(drop=True)

struct=["sqft_final","has_sqft","lot_sqft","beds","baths","parking","storeys","age","age2"]
amen_=list(AMEN.keys())
S0=struct+amen_
S1=S0+rv+["n_images_c","virtual_tour"]
S2=S1+gv
S3=S2+vv
sets=[("Structural\n+ amenity",S0),("+ Rooms\n+ listing",S1),("+ Text\n(LLM)",S2),("+ Vision\n(façade)",S3)]

def ols_r2(cols):
    Xi=df[cols].copy()
    for c in cols: Xi[c]=Xi[c].fillna(df.groupby("ptype")[c].transform("median")).fillna(df[c].median())
    X=pd.concat([Xi.reset_index(drop=True),onehot],axis=1)
    p=cross_val_predict(LinearRegression(),X,y,cv=kf,n_jobs=-1);return r2_score(y,p),p
def gbm_r2(cols):
    X=pd.concat([df[cols].reset_index(drop=True),onehot],axis=1)
    m=HistGradientBoostingRegressor(max_iter=600,learning_rate=0.05,max_leaf_nodes=63,l2_regularization=1.0,random_state=42)
    p=cross_val_predict(m,X,y,cv=kf,n_jobs=-1);return r2_score(y,p),p

ols=[];gbm=[];gpred=None
for nm,cols in sets:
    ro,_=ols_r2(cols);rg,pg=gbm_r2(cols);ols.append(ro);gbm.append(rg)
    if nm.startswith("+ Vision"): gpred=pg
    print(f"{nm.replace(chr(10),' '):<22} OLS={ro:.4f}  GBM={rg:.4f}")

# blend numbers (déjà calculés) : S2 texte, S3 texte+vision
blend={2:0.760,3:0.779}

# ===== Figure 1 : ablation =====
x=np.arange(len(sets))
fig,ax=plt.subplots(figsize=(6.4,3.9))
ax.plot(x,ols,"-o",color=C_OLS,lw=2,label="OLS (linear)")
ax.plot(x,gbm,"-s",color=C_GBM,lw=2,label="Gradient boosting")
ax.plot([2,3],[blend[2],blend[3]],"-D",color=C_BL,lw=2,label="Blend (NN+GBM)")
for xi,(ro,rg) in enumerate(zip(ols,gbm)):
    ax.annotate(f"{rg:.3f}",(xi,rg),textcoords="offset points",xytext=(0,7),ha="center",fontsize=8,color=C_GBM)
    ax.annotate(f"{ro:.3f}",(xi,ro),textcoords="offset points",xytext=(0,-13),ha="center",fontsize=8,color=C_OLS)
ax.annotate(f"{blend[3]:.3f}",(3,blend[3]),textcoords="offset points",xytext=(0,7),ha="center",fontsize=8,color=C_BL)
ax.set_xticks(x);ax.set_xticklabels([s[0] for s in sets])
ax.set_ylabel(r"Out-of-sample $R^2$ (5-fold CV)");ax.set_ylim(0.55,0.80)
ax.legend(frameon=False,loc="lower right",fontsize=9)
ax.set_title("Incremental gain from each feature layer",fontsize=11)
fig.savefig("paper/figures/fig_ablation.pdf");plt.close(fig)
print("-> fig_ablation.pdf")

# ===== Figure 2 : calibration + résidus (blend v5) =====
try: bl=np.load("data/processed/v5_blend_oos.npy")
except Exception: bl=gpred
yp=np.exp(bl);ya=np.exp(y)
fig,axs=plt.subplots(1,2,figsize=(7.6,3.6))
ax=axs[0]
hb=ax.hexbin(ya/1e3,yp/1e3,gridsize=45,bins="log",cmap="Blues",xscale="log",yscale="log",mincnt=1)
lim=[60,6000];ax.plot(lim,lim,"--",color="black",lw=1)
ax.set_xlim(lim);ax.set_ylim(lim);ax.set_xlabel("Actual price (\\$000)");ax.set_ylabel("Predicted price (\\$000)")
ax.set_title(f"Calibration (blend, $R^2$={r2_score(y,bl):.3f})",fontsize=10)
r=(yp/ya-1)*100
ax2=axs[1]
ax2.hist(np.clip(r,-60,60),bins=50,color=C_BL,alpha=0.85,edgecolor="white",lw=0.3)
ax2.axvline(0,color="black",lw=1,ls="--")
ax2.set_xlabel("Prediction error (\\%)");ax2.set_ylabel("Listings")
ax2.set_title(f"Errors (median |.|={np.median(np.abs(r)):.1f}\\%)",fontsize=10)
fig.savefig("paper/figures/fig_calibration.pdf");plt.close(fig)
print("-> fig_calibration.pdf")

# ===== Figure 3 : importance par bloc (GBM v5) =====
Xfull=pd.concat([df[S3].reset_index(drop=True),onehot],axis=1)
gb=HistGradientBoostingRegressor(max_iter=600,learning_rate=0.05,max_leaf_nodes=63,l2_regularization=1.0,random_state=42).fit(Xfull,y)
idx=np.random.RandomState(0).choice(len(Xfull),4000,replace=False)
pi=permutation_importance(gb,Xfull.iloc[idx],y[idx],n_repeats=3,random_state=0,n_jobs=-1)
imp=pd.Series(pi.importances_mean,index=Xfull.columns).clip(lower=0)
def block(c):
    if c.startswith("city_fe") or c.startswith("ptype_fe") or c.startswith("source"): return "Location/type FE"
    if c in ("sqft_final","has_sqft","lot_sqft"): return "Size"
    if c in ("beds","baths","parking","storeys","age","age2"): return "Structural"
    if c in amen_: return "Amenities (rules)"
    if c in rv or c in ("n_images_c","virtual_tour"): return "Rooms / listing"
    if c in gv: return "Text (LLM)"
    if c in vv: return "Vision (façade)"
    return "Other"
bl_imp=imp.groupby(imp.index.map(block)).sum().sort_values()
bl_imp=bl_imp/bl_imp.sum()*100
fig,ax=plt.subplots(figsize=(6.4,3.6))
colors=["#8899AA"]*len(bl_imp)
for i,n in enumerate(bl_imp.index):
    if n=="Vision (façade)": colors[i]=C_GBM
    if n=="Text (LLM)": colors[i]=C_BL
ax.barh(bl_imp.index,bl_imp.values,color=colors,edgecolor="white")
for i,v in enumerate(bl_imp.values): ax.text(v+0.4,i,f"{v:.0f}\\%",va="center",fontsize=8)
ax.set_xlabel("Share of permutation importance (\\%)")
ax.set_title("Feature-block importance (gradient boosting, v5)",fontsize=11)
ax.set_xlim(0,max(bl_imp.values)*1.15)
fig.savefig("paper/figures/fig_importance.pdf");plt.close(fig)
print("-> fig_importance.pdf")
print("blocs:",{k:round(v,1) for k,v in bl_imp.sort_values(ascending=False).items()})

# ===== Table ablation détaillée =====
def med(p): return np.median(np.abs(np.exp(p-y)-1))*100
rowsL=[]
labels=["Structural + amenity","\\quad + Rooms + listing","\\quad + Text (LLM)","\\quad + Vision (façade)"]
for i,(nm,cols) in enumerate(sets):
    ro,po=ols_r2(cols);rg,pg=gbm_r2(cols)
    bcell = f"{blend[i]:.3f}" if i in blend else "---"
    rowsL.append(f"      {labels[i]} & {len(cols)} & {ro:.3f} & {rg:.3f} & {bcell} \\\\")
tab=r"""% Auto-generated v5
\begin{table}[htbp]
  \centering
  \begin{threeparttable}
    \caption{Feature Ablation: Out-of-Sample $R^2$ by Feature Layer}
    \label{tab:ablation}
    \begin{tabular}{@{}lcccc@{}}
      \toprule
      Feature set & \# feats & OLS & GBM & Blend \\
      \midrule
"""+"\n".join(rowsL)+r"""
      \bottomrule
    \end{tabular}
    \begin{tablenotes}[flushleft]
      \footnotesize
      \item \textit{Notes:} Each row adds a feature layer to the row above (nested).
      Entries are out-of-sample $R^2$ on $\ln(\text{Price})$ from 5-fold cross-validation
      with a common seed. ``\# feats'' counts numeric/indicator features before the city,
      type, and source fixed effects (added in every specification). The blend averages
      the robust neural network and gradient boosting and is reported where both were
      trained. GBM and the neural network handle missing living area natively; OLS uses
      within-type median imputation.
    \end{tablenotes}
  \end{threeparttable}
\end{table}
"""
open("paper/tables/ablation.tex","w").write(tab)
print("-> paper/tables/ablation.tex")