"""
Visualization functions for directrs.

All functions return matplotlib.Figure objects.
The caller is responsible for saving or displaying them.
"""

from __future__ import annotations

import numpy as np
import matplotlib
matplotlib.use("Agg")  # non-interactive backend
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.figure import Figure
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix


def _feature_labels(feature_names: list[str] | None, n: int) -> list[str]:
    if feature_names is not None and len(feature_names) == n:
        return list(feature_names)
    return [f"x{i}" for i in range(n)]


def plot_g_matrix(
    G: np.ndarray,
    feature_names: list[str] | None = None,
) -> Figure:
    """Heatmap of the G geometry matrix (full + off-diagonal highlighted)."""
    labels = _feature_labels(feature_names, G.shape[0])
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))

    # Full G
    sns.heatmap(
        G, ax=axes[0], xticklabels=labels, yticklabels=labels,
        cmap="viridis", annot=len(labels) <= 15, fmt=".2f",
    )
    axes[0].set_title("G matrix (full)")

    # Off-diagonal (interactions)
    G_off = G.copy()
    np.fill_diagonal(G_off, 0.0)
    sns.heatmap(
        G_off, ax=axes[1], xticklabels=labels, yticklabels=labels,
        cmap="viridis", annot=len(labels) <= 15, fmt=".2f",
    )
    axes[1].set_title("G matrix (off-diagonal interactions)")

    fig.suptitle("Geometry Matrix G", fontsize=14, fontweight="bold")
    fig.tight_layout()
    return fig


def plot_stretch_spectrum(
    S_prime: np.ndarray,
    feature_names: list[str] | None = None,
) -> Figure:
    """Eigenvalue spectrum of S' and per-feature diagonal values."""
    labels = _feature_labels(feature_names, S_prime.shape[0])
    eigenvalues = np.linalg.eigvalsh(S_prime)[::-1]  # descending
    diag_vals = np.diag(S_prime)

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

    # Eigenvalue spectrum
    axes[0].bar(range(len(eigenvalues)), eigenvalues, color="steelblue")
    axes[0].set_xlabel("Component index")
    axes[0].set_ylabel("Eigenvalue")
    axes[0].set_title("S' eigenvalue spectrum")
    axes[0].set_xticks(range(len(eigenvalues)))

    # Per-feature diagonal
    axes[1].bar(labels, diag_vals, color="darkorange")
    axes[1].set_xlabel("Feature")
    axes[1].set_ylabel("S'[i,i]")
    axes[1].set_title("Per-feature stretch (diagonal of S')")
    axes[1].tick_params(axis="x", rotation=45)

    fig.suptitle("Stretch Matrix S' Spectrum", fontsize=14, fontweight="bold")
    fig.tight_layout()
    return fig


def plot_feature_loadings(
    S_prime: np.ndarray,
    feature_names: list[str] | None = None,
) -> Figure:
    """Heatmap of eigenvector loadings of S'."""
    labels = _feature_labels(feature_names, S_prime.shape[0])
    eigenvalues, eigenvectors = np.linalg.eigh(S_prime)
    # Sort descending
    order = np.argsort(eigenvalues)[::-1]
    eigenvectors = eigenvectors[:, order]
    eigenvalues = eigenvalues[order]

    comp_labels = [f"PC{i+1}\n(λ={eigenvalues[i]:.2f})" for i in range(len(eigenvalues))]

    fig, ax = plt.subplots(figsize=(max(8, len(labels)), max(6, len(labels) // 2)))
    sns.heatmap(
        eigenvectors.T, ax=ax,
        xticklabels=labels, yticklabels=comp_labels,
        cmap="RdBu_r", center=0, annot=len(labels) <= 12, fmt=".2f",
    )
    ax.set_title("S' Eigenvector Loadings", fontsize=13, fontweight="bold")
    ax.set_xlabel("Feature")
    ax.set_ylabel("Component")
    fig.tight_layout()
    return fig


def plot_feature_importance(
    tree,
    S_prime: np.ndarray,
    feature_names: list[str] | None = None,
    task: str = "classification",
) -> Figure:
    """Side-by-side bar charts: tree gain importance vs geometry (diag S') importance."""
    labels = _feature_labels(feature_names, S_prime.shape[0])
    tree_imp = tree.feature_importances_

    diag = np.abs(np.diag(S_prime))
    total = diag.sum()
    geom_imp = diag / total if total > 0 else diag

    fig, axes = plt.subplots(1, 2, figsize=(14, max(5, len(labels) // 2 + 3)))

    order_tree = np.argsort(tree_imp)[::-1]
    order_geom = np.argsort(geom_imp)[::-1]

    axes[0].barh(
        [labels[i] for i in order_tree[::-1]],
        tree_imp[order_tree[::-1]],
        color="steelblue",
    )
    axes[0].set_title("Boundary Importance\n(tree gain)")
    axes[0].set_xlabel("Importance")

    axes[1].barh(
        [labels[i] for i in order_geom[::-1]],
        geom_imp[order_geom[::-1]],
        color="darkorange",
    )
    axes[1].set_title("Geometry Importance\n(diagonal of S')")
    axes[1].set_xlabel("Importance")

    fig.suptitle(f"Feature Importance ({task.capitalize()})", fontsize=14, fontweight="bold")
    fig.tight_layout()
    return fig


def plot_predictions_vs_actual(
    y_true: np.ndarray,
    y_pred: np.ndarray,
    split_name: str = "Train",
    task: str = "classification",
) -> Figure:
    """
    Classification: confusion matrix.
    Regression: scatter actual vs predicted + residual histogram.
    """
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)

    if task == "classification":
        cm = confusion_matrix(y_true, y_pred)
        fig, ax = plt.subplots(figsize=(6, 5))
        disp = ConfusionMatrixDisplay(confusion_matrix=cm)
        disp.plot(ax=ax, colorbar=False)
        ax.set_title(f"Confusion Matrix — {split_name}", fontsize=13, fontweight="bold")
        fig.tight_layout()
        return fig
    else:
        residuals = y_true - y_pred
        fig, axes = plt.subplots(1, 2, figsize=(13, 5))

        # Scatter
        mn = min(y_true.min(), y_pred.min())
        mx = max(y_true.max(), y_pred.max())
        axes[0].scatter(y_true, y_pred, alpha=0.4, s=20, color="steelblue")
        axes[0].plot([mn, mx], [mn, mx], "r--", linewidth=1)
        axes[0].set_xlabel("Actual")
        axes[0].set_ylabel("Predicted")
        axes[0].set_title(f"Actual vs Predicted — {split_name}")

        # Residuals
        axes[1].hist(residuals, bins=30, color="darkorange", edgecolor="white")
        axes[1].axvline(0, color="red", linestyle="--")
        axes[1].set_xlabel("Residual")
        axes[1].set_ylabel("Count")
        axes[1].set_title(f"Residuals — {split_name}")

        fig.suptitle(f"Predictions ({split_name})", fontsize=14, fontweight="bold")
        fig.tight_layout()
        return fig


def plot_metrics_comparison(
    metrics_dict: dict[str, dict[str, float]],
    split_names: list[str],
) -> Figure:
    """
    Bar chart comparing Tree baseline vs DirectRS metrics across train/test splits.

    metrics_dict structure:
      { split_name: { metric_name: value, ... }, ... }
    """
    # Gather all metric names (excluding per-class F1 entries)
    all_metrics: set[str] = set()
    for split in split_names:
        all_metrics.update(metrics_dict[split].keys())
    metric_names = sorted(all_metrics)

    n_metrics = len(metric_names)
    n_splits = len(split_names)
    bar_width = 0.8 / n_splits
    x = np.arange(n_metrics)

    fig, ax = plt.subplots(figsize=(max(10, n_metrics * 2), 5))
    colors = plt.cm.tab10(np.linspace(0, 1, n_splits))

    for idx, (split, color) in enumerate(zip(split_names, colors)):
        values = [metrics_dict[split].get(m, float("nan")) for m in metric_names]
        offset = (idx - n_splits / 2 + 0.5) * bar_width
        bars = ax.bar(x + offset, values, width=bar_width, label=split, color=color, alpha=0.8)
        for bar, val in zip(bars, values):
            if not np.isnan(val):
                ax.text(
                    bar.get_x() + bar.get_width() / 2,
                    bar.get_height() + 0.005,
                    f"{val:.3f}",
                    ha="center", va="bottom", fontsize=8, rotation=45,
                )

    ax.set_xticks(x)
    ax.set_xticklabels(metric_names, rotation=30, ha="right")
    ax.set_ylabel("Metric value")
    ax.set_title("Metrics Comparison (Train vs Test)", fontsize=13, fontweight="bold")
    ax.legend()
    fig.tight_layout()
    return fig


def plot_embedding_scatter(
    X: np.ndarray,
    S_prime: np.ndarray,
    y: np.ndarray,
    feature_names: list[str] | None = None,
    task: str = "classification",
) -> Figure:
    """
    2-D scatter of S'x (first two components after PCA-like projection) colored by target.
    """
    labels = _feature_labels(feature_names, S_prime.shape[0])
    X = np.asarray(X, dtype=float)

    # Project with S'
    X_stretch = X @ S_prime.T  # (n, d)

    # Use first two principal components of S'x for visualization
    eigenvalues, eigenvectors = np.linalg.eigh(S_prime)
    order = np.argsort(eigenvalues)[::-1]
    eigenvectors = eigenvectors[:, order]

    if X_stretch.shape[1] >= 2:
        Z = X_stretch @ eigenvectors[:, :2]
    else:
        # Fallback: pad with zeros
        Z = np.column_stack([X_stretch[:, 0], np.zeros(len(X))])

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

    if task == "classification":
        classes = np.unique(y)
        palette = plt.cm.tab10(np.linspace(0, 1, len(classes)))
        for cls, color in zip(classes, palette):
            mask = y == cls
            ax.scatter(Z[mask, 0], Z[mask, 1], label=str(cls), alpha=0.6, s=20, color=color)
        ax.legend(title="Class", bbox_to_anchor=(1.02, 1), loc="upper left")
    else:
        sc = ax.scatter(Z[:, 0], Z[:, 1], c=y, cmap="viridis", alpha=0.6, s=20)
        plt.colorbar(sc, ax=ax, label="Target")

    ax.set_xlabel("S'x — Component 1")
    ax.set_ylabel("S'x — Component 2")
    ax.set_title("Embedding Scatter (S'x projected)", fontsize=13, fontweight="bold")
    fig.tight_layout()
    return fig
