"""
shap_analysis.py
================
SHAP-based interpretability analysis for hedonic pricing models.

Provides SHAP value computation, publication-quality summary / dependence /
interaction / waterfall plots, partial-dependence plots, accumulated local
effects (ALE), and hedonic premium interpretation.  Gracefully degrades to
permutation importance and manual PDP when the ``shap`` library is not
installed.

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

from __future__ import annotations

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

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.inspection import PartialDependenceDisplay, partial_dependence
from sklearn.metrics import r2_score

logger = logging.getLogger(__name__)

# Publication-quality defaults ------------------------------------------------
matplotlib.rcParams.update(
    {
        "figure.dpi": 150,
        "savefig.dpi": 300,
        "font.family": "serif",
        "font.size": 11,
        "axes.titlesize": 13,
        "axes.labelsize": 12,
        "xtick.labelsize": 10,
        "ytick.labelsize": 10,
        "legend.fontsize": 10,
        "figure.figsize": (10, 7),
        "savefig.bbox": "tight",
        "savefig.pad_inches": 0.1,
    }
)

# Try importing shap ----------------------------------------------------------
_SHAP_AVAILABLE = False
try:
    import shap

    _SHAP_AVAILABLE = True
except ImportError:
    logger.info(
        "shap library not installed. Falling back to permutation importance "
        "and manual PDP implementations."
    )


# ---------------------------------------------------------------------------
# Utility helpers
# ---------------------------------------------------------------------------


def _ensure_dir(path: str) -> None:
    """Create parent directories for *path* if they do not exist."""
    Path(path).parent.mkdir(parents=True, exist_ok=True)


def _save_figure(fig: plt.Figure, save_path: str) -> None:
    """Save a matplotlib figure in both PNG and PDF."""
    _ensure_dir(save_path)
    base, _ = os.path.splitext(save_path)
    fig.savefig(f"{base}.png", format="png")
    fig.savefig(f"{base}.pdf", format="pdf")
    logger.info("Figure saved: %s.png / .pdf", base)
    plt.close(fig)


def _to_dataframe(X: np.ndarray, feature_names: List[str]) -> pd.DataFrame:
    """Convert an ndarray to a DataFrame with proper column names."""
    return pd.DataFrame(X, columns=feature_names)


# ---------------------------------------------------------------------------
# Permutation-importance fallback (when shap is unavailable)
# ---------------------------------------------------------------------------


def _permutation_importance(
    model: Any,
    X: np.ndarray,
    y: Optional[np.ndarray] = None,
    n_repeats: int = 10,
    random_state: int = 42,
) -> np.ndarray:
    """Compute permutation importance as a (n_features,) array of mean
    absolute importance values.

    When *y* is not provided, variance of prediction changes is used instead.
    """
    rng = np.random.RandomState(random_state)
    baseline = model.predict(X)

    if y is not None:
        baseline_score = r2_score(y, baseline)
    else:
        baseline_score = None

    importances = np.zeros(X.shape[1])
    for j in range(X.shape[1]):
        scores: List[float] = []
        for _ in range(n_repeats):
            X_perm = X.copy()
            X_perm[:, j] = rng.permutation(X_perm[:, j])
            preds_perm = model.predict(X_perm)
            if y is not None and baseline_score is not None:
                scores.append(baseline_score - r2_score(y, preds_perm))
            else:
                scores.append(float(np.mean(np.abs(baseline - preds_perm))))
        importances[j] = float(np.mean(scores))
    return importances


# ---------------------------------------------------------------------------
# SHAP value computation
# ---------------------------------------------------------------------------


def compute_shap_values(
    model: Any,
    X: np.ndarray,
    feature_names: List[str],
    max_samples: int = 500,
) -> Union["shap.Explanation", np.ndarray]:
    """Compute SHAP values for the given model and data.

    Strategy:
    1. If ``shap`` is available, use ``TreeExplainer`` for tree-based models;
       otherwise fall back to ``KernelExplainer`` (on a subsample for speed).
    2. If ``shap`` is **not** installed, return a permutation-importance
       matrix shaped ``(n_samples, n_features)`` as a rough substitute.

    Parameters
    ----------
    model : estimator
        Fitted sklearn-compatible model.
    X : np.ndarray
        Feature matrix.
    feature_names : list[str]
        Column names.
    max_samples : int
        Maximum number of background samples for ``KernelExplainer``.

    Returns
    -------
    shap.Explanation | np.ndarray
        SHAP values object or permutation-based importance matrix.
    """
    if _SHAP_AVAILABLE:
        model_type = type(model).__name__
        tree_types = {
            "RandomForestRegressor",
            "GradientBoostingRegressor",
            "XGBRegressor",
            "LGBMRegressor",
            "CatBoostRegressor",
        }

        X_df = _to_dataframe(X, feature_names)

        if model_type in tree_types:
            try:
                explainer = shap.TreeExplainer(model)
                sv = explainer(X_df)
                logger.info(
                    "SHAP TreeExplainer computed for %s (%d samples).",
                    model_type,
                    X.shape[0],
                )
                return sv
            except Exception:
                logger.warning(
                    "TreeExplainer failed for %s — falling back to KernelExplainer.",
                    model_type,
                )

        # KernelExplainer fallback.
        n_bg = min(max_samples, X.shape[0])
        bg = shap.sample(X_df, n_bg)
        explainer = shap.KernelExplainer(model.predict, bg)
        sv = explainer(X_df)
        logger.info(
            "SHAP KernelExplainer computed for %s (%d samples, %d bg).",
            model_type,
            X.shape[0],
            n_bg,
        )
        return sv
    else:
        # Permutation-importance fallback: tile per-feature importance across
        # samples to produce an (n_samples, n_features) matrix.
        logger.warning(
            "shap not installed — computing permutation importance as fallback."
        )
        imp = _permutation_importance(model, X)
        # Normalise to approximate per-sample attribution.
        imp_matrix = np.tile(imp, (X.shape[0], 1))
        return imp_matrix


# ---------------------------------------------------------------------------
# SHAP Plots
# ---------------------------------------------------------------------------


def _extract_values_array(
    shap_values: Union[Any, np.ndarray],
) -> np.ndarray:
    """Get a plain (n_samples, n_features) ndarray from shap Explanation or raw array."""
    if hasattr(shap_values, "values"):
        return np.asarray(shap_values.values)
    return np.asarray(shap_values)


def plot_shap_summary(
    shap_values: Union[Any, np.ndarray],
    X: np.ndarray,
    feature_names: List[str],
    save_path: str,
    max_display: int = 20,
) -> None:
    """SHAP beeswarm summary plot (or bar chart fallback).

    Parameters
    ----------
    shap_values : shap.Explanation | np.ndarray
        SHAP values.
    X : np.ndarray
        Feature matrix (for colouring dots).
    feature_names : list[str]
        Column names.
    save_path : str
        Output file path (without extension; both .png and .pdf are saved).
    max_display : int
        Maximum number of features to show.
    """
    X_df = _to_dataframe(X, feature_names)

    if _SHAP_AVAILABLE and hasattr(shap_values, "values"):
        fig, ax = plt.subplots(figsize=(10, max(7, max_display * 0.35)))
        shap.summary_plot(
            shap_values,
            X_df,
            max_display=max_display,
            show=False,
        )
        fig = plt.gcf()
        fig.suptitle("SHAP Feature Importance", y=1.02, fontweight="bold")
        _save_figure(fig, save_path)
    else:
        # Fallback: horizontal bar chart of mean |values|.
        vals = _extract_values_array(shap_values)
        mean_abs = np.mean(np.abs(vals), axis=0)
        order = np.argsort(mean_abs)[::-1][:max_display]

        fig, ax = plt.subplots(figsize=(10, max(6, max_display * 0.35)))
        ax.barh(
            range(len(order)),
            mean_abs[order][::-1],
            color="#1f77b4",
            edgecolor="white",
        )
        ax.set_yticks(range(len(order)))
        ax.set_yticklabels([feature_names[i] for i in order][::-1])
        ax.set_xlabel("Mean |SHAP value| (or permutation importance)")
        ax.set_title("Feature Importance", fontweight="bold")
        fig.tight_layout()
        _save_figure(fig, save_path)


def plot_shap_dependence(
    shap_values: Union[Any, np.ndarray],
    X: np.ndarray,
    feature: str,
    feature_names: List[str],
    save_path: str,
    interaction_feature: Optional[str] = None,
) -> None:
    """SHAP dependence plot for a single feature.

    Parameters
    ----------
    shap_values : shap.Explanation | np.ndarray
        SHAP values.
    X : np.ndarray
        Feature matrix.
    feature : str
        Target feature name.
    feature_names : list[str]
        All column names.
    save_path : str
        Output file path.
    interaction_feature : str | None
        Optional second feature for colouring.
    """
    if feature not in feature_names:
        raise ValueError(f"Feature '{feature}' not found in feature_names.")

    j = feature_names.index(feature)
    vals = _extract_values_array(shap_values)

    fig, ax = plt.subplots(figsize=(8, 6))

    if interaction_feature and interaction_feature in feature_names:
        k = feature_names.index(interaction_feature)
        scatter = ax.scatter(
            X[:, j],
            vals[:, j],
            c=X[:, k],
            cmap="coolwarm",
            alpha=0.6,
            s=10,
            edgecolors="none",
        )
        cbar = fig.colorbar(scatter, ax=ax)
        cbar.set_label(interaction_feature)
    else:
        ax.scatter(X[:, j], vals[:, j], alpha=0.4, s=10, color="#1f77b4", edgecolors="none")

    ax.set_xlabel(feature)
    ax.set_ylabel(f"SHAP value for {feature}")
    ax.set_title(f"SHAP Dependence: {feature}", fontweight="bold")
    ax.axhline(0, color="grey", linewidth=0.8, linestyle="--")
    fig.tight_layout()
    _save_figure(fig, save_path)


def plot_shap_interaction(
    shap_values: Union[Any, np.ndarray],
    X: np.ndarray,
    feature1: str,
    feature2: str,
    save_path: str,
    feature_names: Optional[List[str]] = None,
) -> None:
    """Interaction plot between two features coloured by SHAP contribution.

    Parameters
    ----------
    shap_values : shap.Explanation | np.ndarray
        SHAP values.
    X : np.ndarray
        Feature matrix.
    feature1, feature2 : str
        Feature names.
    save_path : str
        Output file path.
    feature_names : list[str] | None
        Column names.
    """
    if feature_names is None:
        raise ValueError("feature_names must be provided.")

    j1 = feature_names.index(feature1)
    j2 = feature_names.index(feature2)
    vals = _extract_values_array(shap_values)

    fig, ax = plt.subplots(figsize=(8, 6))
    scatter = ax.scatter(
        X[:, j1],
        X[:, j2],
        c=vals[:, j1] + vals[:, j2],
        cmap="RdBu_r",
        alpha=0.6,
        s=12,
        edgecolors="none",
    )
    cbar = fig.colorbar(scatter, ax=ax)
    cbar.set_label("Combined SHAP value")
    ax.set_xlabel(feature1)
    ax.set_ylabel(feature2)
    ax.set_title(
        f"SHAP Interaction: {feature1} x {feature2}",
        fontweight="bold",
    )
    fig.tight_layout()
    _save_figure(fig, save_path)


# ---------------------------------------------------------------------------
# SHAP Importance Table
# ---------------------------------------------------------------------------


def get_shap_importance(
    shap_values: Union[Any, np.ndarray],
    feature_names: List[str],
) -> pd.DataFrame:
    """Compute mean |SHAP| per feature and return a ranked DataFrame.

    Parameters
    ----------
    shap_values : shap.Explanation | np.ndarray
        SHAP values.
    feature_names : list[str]
        Column names.

    Returns
    -------
    pd.DataFrame
        Columns: ``Feature``, ``mean_abs_shap``, sorted descending.
    """
    vals = _extract_values_array(shap_values)
    mean_abs = np.mean(np.abs(vals), axis=0)
    df = (
        pd.DataFrame({"Feature": feature_names, "mean_abs_shap": mean_abs})
        .sort_values("mean_abs_shap", ascending=False)
        .reset_index(drop=True)
    )
    return df


# ---------------------------------------------------------------------------
# Partial Dependence Plots (sklearn-based)
# ---------------------------------------------------------------------------


def plot_partial_dependence(
    model: Any,
    X: np.ndarray,
    features: List[str],
    feature_names: List[str],
    save_path: str,
    grid_resolution: int = 50,
) -> None:
    """Partial dependence plots for selected features using sklearn.

    Parameters
    ----------
    model : estimator
        Fitted model.
    X : np.ndarray
        Feature matrix.
    features : list[str]
        Feature names to plot (max 9 per figure).
    feature_names : list[str]
        All column names.
    save_path : str
        Output file path.
    grid_resolution : int
        Number of grid points.
    """
    indices = [feature_names.index(f) for f in features if f in feature_names]
    if not indices:
        logger.warning("No valid features found for partial dependence plot.")
        return

    X_df = _to_dataframe(X, feature_names)

    n_feats = len(indices)
    ncols = min(3, n_feats)
    nrows = int(np.ceil(n_feats / ncols))

    fig, axes = plt.subplots(
        nrows,
        ncols,
        figsize=(5 * ncols, 4 * nrows),
        constrained_layout=True,
    )
    if n_feats == 1:
        axes = np.array([axes])
    axes = np.atleast_2d(axes)

    for ax_idx, feat_idx in enumerate(indices):
        row, col = divmod(ax_idx, ncols)
        ax = axes[row, col]

        try:
            pd_result = partial_dependence(
                model,
                X_df,
                features=[feat_idx],
                grid_resolution=grid_resolution,
                kind="average",
            )
            grid_vals = pd_result["grid_values"][0]
            avg_pred = pd_result["average"][0]

            ax.plot(grid_vals, avg_pred, color="#1f77b4", linewidth=2)
            ax.set_xlabel(feature_names[feat_idx])
            ax.set_ylabel("Partial dependence")
            ax.set_title(feature_names[feat_idx], fontweight="bold")
            ax.grid(True, alpha=0.3)
        except Exception as exc:
            logger.warning("PDP failed for feature %s: %s", feature_names[feat_idx], exc)
            ax.set_visible(False)

    # Hide unused axes.
    for ax_idx in range(n_feats, nrows * ncols):
        row, col = divmod(ax_idx, ncols)
        axes[row, col].set_visible(False)

    fig.suptitle("Partial Dependence Plots", fontweight="bold", fontsize=14)
    _save_figure(fig, save_path)


# ---------------------------------------------------------------------------
# Accumulated Local Effects (ALE) — manual implementation
# ---------------------------------------------------------------------------


def plot_ale(
    model: Any,
    X: np.ndarray,
    feature: str,
    feature_names: List[str],
    save_path: str,
    n_bins: int = 40,
) -> None:
    """Accumulated Local Effects plot for a single feature.

    ALE is computed manually following Apley & Zhu (2020): the feature range
    is binned, and the average finite-difference in predictions within each
    bin is accumulated.

    Parameters
    ----------
    model : estimator
        Fitted model.
    X : np.ndarray
        Feature matrix.
    feature : str
        Target feature name.
    feature_names : list[str]
        All column names.
    save_path : str
        Output file path.
    n_bins : int
        Number of quantile bins.
    """
    j = feature_names.index(feature)
    x_feat = X[:, j].copy()

    # Quantile bin edges.
    percentiles = np.linspace(0, 100, n_bins + 1)
    edges = np.unique(np.percentile(x_feat, percentiles))
    if len(edges) < 3:
        logger.warning(
            "Feature '%s' has too few unique values for ALE — skipping.",
            feature,
        )
        return

    n_actual_bins = len(edges) - 1
    ale_values = np.zeros(n_actual_bins)
    bin_counts = np.zeros(n_actual_bins)

    for b in range(n_actual_bins):
        lower, upper = edges[b], edges[b + 1]
        if b < n_actual_bins - 1:
            mask = (x_feat >= lower) & (x_feat < upper)
        else:
            mask = (x_feat >= lower) & (x_feat <= upper)

        if mask.sum() == 0:
            continue

        X_low = X[mask].copy()
        X_high = X[mask].copy()
        X_low[:, j] = lower
        X_high[:, j] = upper

        pred_diff = model.predict(X_high) - model.predict(X_low)
        ale_values[b] = float(np.mean(pred_diff))
        bin_counts[b] = mask.sum()

    # Accumulate.
    ale_cumulative = np.cumsum(ale_values)
    # Centre so that the weighted mean is zero.
    total = bin_counts.sum()
    if total > 0:
        ale_cumulative -= np.sum(ale_cumulative * bin_counts) / total

    bin_centres = (edges[:-1] + edges[1:]) / 2.0

    fig, ax = plt.subplots(figsize=(8, 5))
    ax.plot(bin_centres, ale_cumulative, color="#d62728", linewidth=2)
    ax.fill_between(bin_centres, ale_cumulative, alpha=0.15, color="#d62728")
    ax.axhline(0, color="grey", linewidth=0.8, linestyle="--")
    ax.set_xlabel(feature)
    ax.set_ylabel("ALE")
    ax.set_title(f"Accumulated Local Effects: {feature}", fontweight="bold")
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    _save_figure(fig, save_path)


# ---------------------------------------------------------------------------
# Hedonic Premium Interpretation
# ---------------------------------------------------------------------------


def interpret_hedonic_premiums(
    shap_values: Union[Any, np.ndarray],
    X: np.ndarray,
    feature_names: List[str],
) -> pd.DataFrame:
    """Convert SHAP values into interpretable hedonic premiums.

    For each feature the function reports:

    * **mean_shap** -- average SHAP contribution (signed; positive = premium,
      negative = discount).
    * **mean_abs_shap** -- average absolute contribution (importance).
    * **std_shap** -- variability of the SHAP contribution across listings.
    * **median_shap** -- median contribution.
    * **premium_pct** -- percentage of listings for which the feature adds
      a positive premium.
    * **feature_mean** -- mean raw feature value (for context).
    * **feature_std** -- standard deviation of the raw feature.

    Parameters
    ----------
    shap_values : shap.Explanation | np.ndarray
        SHAP values.
    X : np.ndarray
        Feature matrix.
    feature_names : list[str]
        Column names.

    Returns
    -------
    pd.DataFrame
        One row per feature, sorted by ``mean_abs_shap`` descending.
    """
    vals = _extract_values_array(shap_values)
    records: List[Dict[str, Any]] = []

    for j, fname in enumerate(feature_names):
        sv = vals[:, j]
        records.append(
            {
                "Feature": fname,
                "mean_shap": float(np.mean(sv)),
                "mean_abs_shap": float(np.mean(np.abs(sv))),
                "std_shap": float(np.std(sv)),
                "median_shap": float(np.median(sv)),
                "premium_pct": float(np.mean(sv > 0) * 100),
                "feature_mean": float(np.mean(X[:, j])),
                "feature_std": float(np.std(X[:, j])),
            }
        )

    df = (
        pd.DataFrame(records)
        .sort_values("mean_abs_shap", ascending=False)
        .reset_index(drop=True)
    )
    logger.info("Hedonic premiums computed for %d features.", len(feature_names))
    return df


# ---------------------------------------------------------------------------
# Single-Listing Explanation
# ---------------------------------------------------------------------------


def explain_single_listing(
    model: Any,
    shap_values: Union[Any, np.ndarray],
    X: np.ndarray,
    idx: int,
    feature_names: List[str],
    save_path: Optional[str] = None,
    top_n: int = 15,
) -> pd.DataFrame:
    """Explain the predicted price of a single listing via a waterfall chart.

    Parameters
    ----------
    model : estimator
        Fitted model (used only to obtain the base prediction).
    shap_values : shap.Explanation | np.ndarray
        SHAP values.
    X : np.ndarray
        Feature matrix.
    idx : int
        Row index of the listing to explain.
    feature_names : list[str]
        Column names.
    save_path : str | None
        If provided, save the waterfall plot.
    top_n : int
        Number of top features to display.

    Returns
    -------
    pd.DataFrame
        Per-feature SHAP contributions for the listing, sorted by absolute
        magnitude.
    """
    vals = _extract_values_array(shap_values)
    sv_row = vals[idx]

    explanation_df = (
        pd.DataFrame(
            {
                "Feature": feature_names,
                "Feature_Value": X[idx],
                "SHAP_Value": sv_row,
            }
        )
        .assign(abs_shap=lambda d: np.abs(d["SHAP_Value"]))
        .sort_values("abs_shap", ascending=False)
        .drop(columns="abs_shap")
        .reset_index(drop=True)
    )

    predicted = float(model.predict(X[idx : idx + 1])[0])
    base_value: float = predicted - float(np.sum(sv_row))
    logger.info(
        "Listing %d: predicted=%.2f  base=%.2f  sum(SHAP)=%.2f",
        idx,
        predicted,
        base_value,
        np.sum(sv_row),
    )

    # Waterfall / bar chart.
    if _SHAP_AVAILABLE and hasattr(shap_values, "values"):
        try:
            fig = plt.figure(figsize=(10, 7))
            single_exp = shap_values[idx]
            shap.waterfall_plot(single_exp, max_display=top_n, show=False)
            fig = plt.gcf()
            fig.suptitle(
                f"Listing {idx} Explanation (predicted=${predicted:,.0f})",
                y=1.02,
                fontweight="bold",
            )
            if save_path:
                _save_figure(fig, save_path)
            else:
                plt.close(fig)
            return explanation_df
        except Exception:
            logger.warning("SHAP waterfall_plot failed — using manual bar chart.")

    # Manual waterfall-style horizontal bar chart.
    top = explanation_df.head(top_n)
    fig, ax = plt.subplots(figsize=(10, max(6, top_n * 0.4)))
    colours = ["#d62728" if v < 0 else "#2ca02c" for v in top["SHAP_Value"]]
    y_pos = range(len(top) - 1, -1, -1)
    ax.barh(
        y_pos,
        top["SHAP_Value"].values,
        color=colours,
        edgecolor="white",
        height=0.7,
    )
    labels = [
        f"{row.Feature} = {row.Feature_Value:.2g}"
        for _, row in top.iterrows()
    ]
    ax.set_yticks(list(y_pos))
    ax.set_yticklabels(labels)
    ax.axvline(0, color="grey", linewidth=0.8)
    ax.set_xlabel("SHAP value (impact on prediction)")
    ax.set_title(
        f"Listing {idx} Explanation  |  Base: ${base_value:,.0f}  |  Predicted: ${predicted:,.0f}",
        fontweight="bold",
    )
    fig.tight_layout()

    if save_path:
        _save_figure(fig, save_path)
    else:
        plt.close(fig)

    return explanation_df