"""
ml_models.py
============
Machine learning models for hedonic pricing analysis of Airbnb listings.

Provides training, evaluation, cross-validation, feature importance extraction,
and hyperparameter tuning for a battery of regression models including
ElasticNet, Random Forest, Gradient Boosting, XGBoost, LightGBM, CatBoost,
and Generalized Additive Models (GAM).

Author: Simon-Pierre Boucher
Date: 2026-05-20
"""

from __future__ import annotations

import logging
import warnings
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
from sklearn.linear_model import ElasticNetCV
from sklearn.metrics import (
    mean_absolute_error,
    mean_absolute_percentage_error,
    mean_squared_error,
    r2_score,
)
from sklearn.model_selection import (
    KFold,
    RandomizedSearchCV,
    cross_val_score,
    train_test_split,
)
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import PolynomialFeatures, StandardScaler

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Data Preparation
# ---------------------------------------------------------------------------


def prepare_ml_data(
    df: pd.DataFrame,
    y_col: str,
    x_cols: List[str],
    test_size: float = 0.2,
    random_state: int = 42,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[str]]:
    """Split data into train/test sets with NaN handling.

    Parameters
    ----------
    df : pd.DataFrame
        Source dataframe.
    y_col : str
        Name of the target (price) column.
    x_cols : list[str]
        Feature column names.
    test_size : float
        Fraction held out for testing.
    random_state : int
        Seed for reproducibility.

    Returns
    -------
    X_train, X_test, y_train, y_test : np.ndarray
        Split arrays.
    feature_names : list[str]
        Column names retained after NaN removal.
    """
    subset = df[x_cols + [y_col]].copy()

    n_before = len(subset)
    subset = subset.dropna()
    n_after = len(subset)
    if n_before != n_after:
        logger.info(
            "Dropped %d rows containing NaN (%d -> %d).",
            n_before - n_after,
            n_before,
            n_after,
        )

    X = subset[x_cols].values.astype(np.float64)
    y = subset[y_col].values.astype(np.float64)
    feature_names = list(x_cols)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=random_state
    )

    logger.info(
        "Train: %d samples, Test: %d samples, Features: %d",
        X_train.shape[0],
        X_test.shape[0],
        X_train.shape[1],
    )
    return X_train, X_test, y_train, y_test, feature_names


# ---------------------------------------------------------------------------
# Individual Model Trainers
# ---------------------------------------------------------------------------


def train_elastic_net(
    X_train: np.ndarray,
    y_train: np.ndarray,
    alphas: Optional[List[float]] = None,
    l1_ratios: Optional[List[float]] = None,
    cv: int = 5,
    random_state: int = 42,
) -> ElasticNetCV:
    """Train an ElasticNet with built-in cross-validated alpha selection.

    Parameters
    ----------
    X_train, y_train : np.ndarray
        Training data.
    alphas : list[float] | None
        Candidate regularisation strengths.  Defaults to a log-spaced grid.
    l1_ratios : list[float] | None
        L1/L2 mixing ratios to try.  Defaults to ``[0.1, 0.5, 0.7, 0.9, 0.95, 1.0]``.
    cv : int
        Number of CV folds.
    random_state : int
        Seed for reproducibility.

    Returns
    -------
    ElasticNetCV
        Fitted model.
    """
    if alphas is None:
        alphas = np.logspace(-4, 2, 50).tolist()
    if l1_ratios is None:
        l1_ratios = [0.1, 0.5, 0.7, 0.9, 0.95, 1.0]

    model = ElasticNetCV(
        alphas=alphas,
        l1_ratio=l1_ratios,
        cv=cv,
        random_state=random_state,
        max_iter=10_000,
        n_jobs=-1,
    )
    model.fit(X_train, y_train)
    logger.info(
        "ElasticNet  alpha=%.6f  l1_ratio=%.2f",
        model.alpha_,
        model.l1_ratio_,
    )
    return model


def train_random_forest(
    X_train: np.ndarray,
    y_train: np.ndarray,
    n_estimators: int = 500,
    max_depth: Optional[int] = None,
    min_samples_leaf: int = 5,
    random_state: int = 42,
) -> RandomForestRegressor:
    """Train a Random Forest regressor.

    Parameters
    ----------
    X_train, y_train : np.ndarray
        Training data.
    n_estimators : int
        Number of trees.
    max_depth : int | None
        Maximum tree depth (``None`` = unlimited).
    min_samples_leaf : int
        Minimum samples per leaf node.
    random_state : int
        Seed for reproducibility.

    Returns
    -------
    RandomForestRegressor
        Fitted model.
    """
    model = RandomForestRegressor(
        n_estimators=n_estimators,
        max_depth=max_depth,
        min_samples_leaf=min_samples_leaf,
        random_state=random_state,
        n_jobs=-1,
    )
    model.fit(X_train, y_train)
    logger.info("RandomForest  n_estimators=%d  fitted.", n_estimators)
    return model


def train_gradient_boosting(
    X_train: np.ndarray,
    y_train: np.ndarray,
    n_estimators: int = 500,
    learning_rate: float = 0.05,
    max_depth: int = 5,
    subsample: float = 0.8,
    random_state: int = 42,
) -> GradientBoostingRegressor:
    """Train a scikit-learn GradientBoostingRegressor.

    Parameters
    ----------
    X_train, y_train : np.ndarray
        Training data.
    n_estimators : int
        Boosting rounds.
    learning_rate : float
        Shrinkage factor.
    max_depth : int
        Maximum tree depth per stage.
    subsample : float
        Row subsampling fraction.
    random_state : int
        Seed for reproducibility.

    Returns
    -------
    GradientBoostingRegressor
        Fitted model.
    """
    model = GradientBoostingRegressor(
        n_estimators=n_estimators,
        learning_rate=learning_rate,
        max_depth=max_depth,
        subsample=subsample,
        random_state=random_state,
    )
    model.fit(X_train, y_train)
    logger.info("GradientBoosting  fitted (n_estimators=%d).", n_estimators)
    return model


def train_xgboost(
    X_train: np.ndarray,
    y_train: np.ndarray,
    n_estimators: int = 500,
    learning_rate: float = 0.05,
    max_depth: int = 6,
    subsample: float = 0.8,
    colsample_bytree: float = 0.8,
    random_state: int = 42,
) -> Any:
    """Train an XGBRegressor.  Falls back to sklearn GradientBoostingRegressor
    if ``xgboost`` is not installed.

    Returns
    -------
    model
        Fitted regressor (XGBRegressor or GradientBoostingRegressor).
    """
    try:
        from xgboost import XGBRegressor

        model = XGBRegressor(
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            max_depth=max_depth,
            subsample=subsample,
            colsample_bytree=colsample_bytree,
            random_state=random_state,
            n_jobs=-1,
            verbosity=0,
        )
        model.fit(X_train, y_train)
        logger.info("XGBoost  fitted (n_estimators=%d).", n_estimators)
        return model
    except ImportError:
        logger.warning("xgboost not installed — falling back to sklearn GBR.")
        return train_gradient_boosting(
            X_train,
            y_train,
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            max_depth=max_depth,
            subsample=subsample,
            random_state=random_state,
        )


def train_lightgbm(
    X_train: np.ndarray,
    y_train: np.ndarray,
    n_estimators: int = 500,
    learning_rate: float = 0.05,
    num_leaves: int = 31,
    subsample: float = 0.8,
    colsample_bytree: float = 0.8,
    random_state: int = 42,
) -> Any:
    """Train an LGBMRegressor.  Falls back to sklearn GBR if ``lightgbm`` is
    not installed.

    Returns
    -------
    model
        Fitted regressor.
    """
    try:
        from lightgbm import LGBMRegressor

        model = LGBMRegressor(
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            num_leaves=num_leaves,
            subsample=subsample,
            colsample_bytree=colsample_bytree,
            random_state=random_state,
            n_jobs=-1,
            verbose=-1,
        )
        model.fit(X_train, y_train)
        logger.info("LightGBM  fitted (n_estimators=%d).", n_estimators)
        return model
    except ImportError:
        logger.warning("lightgbm not installed — falling back to sklearn GBR.")
        return train_gradient_boosting(
            X_train,
            y_train,
            n_estimators=n_estimators,
            learning_rate=learning_rate,
            max_depth=5,
            subsample=subsample,
            random_state=random_state,
        )


def train_catboost(
    X_train: np.ndarray,
    y_train: np.ndarray,
    iterations: int = 500,
    learning_rate: float = 0.05,
    depth: int = 6,
    random_state: int = 42,
) -> Any:
    """Train a CatBoostRegressor.  Falls back to sklearn GBR if ``catboost``
    is not installed.

    Returns
    -------
    model
        Fitted regressor.
    """
    try:
        from catboost import CatBoostRegressor

        model = CatBoostRegressor(
            iterations=iterations,
            learning_rate=learning_rate,
            depth=depth,
            random_state=random_state,
            verbose=0,
        )
        model.fit(X_train, y_train)
        logger.info("CatBoost  fitted (iterations=%d).", iterations)
        return model
    except ImportError:
        logger.warning("catboost not installed — falling back to sklearn GBR.")
        return train_gradient_boosting(
            X_train,
            y_train,
            n_estimators=iterations,
            learning_rate=learning_rate,
            max_depth=depth,
            random_state=random_state,
        )


def train_gam(
    X_train: np.ndarray,
    y_train: np.ndarray,
    n_splines: int = 20,
    degree: int = 2,
) -> Any:
    """Train a Generalized Additive Model.

    Attempts to use ``pygam.LinearGAM``; if unavailable, builds a polynomial-
    feature pipeline as a lightweight approximation.

    Returns
    -------
    model
        Fitted GAM or polynomial pipeline.
    """
    try:
        from pygam import LinearGAM, s

        # Build one spline term per feature.
        terms = s(0, n_splines=n_splines)
        for j in range(1, X_train.shape[1]):
            terms += s(j, n_splines=n_splines)

        model = LinearGAM(terms)
        model.gridsearch(X_train, y_train, progress=False)
        logger.info("pygam LinearGAM  fitted (n_splines=%d).", n_splines)
        return model
    except ImportError:
        logger.warning(
            "pygam not installed — approximating GAM with polynomial features (degree=%d).",
            degree,
        )
        pipe = Pipeline(
            [
                ("poly", PolynomialFeatures(degree=degree, interaction_only=True, include_bias=False)),
                ("scaler", StandardScaler()),
                (
                    "enet",
                    ElasticNetCV(
                        l1_ratio=[0.5, 0.7, 0.9, 0.95, 1.0],
                        cv=5,
                        max_iter=10_000,
                        n_jobs=-1,
                    ),
                ),
            ]
        )
        pipe.fit(X_train, y_train)
        logger.info("Polynomial GAM approximation  fitted (degree=%d).", degree)
        return pipe


# ---------------------------------------------------------------------------
# Train All Models
# ---------------------------------------------------------------------------


def train_all_models(
    X_train: np.ndarray,
    y_train: np.ndarray,
) -> Dict[str, Any]:
    """Train every available model and return a name -> fitted model dict.

    Models that fail to train (e.g. convergence issues) are logged and
    skipped rather than raising.

    Returns
    -------
    dict[str, model]
        Mapping of model name to fitted estimator.
    """
    trainers: Dict[str, Any] = {
        "ElasticNet": train_elastic_net,
        "RandomForest": train_random_forest,
        "GradientBoosting": train_gradient_boosting,
        "XGBoost": train_xgboost,
        "LightGBM": train_lightgbm,
        "CatBoost": train_catboost,
        "GAM": train_gam,
    }

    models: Dict[str, Any] = {}
    for name, trainer in trainers.items():
        try:
            logger.info("Training %s ...", name)
            models[name] = trainer(X_train, y_train)
        except Exception:
            logger.exception("Failed to train %s — skipping.", name)
    return models


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------


def evaluate_model(
    model: Any,
    X_test: np.ndarray,
    y_test: np.ndarray,
) -> Dict[str, float]:
    """Compute regression metrics for a fitted model.

    Returns
    -------
    dict
        Keys: ``R2``, ``RMSE``, ``MAE``, ``MAPE``.
    """
    y_pred = model.predict(X_test)
    r2 = r2_score(y_test, y_pred)
    rmse = float(np.sqrt(mean_squared_error(y_test, y_pred)))
    mae = mean_absolute_error(y_test, y_pred)

    # Guard against zero prices when computing MAPE.
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        mask = y_test != 0
        if mask.sum() > 0:
            mape = float(np.mean(np.abs((y_test[mask] - y_pred[mask]) / y_test[mask]))) * 100
        else:
            mape = np.nan

    metrics = {"R2": r2, "RMSE": rmse, "MAE": mae, "MAPE": mape}
    logger.info(
        "Evaluation  R2=%.4f  RMSE=%.2f  MAE=%.2f  MAPE=%.2f%%",
        r2,
        rmse,
        mae,
        mape,
    )
    return metrics


def compare_all_models(
    models_dict: Dict[str, Any],
    X_test: np.ndarray,
    y_test: np.ndarray,
) -> pd.DataFrame:
    """Evaluate every model and return a comparison DataFrame sorted by R2.

    Parameters
    ----------
    models_dict : dict[str, model]
        Output of :func:`train_all_models`.
    X_test, y_test : np.ndarray
        Held-out test data.

    Returns
    -------
    pd.DataFrame
        Columns: ``Model``, ``R2``, ``RMSE``, ``MAE``, ``MAPE``.
    """
    rows: List[Dict[str, Any]] = []
    for name, model in models_dict.items():
        try:
            metrics = evaluate_model(model, X_test, y_test)
            rows.append({"Model": name, **metrics})
        except Exception:
            logger.exception("Evaluation failed for %s.", name)

    comparison = pd.DataFrame(rows).sort_values("R2", ascending=False).reset_index(drop=True)
    logger.info("\n%s", comparison.to_string(index=False))
    return comparison


# ---------------------------------------------------------------------------
# Cross-Validation
# ---------------------------------------------------------------------------


def cross_validate_model(
    model: Any,
    X: np.ndarray,
    y: np.ndarray,
    k: int = 5,
    random_state: int = 42,
) -> Dict[str, float]:
    """K-fold cross-validation returning mean and std of R2, RMSE, and MAE.

    Parameters
    ----------
    model : estimator
        An **unfitted** or fitted sklearn-compatible estimator. A clone is
        created internally, so the caller's model is not mutated.
    X, y : np.ndarray
        Full dataset (train + test).
    k : int
        Number of folds.
    random_state : int
        Seed for KFold split.

    Returns
    -------
    dict
        Keys: ``R2_mean``, ``R2_std``, ``RMSE_mean``, ``RMSE_std``,
        ``MAE_mean``, ``MAE_std``.
    """
    from sklearn.base import clone

    kf = KFold(n_splits=k, shuffle=True, random_state=random_state)
    r2_scores: List[float] = []
    rmse_scores: List[float] = []
    mae_scores: List[float] = []

    for train_idx, val_idx in kf.split(X):
        X_tr, X_val = X[train_idx], X[val_idx]
        y_tr, y_val = y[train_idx], y[val_idx]

        m = clone(model)
        m.fit(X_tr, y_tr)
        y_pred = m.predict(X_val)

        r2_scores.append(r2_score(y_val, y_pred))
        rmse_scores.append(float(np.sqrt(mean_squared_error(y_val, y_pred))))
        mae_scores.append(mean_absolute_error(y_val, y_pred))

    results = {
        "R2_mean": float(np.mean(r2_scores)),
        "R2_std": float(np.std(r2_scores)),
        "RMSE_mean": float(np.mean(rmse_scores)),
        "RMSE_std": float(np.std(rmse_scores)),
        "MAE_mean": float(np.mean(mae_scores)),
        "MAE_std": float(np.std(mae_scores)),
    }
    logger.info(
        "CV(%d)  R2=%.4f +/- %.4f  RMSE=%.2f +/- %.2f",
        k,
        results["R2_mean"],
        results["R2_std"],
        results["RMSE_mean"],
        results["RMSE_std"],
    )
    return results


# ---------------------------------------------------------------------------
# Feature Importance
# ---------------------------------------------------------------------------


def get_feature_importance(
    model: Any,
    feature_names: List[str],
) -> pd.DataFrame:
    """Extract feature importances from a fitted model.

    Supports tree-based models (``feature_importances_``), linear models
    (``coef_``), and ``pygam`` (partial dependence variance).

    Parameters
    ----------
    model : estimator
        Fitted model.
    feature_names : list[str]
        Feature column names.

    Returns
    -------
    pd.DataFrame
        Columns: ``Feature``, ``Importance``, sorted descending.
    """
    importances: Optional[np.ndarray] = None

    # Tree-based models (RF, GBR, XGB, LGBM, CatBoost).
    if hasattr(model, "feature_importances_"):
        importances = np.asarray(model.feature_importances_)

    # Linear models (ElasticNet, Ridge, Lasso).
    elif hasattr(model, "coef_"):
        importances = np.abs(np.asarray(model.coef_))

    # pygam — approximate importance via partial dependence variance.
    elif hasattr(model, "partial_dependence"):
        try:
            variances = []
            for j in range(len(feature_names)):
                XX = model.generate_X_grid(term=j)
                pdep, _ = model.partial_dependence(term=j, X=XX, width=0.95)
                variances.append(float(np.var(pdep)))
            importances = np.array(variances)
        except Exception:
            logger.warning("Could not extract importance from pygam model.")

    # Pipeline fallback — try last step.
    elif hasattr(model, "named_steps"):
        last_step = list(model.named_steps.values())[-1]
        if hasattr(last_step, "coef_"):
            # Polynomial pipeline: importances don't map 1-to-1 to original
            # features.  Aggregate absolute coefficients.
            poly = model.named_steps.get("poly")
            if poly is not None:
                coefs = np.abs(last_step.coef_)
                powers = poly.powers_
                agg = np.zeros(len(feature_names))
                for idx, power_vec in enumerate(powers):
                    for feat_idx in np.nonzero(power_vec)[0]:
                        if feat_idx < len(feature_names):
                            agg[feat_idx] += coefs[idx] if idx < len(coefs) else 0.0
                importances = agg
            else:
                importances = np.abs(np.asarray(last_step.coef_))

    if importances is None:
        logger.warning(
            "Cannot extract importance from model type %s.",
            type(model).__name__,
        )
        return pd.DataFrame(columns=["Feature", "Importance"])

    # Truncate / pad to match feature_names length.
    if len(importances) > len(feature_names):
        importances = importances[: len(feature_names)]
    elif len(importances) < len(feature_names):
        importances = np.pad(importances, (0, len(feature_names) - len(importances)))

    df_imp = (
        pd.DataFrame({"Feature": feature_names, "Importance": importances})
        .sort_values("Importance", ascending=False)
        .reset_index(drop=True)
    )
    return df_imp


# ---------------------------------------------------------------------------
# Hyperparameter Tuning
# ---------------------------------------------------------------------------


def tune_best_model(
    X_train: np.ndarray,
    y_train: np.ndarray,
    model_type: str = "xgboost",
    n_iter: int = 50,
    cv: int = 5,
    random_state: int = 42,
    scoring: str = "neg_root_mean_squared_error",
) -> Tuple[Any, Dict[str, Any]]:
    """Randomised hyperparameter search for a given model type.

    Parameters
    ----------
    X_train, y_train : np.ndarray
        Training data.
    model_type : str
        One of ``'xgboost'``, ``'lightgbm'``, ``'catboost'``,
        ``'random_forest'``, ``'gradient_boosting'``.
    n_iter : int
        Number of parameter settings sampled.
    cv : int
        Cross-validation folds.
    random_state : int
        Seed for reproducibility.
    scoring : str
        Metric for ``RandomizedSearchCV``.

    Returns
    -------
    best_model
        Refitted model with the best parameter set.
    best_params : dict
        Best parameters found.
    """

    # --- estimator + param distributions -----------------------------------
    if model_type == "xgboost":
        try:
            from xgboost import XGBRegressor

            estimator = XGBRegressor(random_state=random_state, n_jobs=-1, verbosity=0)
        except ImportError:
            logger.warning("xgboost unavailable — tuning sklearn GBR instead.")
            estimator = GradientBoostingRegressor(random_state=random_state)
            model_type = "gradient_boosting"
    elif model_type == "lightgbm":
        try:
            from lightgbm import LGBMRegressor

            estimator = LGBMRegressor(random_state=random_state, n_jobs=-1, verbose=-1)
        except ImportError:
            logger.warning("lightgbm unavailable — tuning sklearn GBR instead.")
            estimator = GradientBoostingRegressor(random_state=random_state)
            model_type = "gradient_boosting"
    elif model_type == "catboost":
        try:
            from catboost import CatBoostRegressor

            estimator = CatBoostRegressor(random_state=random_state, verbose=0)
        except ImportError:
            logger.warning("catboost unavailable — tuning sklearn GBR instead.")
            estimator = GradientBoostingRegressor(random_state=random_state)
            model_type = "gradient_boosting"
    elif model_type == "random_forest":
        estimator = RandomForestRegressor(random_state=random_state, n_jobs=-1)
    else:
        estimator = GradientBoostingRegressor(random_state=random_state)
        model_type = "gradient_boosting"

    param_grids: Dict[str, Dict[str, Any]] = {
        "xgboost": {
            "n_estimators": [200, 500, 800, 1000],
            "learning_rate": [0.01, 0.03, 0.05, 0.1],
            "max_depth": [3, 4, 5, 6, 8],
            "subsample": [0.6, 0.7, 0.8, 0.9, 1.0],
            "colsample_bytree": [0.6, 0.7, 0.8, 0.9, 1.0],
            "min_child_weight": [1, 3, 5, 7],
            "reg_alpha": [0, 0.01, 0.1, 1],
            "reg_lambda": [0.5, 1, 2, 5],
        },
        "lightgbm": {
            "n_estimators": [200, 500, 800, 1000],
            "learning_rate": [0.01, 0.03, 0.05, 0.1],
            "num_leaves": [15, 31, 63, 127],
            "subsample": [0.6, 0.7, 0.8, 0.9, 1.0],
            "colsample_bytree": [0.6, 0.7, 0.8, 0.9, 1.0],
            "min_child_samples": [5, 10, 20, 50],
            "reg_alpha": [0, 0.01, 0.1, 1],
            "reg_lambda": [0.5, 1, 2, 5],
        },
        "catboost": {
            "iterations": [200, 500, 800, 1000],
            "learning_rate": [0.01, 0.03, 0.05, 0.1],
            "depth": [4, 6, 8, 10],
            "l2_leaf_reg": [1, 3, 5, 7, 9],
            "bagging_temperature": [0, 0.5, 1],
        },
        "random_forest": {
            "n_estimators": [200, 500, 800, 1000],
            "max_depth": [None, 10, 20, 30, 50],
            "min_samples_leaf": [1, 2, 5, 10],
            "min_samples_split": [2, 5, 10],
            "max_features": ["sqrt", "log2", 0.5, 0.8, None],
        },
        "gradient_boosting": {
            "n_estimators": [200, 500, 800, 1000],
            "learning_rate": [0.01, 0.03, 0.05, 0.1],
            "max_depth": [3, 4, 5, 6, 8],
            "subsample": [0.6, 0.7, 0.8, 0.9, 1.0],
            "min_samples_leaf": [1, 2, 5, 10],
        },
    }

    param_dist = param_grids.get(model_type, param_grids["gradient_boosting"])

    search = RandomizedSearchCV(
        estimator,
        param_distributions=param_dist,
        n_iter=n_iter,
        scoring=scoring,
        cv=cv,
        random_state=random_state,
        n_jobs=-1,
        verbose=0,
    )
    search.fit(X_train, y_train)

    logger.info(
        "Best %s params: %s  (score=%.4f)",
        model_type,
        search.best_params_,
        search.best_score_,
    )
    return search.best_estimator_, search.best_params_