"""
Visualization utilities for feature space exploration.

PCA and t-SNE plots to inspect class separability in the
extracted feature representations.
"""

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE


def plot_pca_tsne(X, Y, legend_names=None, perplexity=35,
                  random_state=42, output_path=None, figsize=(12, 5),
                  dpi=150):
    """Generate side-by-side PCA and t-SNE plots of feature vectors.

    Parameters
    ----------
    X : np.ndarray of shape (n_samples, n_features)
        Feature matrix (e.g. VGG16 outputs).
    Y : np.ndarray of shape (n_samples,)
        Class labels.
    legend_names : dict, optional
        Mapping from raw label -> display label, e.g.
        {"Bengin cases": "Benign", ...}.
    perplexity : int
        t-SNE perplexity (default 35).
    random_state : int
        Seed for reproducibility.
    output_path : str, optional
        Path to save the figure. If None, saved as "pca_tsne.png".
    figsize : tuple
        Figure dimensions.
    dpi : int
        Output resolution.

    Returns
    -------
    fig : matplotlib Figure
    """
    if legend_names is None:
        legend_names = {}

    n_samples, n_features = X.shape

    # --- PCA: 50D → 2D ---
    n_pca = min(50, n_samples - 1, n_features)
    pca_pre = PCA(n_components=n_pca, random_state=random_state)
    X_pca50 = pca_pre.fit_transform(X)

    pca_2d = PCA(n_components=2, random_state=random_state)
    X_pca2 = pca_2d.fit_transform(X_pca50)
    explained_variance = pca_2d.explained_variance_ratio_ * 100

    # --- t-SNE: 50D → 2D ---
    tsne = TSNE(
        n_components=2,
        perplexity=perplexity,
        learning_rate="auto",
        init="pca",
        random_state=random_state
    )
    X_tsne = tsne.fit_transform(X_pca50)

    # --- Plot ---
    fig, axes = plt.subplots(1, 2, figsize=figsize)
    unique_classes = np.unique(Y)

    for class_name in unique_classes:
        idx = Y == class_name
        label = legend_names.get(class_name, class_name)

        axes[0].scatter(
            X_pca2[idx, 0], X_pca2[idx, 1],
            label=label, alpha=0.7, s=8
        )
    axes[0].set_xlabel(f"PC1 ({explained_variance[0]:.1f}%)")
    axes[0].set_ylabel(f"PC2 ({explained_variance[1]:.1f}%)")
    axes[0].set_title("PCA")
    axes[0].legend(markerscale=2)

    for class_name in unique_classes:
        idx = Y == class_name
        label = legend_names.get(class_name, class_name)

        axes[1].scatter(
            X_tsne[idx, 0], X_tsne[idx, 1],
            label=label, alpha=0.7, s=8
        )
    axes[1].set_xlabel("t-SNE dimension 1")
    axes[1].set_ylabel("t-SNE dimension 2")
    axes[1].set_title("t-SNE")
    axes[1].legend(markerscale=2)

    plt.tight_layout()

    if output_path is None:
        output_path = "pca_tsne.png"
    plt.savefig(output_path, dpi=dpi)
    plt.close()
    print(f"  Saved PCA/t-SNE plot to {output_path}")

    return fig


def plot_cv_accuracy(cv_results, split_type="", output_path=None,
                     figsize=(7, 5), dpi=150):
    """Plot cross-validation accuracy vs number of selected features.

    Parameters
    ----------
    cv_results : list of dict
        Each dict has keys "nfeatures" and "best_cv_accuracy".
    split_type : str
        Label for the plot title, e.g. "Image-level split".
    output_path : str, optional
        Output file path.
    figsize : tuple
    dpi : int

    Returns
    -------
    fig : matplotlib Figure
    """
    fig, ax = plt.subplots(figsize=figsize)

    nfeats = [r["nfeatures"] for r in cv_results]
    accs = [r["best_cv_accuracy"] for r in cv_results]

    ax.plot(nfeats, accs, marker="o")
    ax.set_xscale("log")
    ax.set_xlabel("Number of selected features")
    ax.set_ylabel("Mean 5-fold validation accuracy")
    if split_type:
        ax.set_title(split_type)
    plt.tight_layout()

    if output_path is None:
        output_path = f"{split_type.replace(' ', '_').replace('-', '_')}_cv.png"
    plt.savefig(output_path, dpi=dpi)
    plt.close()
    print(f"  Saved CV accuracy plot to {output_path}")

    return fig