"""
figures.py - Publication-quality figure generation for hedonic pricing analysis.

All figures use matplotlib + seaborn with a consistent 'whitegrid' style,
(10, 6) default size, 300 DPI, and are saved as both PNG and PDF for
LaTeX / journal submission compatibility.

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

from __future__ import annotations

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

import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import numpy as np
import pandas as pd
import seaborn as sns

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Global style defaults
# ---------------------------------------------------------------------------

_DEFAULT_FIGSIZE: Tuple[int, int] = (10, 6)
_DEFAULT_DPI: int = 300
_STYLE: str = "whitegrid"
_PALETTE: str = "muted"

# Publication-quality font sizes
_RC_PARAMS: Dict[str, Any] = {
    "font.size": 11,
    "axes.titlesize": 13,
    "axes.labelsize": 12,
    "xtick.labelsize": 10,
    "ytick.labelsize": 10,
    "legend.fontsize": 10,
    "figure.titlesize": 14,
    "font.family": "serif",
}


def _apply_style() -> None:
    """Apply consistent seaborn + matplotlib style."""
    sns.set_style(_STYLE)
    sns.set_palette(_PALETTE)
    plt.rcParams.update(_RC_PARAMS)


def _save_figure(
    fig: plt.Figure,
    filename: str,
    save_dir: str,
) -> List[str]:
    """Save a figure as PNG and PDF, return list of saved paths."""
    os.makedirs(save_dir, exist_ok=True)
    paths: List[str] = []
    for ext in ("png", "pdf"):
        stem = Path(filename).stem
        fpath = os.path.join(save_dir, f"{stem}.{ext}")
        fig.savefig(fpath, dpi=_DEFAULT_DPI, bbox_inches="tight",
                    facecolor="white", edgecolor="none")
        paths.append(fpath)
        logger.info("Saved figure: %s", fpath)
    plt.close(fig)
    return paths


# ---------------------------------------------------------------------------
# Figure 01 - Price distribution
# ---------------------------------------------------------------------------

def plot_price_distribution(
    df: pd.DataFrame,
    save_dir: str,
    price_col: str = "price_numeric",
    figsize: Tuple[int, int] = _DEFAULT_FIGSIZE,
) -> List[str]:
    """
    01_price_distribution.png -- Histogram + KDE of raw prices.

    Parameters
    ----------
    df : pd.DataFrame
    save_dir : str
    price_col : str
    figsize : tuple

    Returns
    -------
    list[str]
        Paths to saved files.
    """
    _apply_style()
    prices = df[price_col].dropna()

    fig, ax = plt.subplots(figsize=figsize)
    sns.histplot(prices, bins=80, kde=True, color="#4C72B0",
                 edgecolor="white", linewidth=0.5, ax=ax)
    ax.set_xlabel("Nightly Price (CAD)")
    ax.set_ylabel("Frequency")
    ax.set_title("Distribution of Airbnb Listing Prices")

    # Annotate key statistics
    med = prices.median()
    mean = prices.mean()
    ax.axvline(med, color="#C44E52", linestyle="--", linewidth=1.2,
               label=f"Median = ${med:,.0f}")
    ax.axvline(mean, color="#DD8452", linestyle="--", linewidth=1.2,
               label=f"Mean = ${mean:,.0f}")
    ax.legend(frameon=True, fancybox=False, edgecolor="0.8")

    ax.xaxis.set_major_formatter(mticker.StrMethodFormatter("${x:,.0f}"))
    fig.tight_layout()
    return _save_figure(fig, "01_price_distribution", save_dir)


# ---------------------------------------------------------------------------
# Figure 02 - Log-price distribution
# ---------------------------------------------------------------------------

def plot_log_price_distribution(
    df: pd.DataFrame,
    save_dir: str,
    price_col: str = "price_numeric",
    figsize: Tuple[int, int] = _DEFAULT_FIGSIZE,
) -> List[str]:
    """02_log_price_distribution.png -- Histogram + KDE of ln(price)."""
    _apply_style()
    prices = df[price_col].dropna()
    log_prices = np.log(prices[prices > 0])

    fig, ax = plt.subplots(figsize=figsize)
    sns.histplot(log_prices, bins=60, kde=True, color="#55A868",
                 edgecolor="white", linewidth=0.5, ax=ax)
    ax.set_xlabel("ln(Nightly Price)")
    ax.set_ylabel("Frequency")
    ax.set_title("Distribution of Log-Transformed Prices\n"
                 "(Semi-log hedonic model dependent variable)")

    med = log_prices.median()
    ax.axvline(med, color="#C44E52", linestyle="--", linewidth=1.2,
               label=f"Median ln(price) = {med:.2f}")
    ax.legend(frameon=True, fancybox=False, edgecolor="0.8")

    # Shapiro-Wilk note (sample if too large)
    from scipy import stats as sp_stats
    sample = log_prices.sample(min(5000, len(log_prices)), random_state=42)
    _, sw_p = sp_stats.shapiro(sample)
    ax.annotate(f"Shapiro-Wilk p = {sw_p:.4f} (n={len(sample)})",
                xy=(0.98, 0.95), xycoords="axes fraction",
                ha="right", va="top", fontsize=9,
                bbox=dict(boxstyle="round,pad=0.3", fc="white",
                          ec="0.8", alpha=0.9))

    fig.tight_layout()
    return _save_figure(fig, "02_log_price_distribution", save_dir)


# ---------------------------------------------------------------------------
# Figure 03 - Price map
# ---------------------------------------------------------------------------

def plot_price_map(
    df: pd.DataFrame,
    save_dir: str,
    price_col: str = "price_numeric",
    lat_col: str = "lat",
    lon_col: str = "long",
    figsize: Tuple[int, int] = (12, 10),
) -> List[str]:
    """03_price_map.png -- Scatter on lat/long coloured by price."""
    _apply_style()
    sub = df.dropna(subset=[price_col, lat_col, lon_col]).copy()

    # Cap at 99th percentile for colour scale
    p99 = sub[price_col].quantile(0.99)
    sub["price_capped"] = sub[price_col].clip(upper=p99)

    fig, ax = plt.subplots(figsize=figsize)
    scatter = ax.scatter(
        sub[lon_col], sub[lat_col],
        c=sub["price_capped"], cmap="RdYlGn_r",
        s=8, alpha=0.6, edgecolors="none",
    )
    cbar = fig.colorbar(scatter, ax=ax, pad=0.02, shrink=0.8)
    cbar.set_label("Nightly Price (CAD)", fontsize=11)

    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    ax.set_title("Spatial Distribution of Airbnb Prices -- Quebec")

    fig.tight_layout()
    return _save_figure(fig, "03_price_map", save_dir)


# ---------------------------------------------------------------------------
# Figure 04 - Price by neighbourhood / city
# ---------------------------------------------------------------------------

def plot_price_by_neighbourhood(
    df: pd.DataFrame,
    save_dir: str,
    price_col: str = "price_numeric",
    geo_col: str = "city",
    top_n: int = 20,
    figsize: Tuple[int, int] = (12, 7),
) -> List[str]:
    """04_price_by_neighbourhood.png -- Top 20 cities/neighbourhoods."""
    _apply_style()
    sub = df.dropna(subset=[price_col, geo_col]).copy()

    # Top N by listing count
    top_areas = (
        sub[geo_col].value_counts().head(top_n).index.tolist()
    )
    sub = sub[sub[geo_col].isin(top_areas)]

    # Order by median price
    order = (
        sub.groupby(geo_col, observed=True)[price_col]
        .median().sort_values(ascending=False).index.tolist()
    )

    fig, ax = plt.subplots(figsize=figsize)
    sns.boxplot(
        data=sub, y=geo_col, x=price_col, order=order,
        palette="coolwarm", showfliers=False, ax=ax,
        linewidth=0.8,
    )
    ax.set_xlabel("Nightly Price (CAD)")
    ax.set_ylabel("")
    ax.set_title(f"Price Distribution by Top {top_n} Cities / Neighbourhoods")
    ax.xaxis.set_major_formatter(mticker.StrMethodFormatter("${x:,.0f}"))

    fig.tight_layout()
    return _save_figure(fig, "04_price_by_neighbourhood", save_dir)


# ---------------------------------------------------------------------------
# Figure 05 - Price by room type
# ---------------------------------------------------------------------------

def plot_price_by_room_type(
    df: pd.DataFrame,
    save_dir: str,
    price_col: str = "price_numeric",
    room_col: str = "property_type",
    figsize: Tuple[int, int] = _DEFAULT_FIGSIZE,
) -> List[str]:
    """05_price_by_room_type.png -- Violin + strip plot by room/property type."""
    _apply_style()
    sub = df.dropna(subset=[price_col, room_col]).copy()

    # Keep types with enough observations
    counts = sub[room_col].value_counts()
    valid_types = counts[counts >= 10].index.tolist()
    sub = sub[sub[room_col].isin(valid_types)]

    order = (
        sub.groupby(room_col, observed=True)[price_col]
        .median().sort_values(ascending=False).index.tolist()
    )

    fig, ax = plt.subplots(figsize=figsize)
    sns.violinplot(
        data=sub, y=room_col, x=price_col, order=order,
        palette="Set2", inner="quartile", cut=0, ax=ax,
        linewidth=0.8,
    )
    ax.set_xlabel("Nightly Price (CAD)")
    ax.set_ylabel("")
    ax.set_title("Price Distribution by Property Type")
    ax.xaxis.set_major_formatter(mticker.StrMethodFormatter("${x:,.0f}"))

    fig.tight_layout()
    return _save_figure(fig, "05_price_by_room_type", save_dir)


# ---------------------------------------------------------------------------
# Figure 06 - Amenity premiums
# ---------------------------------------------------------------------------

def plot_amenity_premiums(
    premiums_df: pd.DataFrame,
    save_dir: str,
    top_n: int = 20,
    figsize: Tuple[int, int] = (10, 8),
) -> List[str]:
    """06_amenity_premiums.png -- Horizontal bar chart of % premiums."""
    _apply_style()
    df = premiums_df.copy()

    # Normalise column names
    if "% Premium" not in df.columns:
        for c in df.columns:
            if "premium" in c.lower() or "effect" in c.lower():
                df = df.rename(columns={c: "% Premium"})
                break
    if "Amenity" not in df.columns:
        for c in df.columns:
            if "variable" in c.lower() or "amenity" in c.lower() or "feature" in c.lower():
                df = df.rename(columns={c: "Amenity"})
                break

    df = df.sort_values("% Premium", ascending=True).tail(top_n)

    fig, ax = plt.subplots(figsize=figsize)
    colours = ["#C44E52" if v < 0 else "#55A868" for v in df["% Premium"]]
    ax.barh(df["Amenity"], df["% Premium"], color=colours, edgecolor="white",
            linewidth=0.5)
    ax.set_xlabel("Price Premium (%)")
    ax.set_ylabel("")
    ax.set_title("Estimated Amenity Price Premiums\n(from semi-log hedonic model)")
    ax.axvline(0, color="black", linewidth=0.8)

    # Value labels
    for i, (val, name) in enumerate(zip(df["% Premium"], df["Amenity"])):
        ax.text(val + (0.3 if val >= 0 else -0.3), i,
                f"{val:+.1f}%", va="center",
                ha="left" if val >= 0 else "right", fontsize=9)

    fig.tight_layout()
    return _save_figure(fig, "06_amenity_premiums", save_dir)


# ---------------------------------------------------------------------------
# Figure 07 - Review score premium
# ---------------------------------------------------------------------------

def plot_review_score_premium(
    df: pd.DataFrame,
    save_dir: str,
    price_col: str = "price_numeric",
    rating_col: str = "rating",
    figsize: Tuple[int, int] = _DEFAULT_FIGSIZE,
) -> List[str]:
    """07_review_score_premium.png -- Price vs. review score."""
    _apply_style()
    sub = df.dropna(subset=[price_col, rating_col]).copy()
    sub = sub[sub[price_col] > 0]

    fig, axes = plt.subplots(1, 2, figsize=(14, 6))

    # Left: scatter + regression line
    ax = axes[0]
    sns.regplot(
        data=sub, x=rating_col, y=price_col,
        scatter_kws={"alpha": 0.15, "s": 10, "color": "#4C72B0"},
        line_kws={"color": "#C44E52", "linewidth": 1.5},
        ax=ax, ci=95,
    )
    ax.set_xlabel("Review Rating")
    ax.set_ylabel("Nightly Price (CAD)")
    ax.set_title("Price vs. Review Score")
    ax.yaxis.set_major_formatter(mticker.StrMethodFormatter("${x:,.0f}"))

    # Right: binned means with error bars
    ax = axes[1]
    sub["rating_bin"] = pd.cut(sub[rating_col], bins=10)
    binned = sub.groupby("rating_bin", observed=True)[price_col].agg(
        ["mean", "sem", "count"]
    ).dropna()
    binned["mid"] = [interval.mid for interval in binned.index]
    ax.errorbar(binned["mid"], binned["mean"], yerr=1.96 * binned["sem"],
                fmt="o-", color="#4C72B0", capsize=3, markersize=5)
    ax.set_xlabel("Review Rating (binned)")
    ax.set_ylabel("Mean Price (CAD)")
    ax.set_title("Mean Price by Rating Bin (95% CI)")
    ax.yaxis.set_major_formatter(mticker.StrMethodFormatter("${x:,.0f}"))

    fig.suptitle("Review Score and Price Relationship", y=1.02, fontsize=14)
    fig.tight_layout()
    return _save_figure(fig, "07_review_score_premium", save_dir)


# ---------------------------------------------------------------------------
# Figure 08 - Distance gradient
# ---------------------------------------------------------------------------

def plot_distance_gradient(
    df: pd.DataFrame,
    save_dir: str,
    price_col: str = "price_numeric",
    distance_col: Optional[str] = None,
    lat_col: str = "lat",
    lon_col: str = "long",
    center_lat: Optional[float] = None,
    center_lon: Optional[float] = None,
    figsize: Tuple[int, int] = _DEFAULT_FIGSIZE,
) -> List[str]:
    """
    08_distance_to_center_gradient.png -- Scatter + LOWESS of price vs
    distance to city centre.

    If no distance column exists, one is computed from lat/long using a
    Haversine approximation to the centroid (or a supplied centre).
    """
    _apply_style()
    sub = df.dropna(subset=[price_col, lat_col, lon_col]).copy()
    sub = sub[sub[price_col] > 0]

    if distance_col and distance_col in sub.columns:
        dist = sub[distance_col]
    else:
        # Compute distance to centroid (Haversine)
        if center_lat is None:
            center_lat = sub[lat_col].median()
        if center_lon is None:
            center_lon = sub[lon_col].median()

        def _haversine_km(lat1: float, lon1: float,
                          lat2: float, lon2: float) -> float:
            R = 6371.0
            dlat = np.radians(lat2 - lat1)
            dlon = np.radians(lon2 - lon1)
            a = (np.sin(dlat / 2) ** 2 +
                 np.cos(np.radians(lat1)) * np.cos(np.radians(lat2)) *
                 np.sin(dlon / 2) ** 2)
            return R * 2 * np.arctan2(np.sqrt(a), np.sqrt(1 - a))

        dist = sub.apply(
            lambda r: _haversine_km(r[lat_col], r[lon_col],
                                    center_lat, center_lon),
            axis=1,
        )
        distance_col = "dist_to_center_km"
        sub[distance_col] = dist

    fig, ax = plt.subplots(figsize=figsize)
    ax.scatter(sub[distance_col], sub[price_col],
               alpha=0.15, s=8, color="#4C72B0", edgecolors="none")

    # LOWESS smoothing
    try:
        from statsmodels.nonparametric.smoothers_lowess import lowess
        sample = sub.sample(min(5000, len(sub)), random_state=42)
        lowess_result = lowess(
            sample[price_col].values,
            sample[distance_col].values,
            frac=0.3,
        )
        ax.plot(lowess_result[:, 0], lowess_result[:, 1],
                color="#C44E52", linewidth=2.5, label="LOWESS")
        ax.legend(frameon=True, fancybox=False, edgecolor="0.8")
    except ImportError:
        logger.warning("statsmodels not available; LOWESS line omitted.")

    ax.set_xlabel("Distance to Centre (km)")
    ax.set_ylabel("Nightly Price (CAD)")
    ax.set_title("Price Gradient: Distance to City Centre\n"
                 "(negative gradient expected from urban economics theory)")
    ax.yaxis.set_major_formatter(mticker.StrMethodFormatter("${x:,.0f}"))

    fig.tight_layout()
    return _save_figure(fig, "08_distance_to_center_gradient", save_dir)


# ---------------------------------------------------------------------------
# Figure 09 - Quantile coefficients
# ---------------------------------------------------------------------------

def plot_quantile_coefficients(
    coef_df: pd.DataFrame,
    save_dir: str,
    variables: Optional[List[str]] = None,
    figsize: Tuple[int, int] = (12, 8),
) -> List[str]:
    """
    09_quantile_coefficients.png -- Coefficient evolution across quantiles.

    Parameters
    ----------
    coef_df : pd.DataFrame
        Must contain columns 'quantile', 'variable', 'coef', and optionally
        'ci_lower', 'ci_upper'.
    """
    _apply_style()

    if variables is None:
        variables = coef_df["variable"].unique().tolist()[:6]

    n_vars = len(variables)
    ncols = min(3, n_vars)
    nrows = int(np.ceil(n_vars / ncols))

    fig, axes = plt.subplots(nrows, ncols, figsize=figsize, squeeze=False)

    for idx, var in enumerate(variables):
        ax = axes[idx // ncols, idx % ncols]
        var_data = coef_df[coef_df["variable"] == var].sort_values("quantile")

        ax.plot(var_data["quantile"], var_data["coef"],
                "o-", color="#4C72B0", markersize=5, linewidth=1.5)

        if "ci_lower" in var_data.columns and "ci_upper" in var_data.columns:
            ax.fill_between(
                var_data["quantile"],
                var_data["ci_lower"],
                var_data["ci_upper"],
                alpha=0.2, color="#4C72B0",
            )

        ax.axhline(0, color="black", linewidth=0.5, linestyle="--")
        ax.set_title(var, fontsize=10)
        ax.set_xlabel("Quantile")
        ax.set_ylabel("Coefficient")

    # Hide unused axes
    for idx in range(n_vars, nrows * ncols):
        axes[idx // ncols, idx % ncols].set_visible(False)

    fig.suptitle("Quantile Regression Coefficients Across the Price Distribution",
                 y=1.02, fontsize=13)
    fig.tight_layout()
    return _save_figure(fig, "09_quantile_coefficients", save_dir)


# ---------------------------------------------------------------------------
# Figure 14 - Predicted vs actual
# ---------------------------------------------------------------------------

def plot_predicted_vs_actual(
    y_true: Union[np.ndarray, pd.Series],
    y_pred: Union[np.ndarray, pd.Series],
    save_dir: str,
    model_name: str = "Model",
    figsize: Tuple[int, int] = _DEFAULT_FIGSIZE,
) -> List[str]:
    """14_predicted_vs_actual.png -- 45-degree line plot."""
    _apply_style()

    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)

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

    ax.scatter(y_true, y_pred, alpha=0.2, s=10, color="#4C72B0",
               edgecolors="none")

    # 45-degree line
    lims = [
        min(y_true.min(), y_pred.min()),
        max(y_true.max(), y_pred.max()),
    ]
    ax.plot(lims, lims, "--", color="#C44E52", linewidth=1.5,
            label="Perfect prediction")

    # Metrics
    from sklearn.metrics import r2_score, mean_absolute_error, \
        root_mean_squared_error
    r2 = r2_score(y_true, y_pred)
    mae = mean_absolute_error(y_true, y_pred)
    rmse = root_mean_squared_error(y_true, y_pred)

    stats_text = f"$R^2$ = {r2:.4f}\nMAE = {mae:.2f}\nRMSE = {rmse:.2f}"
    ax.annotate(stats_text, xy=(0.05, 0.92), xycoords="axes fraction",
                fontsize=10, va="top",
                bbox=dict(boxstyle="round,pad=0.4", fc="white",
                          ec="0.8", alpha=0.9))

    ax.set_xlabel("Actual Values")
    ax.set_ylabel("Predicted Values")
    ax.set_title(f"Predicted vs. Actual -- {model_name}")
    ax.legend(loc="lower right", frameon=True, fancybox=False, edgecolor="0.8")

    fig.tight_layout()
    return _save_figure(fig, "14_predicted_vs_actual", save_dir)


# ---------------------------------------------------------------------------
# Figure 15 - Hedonic decomposition (stacked bar)
# ---------------------------------------------------------------------------

def plot_hedonic_decomposition(
    decomp_df: pd.DataFrame,
    save_dir: str,
    figsize: Tuple[int, int] = (12, 7),
) -> List[str]:
    """
    15_hedonic_decomposition_example.png -- Stacked bar showing how
    predicted price is decomposed into attribute contributions.

    Parameters
    ----------
    decomp_df : pd.DataFrame
        Rows = individual listings (or a small sample), columns = attribute
        categories (e.g., 'Location', 'Size', 'Amenities', 'Host',
        'Intercept'). Values = contribution to predicted ln(price) or price.
    """
    _apply_style()

    df = decomp_df.copy()
    if len(df) > 15:
        df = df.head(15)

    # Ensure listing labels
    if df.index.dtype == "int64":
        df.index = [f"Listing {i+1}" for i in range(len(df))]

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

    # Build stacked bars (positive contributions only for clarity)
    bottom = np.zeros(len(df))
    colours = sns.color_palette("Set2", n_colors=len(df.columns))

    for i, col in enumerate(df.columns):
        vals = df[col].values.astype(float)
        ax.barh(df.index, vals, left=bottom, color=colours[i],
                edgecolor="white", linewidth=0.5, label=col)
        bottom += np.maximum(vals, 0)

    ax.set_xlabel("Contribution to Predicted Price")
    ax.set_ylabel("")
    ax.set_title("Hedonic Price Decomposition\n"
                 "(marginal contribution of each attribute category)")
    ax.legend(bbox_to_anchor=(1.01, 1), loc="upper left",
              frameon=True, fancybox=False, edgecolor="0.8")

    fig.tight_layout()
    return _save_figure(fig, "15_hedonic_decomposition_example", save_dir)


# ---------------------------------------------------------------------------
# Master function
# ---------------------------------------------------------------------------

def generate_all_figures(
    df: pd.DataFrame,
    results: Optional[Dict[str, Any]] = None,
    save_dir: Optional[str] = None,
    price_col: str = "price_numeric",
) -> Dict[str, List[str]]:
    """
    Generate all figures that can be produced from the dataset alone
    (figures requiring model results are skipped if *results* is None).

    Parameters
    ----------
    df : pd.DataFrame
        Full analysis dataset.
    results : dict, optional
        ``{'premiums_df': pd.DataFrame, 'coef_df': pd.DataFrame,
        'y_true': array, 'y_pred': array, 'decomp_df': pd.DataFrame,
        'model_name': str}``.
    save_dir : str, optional
        Output directory; defaults to ``results/figures/``.

    Returns
    -------
    dict
        ``{figure_name: [saved_paths]}``.
    """
    if save_dir is None:
        save_dir = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            "results", "figures",
        )

    if results is None:
        results = {}

    output: Dict[str, List[str]] = {}
    logger.info("Generating all figures to %s", save_dir)

    # Data-only figures
    try:
        output["01_price_distribution"] = plot_price_distribution(
            df, save_dir, price_col=price_col)
    except Exception as e:
        logger.error("Failed 01_price_distribution: %s", e)

    try:
        output["02_log_price_distribution"] = plot_log_price_distribution(
            df, save_dir, price_col=price_col)
    except Exception as e:
        logger.error("Failed 02_log_price_distribution: %s", e)

    try:
        output["03_price_map"] = plot_price_map(
            df, save_dir, price_col=price_col)
    except Exception as e:
        logger.error("Failed 03_price_map: %s", e)

    try:
        output["04_price_by_neighbourhood"] = plot_price_by_neighbourhood(
            df, save_dir, price_col=price_col)
    except Exception as e:
        logger.error("Failed 04_price_by_neighbourhood: %s", e)

    try:
        output["05_price_by_room_type"] = plot_price_by_room_type(
            df, save_dir, price_col=price_col)
    except Exception as e:
        logger.error("Failed 05_price_by_room_type: %s", e)

    # Review score premium (data-only)
    try:
        output["07_review_score_premium"] = plot_review_score_premium(
            df, save_dir, price_col=price_col)
    except Exception as e:
        logger.error("Failed 07_review_score_premium: %s", e)

    # Distance gradient (data-only)
    try:
        output["08_distance_gradient"] = plot_distance_gradient(
            df, save_dir, price_col=price_col)
    except Exception as e:
        logger.error("Failed 08_distance_gradient: %s", e)

    # Model-dependent figures
    if "premiums_df" in results:
        try:
            output["06_amenity_premiums"] = plot_amenity_premiums(
                results["premiums_df"], save_dir)
        except Exception as e:
            logger.error("Failed 06_amenity_premiums: %s", e)

    if "coef_df" in results:
        try:
            output["09_quantile_coefficients"] = plot_quantile_coefficients(
                results["coef_df"], save_dir)
        except Exception as e:
            logger.error("Failed 09_quantile_coefficients: %s", e)

    if "y_true" in results and "y_pred" in results:
        try:
            output["14_predicted_vs_actual"] = plot_predicted_vs_actual(
                results["y_true"], results["y_pred"], save_dir,
                model_name=results.get("model_name", "Model"))
        except Exception as e:
            logger.error("Failed 14_predicted_vs_actual: %s", e)

    if "decomp_df" in results:
        try:
            output["15_hedonic_decomposition"] = plot_hedonic_decomposition(
                results["decomp_df"], save_dir)
        except Exception as e:
            logger.error("Failed 15_hedonic_decomposition: %s", e)

    logger.info("Figure generation complete: %d/%d succeeded.",
                len(output), 11)
    return output