# Author: Simon-Pierre Boucher — contact@spboucher.ai
"""
09_ml_robustness.py
-------------------
Model 6: ML Robustness Models.
Predict log_rent using Airbnb exposure and property characteristics.
Compare OLS, LASSO, Elastic Net, Random Forest, and Gradient Boosting.
Compute SHAP values (if available) or sklearn feature importances.

Outputs
-------
- results/tables/ml_comparison.tex
- results/tables/ml_comparison.csv
- figures/ml_predicted_vs_actual.pdf
- figures/feature_importance.pdf
- figures/shap_summary.pdf  (if shap available)
"""

import sys
import warnings
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.plotting import use_publication_style

use_publication_style()

import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402
from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor  # noqa: E402
from sklearn.linear_model import ElasticNetCV, LassoCV, LinearRegression  # noqa: E402
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score  # noqa: E402
from sklearn.model_selection import train_test_split  # noqa: E402
from sklearn.preprocessing import StandardScaler  # noqa: E402

from src.config import (  # noqa: E402
    MERGED_ANALYSIS,
    TABLE_DIR,
    FIG_DIR,
    RANDOM_STATE,
    require,
)

warnings.filterwarnings("ignore")

PROCESSED_HINT = "Run scripts/04_merge_data.py first."

FEATURES = [
    "airbnb_count_500m",
    "airbnb_density_500m",
    "mean_airbnb_price_500m",
    "share_entire_home_500m",
    "bedrooms",
    "bathrooms",
    "lat",
    "lon",
]
TARGET = "log_rent"

# Linear models use standardized features; tree-based models use raw features
USE_SCALED = {"OLS", "LASSO", "Elastic Net"}


def train_models(X_train, X_test, y_train, y_test, X_train_sc, X_test_sc):
    """Train the five models and return (results DataFrame, fitted dict)."""
    models = {
        "OLS": LinearRegression(),
        "LASSO": LassoCV(cv=5, random_state=RANDOM_STATE, max_iter=10_000),
        "Elastic Net": ElasticNetCV(cv=5, random_state=RANDOM_STATE, max_iter=10_000),
        "Random Forest": RandomForestRegressor(
            n_estimators=500, max_depth=15, random_state=RANDOM_STATE, n_jobs=-1
        ),
        "Gradient Boosting": GradientBoostingRegressor(
            n_estimators=500, max_depth=5, learning_rate=0.05,
            random_state=RANDOM_STATE
        ),
    }

    results = []
    fitted = {}

    for name, model in models.items():
        print(f"Training {name} ...")
        Xtr = X_train_sc if name in USE_SCALED else X_train
        Xte = X_test_sc if name in USE_SCALED else X_test

        model.fit(Xtr, y_train)
        fitted[name] = model

        y_pred_train = model.predict(Xtr)
        y_pred_test = model.predict(Xte)

        r2_train = r2_score(y_train, y_pred_train)
        r2_test = r2_score(y_test, y_pred_test)
        rmse_test = np.sqrt(mean_squared_error(y_test, y_pred_test))
        mae_test = mean_absolute_error(y_test, y_pred_test)

        results.append({
            "Model": name,
            "R2_train": r2_train,
            "R2_test": r2_test,
            "RMSE_test": rmse_test,
            "MAE_test": mae_test,
        })
        print(f"  R2 train={r2_train:.4f}  test={r2_test:.4f}  "
              f"RMSE={rmse_test:.4f}  MAE={mae_test:.4f}")

    return pd.DataFrame(results), fitted


def save_comparison_table(res_df: pd.DataFrame) -> None:
    """Write ml_comparison.csv and ml_comparison.tex."""
    res_df.to_csv(TABLE_DIR / "ml_comparison.csv", index=False)

    latex_rows = []
    latex_rows.append(r"\begin{tabular}{lcccc}")
    latex_rows.append(r"\toprule")
    latex_rows.append(r"Model & $R^2$ (Train) & $R^2$ (Test) & RMSE (Test) & MAE (Test) \\")
    latex_rows.append(r"\midrule")
    for _, row in res_df.iterrows():
        latex_rows.append(
            f"{row['Model']} & {row['R2_train']:.4f} & {row['R2_test']:.4f} "
            f"& {row['RMSE_test']:.4f} & {row['MAE_test']:.4f} \\\\"
        )
    latex_rows.append(r"\bottomrule")
    latex_rows.append(r"\end{tabular}")
    latex_rows.append(r"\begin{tablenotes}\small")
    latex_rows.append(r"\item \textit{Notes:} OLS, LASSO, and Elastic Net use "
                      r"standardized features. Tree-based models use raw features. "
                      r"Train/test split is 80/20 with random\_state=42.")
    latex_rows.append(r"\end{tablenotes}")

    (TABLE_DIR / "ml_comparison.tex").write_text("\n".join(latex_rows) + "\n",
                                                 encoding="utf-8")
    print(f"\nSaved {TABLE_DIR / 'ml_comparison.tex'}")
    print(f"Saved {TABLE_DIR / 'ml_comparison.csv'}")


def plot_predicted_vs_actual(res_df, fitted, X_test, X_test_sc, y_test) -> None:
    """Scatter of predicted vs actual log rent for the best model by test R2."""
    best_name = res_df.loc[res_df["R2_test"].idxmax(), "Model"]
    best_model = fitted[best_name]
    Xte_best = X_test_sc if best_name in USE_SCALED else X_test
    y_pred_best = best_model.predict(Xte_best)

    fig, ax = plt.subplots()
    ax.scatter(y_test, y_pred_best, alpha=0.3, s=8, edgecolors="none", c="#2166ac")
    mn, mx = min(y_test.min(), y_pred_best.min()), max(y_test.max(), y_pred_best.max())
    ax.plot([mn, mx], [mn, mx], "k--", lw=1, label="45-degree line")
    ax.set_xlabel("Actual log(rent)")
    ax.set_ylabel("Predicted log(rent)")
    ax.set_title(f"Predicted vs Actual — {best_name} "
                 f"(test R²={res_df.loc[res_df['R2_test'].idxmax(), 'R2_test']:.4f})")
    ax.legend(loc="upper left", frameon=False)
    fig.savefig(FIG_DIR / "ml_predicted_vs_actual.pdf")
    plt.close(fig)
    print(f"Saved {FIG_DIR / 'ml_predicted_vs_actual.pdf'}")


def plot_importances(fitted, X_test) -> None:
    """SHAP summary + importance plots, with sklearn fallback."""
    shap_available = False
    try:
        import shap
        shap_available = True
        print("\nSHAP library found — computing SHAP values ...")
    except ImportError:
        print("\nSHAP not available — using sklearn feature importances instead.")

    if shap_available:
        # Use the best tree model for SHAP; prefer Gradient Boosting, fall back RF
        for shap_model_name in ["Gradient Boosting", "Random Forest"]:
            if shap_model_name in fitted:
                break
        shap_model = fitted[shap_model_name]
        Xte_shap = X_test  # tree models use raw features

        try:
            explainer = shap.TreeExplainer(shap_model)
            shap_values = explainer.shap_values(Xte_shap)

            # Summary bee-swarm plot
            plt.figure()
            shap.summary_plot(
                shap_values, Xte_shap, feature_names=FEATURES, show=False
            )
            plt.tight_layout()
            plt.savefig(FIG_DIR / "shap_summary.pdf")
            plt.close()
            print(f"Saved {FIG_DIR / 'shap_summary.pdf'}")

            # Bar plot of mean |SHAP|
            plt.figure()
            shap.summary_plot(
                shap_values, Xte_shap, feature_names=FEATURES,
                plot_type="bar", show=False
            )
            plt.tight_layout()
            plt.savefig(FIG_DIR / "feature_importance.pdf")
            plt.close()
            print(f"Saved {FIG_DIR / 'feature_importance.pdf'}")

        except Exception as e:
            print(f"SHAP computation failed ({e}); falling back to sklearn importances.")
            shap_available = False

    if not shap_available:
        # Fallback: sklearn feature importances from tree models
        for imp_model_name in ["Gradient Boosting", "Random Forest"]:
            if imp_model_name in fitted and hasattr(fitted[imp_model_name],
                                                    "feature_importances_"):
                break
        imp_model = fitted[imp_model_name]
        importances = imp_model.feature_importances_
        idx = np.argsort(importances)[::-1]

        fig, ax = plt.subplots(figsize=(7, 5))
        ax.barh(
            range(len(FEATURES)),
            importances[idx[::-1]],
            color="#2166ac",
            edgecolor="white",
        )
        ax.set_yticks(range(len(FEATURES)))
        ax.set_yticklabels([FEATURES[i] for i in idx[::-1]])
        ax.set_xlabel("Feature Importance")
        ax.set_title(f"Feature Importance — {imp_model_name}")
        fig.savefig(FIG_DIR / "feature_importance.pdf")
        plt.close(fig)
        print(f"Saved {FIG_DIR / 'feature_importance.pdf'}")


def main() -> None:
    # 1. Load and prepare data
    print("Loading data ...")
    df = pd.read_parquet(require(MERGED_ANALYSIS, PROCESSED_HINT))

    cols = FEATURES + [TARGET]
    df_ml = df[cols].dropna()
    print(f"Sample size after dropping NaN: {len(df_ml):,}")

    X = df_ml[FEATURES].values
    y = df_ml[TARGET].values

    # 2. Train / test split
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.20, random_state=RANDOM_STATE
    )
    print(f"Train: {len(X_train):,}  |  Test: {len(X_test):,}")

    # 3. Standardize features
    scaler = StandardScaler()
    X_train_sc = scaler.fit_transform(X_train)
    X_test_sc = scaler.transform(X_test)

    # 4. Define and train models
    res_df, fitted = train_models(X_train, X_test, y_train, y_test,
                                  X_train_sc, X_test_sc)

    # 5. Save comparison table
    save_comparison_table(res_df)

    # 6. Predicted vs actual scatter (best model by test R²)
    plot_predicted_vs_actual(res_df, fitted, X_test, X_test_sc, y_test)

    # 7. SHAP values or sklearn feature importances
    plot_importances(fitted, X_test)

    print("\n=== 09_ml_robustness.py complete ===")


if __name__ == "__main__":
    main()