"""
Visualization functions for activation geometry.

All functions return matplotlib figures or data for plotting.
"""

import torch
import numpy as np
from typing import Dict, Any, List, Tuple, Optional
import warnings


def plot_pca_projection(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    n_components: int = 2,
    title: str = "PCA Projection",
) -> Dict[str, Any]:
    """
    Project activations to 2D or 3D using PCA and return plot data.
    """
    from sklearn.decomposition import PCA

    pos = pos_activations.float().cpu().numpy()
    neg = neg_activations.float().cpu().numpy()

    # Combine for fitting PCA
    X = np.vstack([pos, neg])
    labels = np.array([1] * len(pos) + [0] * len(neg))

    pca = PCA(n_components=n_components)
    X_pca = pca.fit_transform(X)

    pos_pca = X_pca[:len(pos)]
    neg_pca = X_pca[len(pos):]

    return {
        "pos_projected": pos_pca,
        "neg_projected": neg_pca,
        "explained_variance": pca.explained_variance_ratio_.tolist(),
        "components": pca.components_,
        "n_components": n_components,
        "title": title,
    }


def plot_diff_vectors(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    title: str = "Difference Vectors",
) -> Dict[str, Any]:
    """
    Visualize the difference vectors projected to 2D.
    """
    from sklearn.decomposition import PCA

    pos = pos_activations.float().cpu().numpy()
    neg = neg_activations.float().cpu().numpy()

    n = min(len(pos), len(neg))
    diffs = pos[:n] - neg[:n]

    # Project diffs to 2D
    pca = PCA(n_components=2)
    diffs_2d = pca.fit_transform(diffs)

    # Mean direction
    mean_diff = diffs.mean(axis=0)
    mean_diff_2d = pca.transform(mean_diff.reshape(1, -1))[0]

    return {
        "diffs_projected": diffs_2d,
        "mean_diff_projected": mean_diff_2d,
        "explained_variance": pca.explained_variance_ratio_.tolist(),
        "title": title,
    }


def plot_norm_distribution(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    title: str = "Norm Distribution",
) -> Dict[str, Any]:
    """
    Plot distribution of activation norms.
    """
    pos = pos_activations.float().cpu().numpy()
    neg = neg_activations.float().cpu().numpy()

    pos_norms = np.linalg.norm(pos, axis=1)
    neg_norms = np.linalg.norm(neg, axis=1)

    return {
        "pos_norms": pos_norms,
        "neg_norms": neg_norms,
        "pos_mean": float(pos_norms.mean()),
        "neg_mean": float(neg_norms.mean()),
        "pos_std": float(pos_norms.std()),
        "neg_std": float(neg_norms.std()),
        "title": title,
    }


def plot_alignment_distribution(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    title: str = "Alignment Distribution",
) -> Dict[str, Any]:
    """
    Plot distribution of how well each diff aligns with mean diff.
    """
    pos = pos_activations.float().cpu().numpy()
    neg = neg_activations.float().cpu().numpy()

    n = min(len(pos), len(neg))
    diffs = pos[:n] - neg[:n]

    # Mean direction
    mean_diff = diffs.mean(axis=0)
    mean_diff_norm = np.linalg.norm(mean_diff)

    if mean_diff_norm < 1e-8:
        return {"error": "mean diff too small"}

    mean_diff_normalized = mean_diff / mean_diff_norm

    # Per-diff alignment
    diff_norms = np.linalg.norm(diffs, axis=1)
    valid = diff_norms > 1e-8
    alignments = np.zeros(n)
    alignments[valid] = (diffs[valid] / diff_norms[valid, np.newaxis]) @ mean_diff_normalized

    return {
        "alignments": alignments,
        "mean": float(alignments.mean()),
        "std": float(alignments.std()),
        "min": float(alignments.min()),
        "max": float(alignments.max()),
        "title": title,
    }


def plot_eigenvalue_spectrum(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    n_components: int = 50,
    title: str = "Eigenvalue Spectrum",
) -> Dict[str, Any]:
    """
    Plot eigenvalue spectrum of the difference vectors.
    """
    from sklearn.decomposition import PCA

    pos = pos_activations.float().cpu().numpy()
    neg = neg_activations.float().cpu().numpy()

    n = min(len(pos), len(neg))
    diffs = pos[:n] - neg[:n]

    n_comp = min(n_components, n - 1, diffs.shape[1])
    pca = PCA(n_components=n_comp)
    pca.fit(diffs)

    return {
        "eigenvalues": pca.explained_variance_.tolist(),
        "explained_variance_ratio": pca.explained_variance_ratio_.tolist(),
        "cumulative_variance": np.cumsum(pca.explained_variance_ratio_).tolist(),
        "title": title,
    }


def plot_pairwise_distances(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    max_samples: int = 200,
    title: str = "Pairwise Distances",
) -> Dict[str, Any]:
    """
    Plot distribution of pairwise distances within and between classes.
    """
    from scipy.spatial.distance import pdist

    pos = pos_activations.float().cpu().numpy()
    neg = neg_activations.float().cpu().numpy()

    # Subsample if needed
    if len(pos) > max_samples:
        idx = np.random.choice(len(pos), max_samples, replace=False)
        pos = pos[idx]
    if len(neg) > max_samples:
        idx = np.random.choice(len(neg), max_samples, replace=False)
        neg = neg[idx]

    # Within-class distances
    pos_dists = pdist(pos)
    neg_dists = pdist(neg)

    # Between-class distances
    from scipy.spatial.distance import cdist
    between_dists = cdist(pos, neg).flatten()

    return {
        "pos_within": pos_dists,
        "neg_within": neg_dists,
        "between": between_dists,
        "pos_within_mean": float(pos_dists.mean()),
        "neg_within_mean": float(neg_dists.mean()),
        "between_mean": float(between_dists.mean()),
        "title": title,
    }


def plot_tsne_projection(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    perplexity: int = 30,
    title: str = "t-SNE Projection",
) -> Dict[str, Any]:
    """
    Project activations using t-SNE.
    """
    try:
        from sklearn.manifold import TSNE
    except ImportError:
        return {"error": "sklearn not available"}

    pos = pos_activations.float().cpu().numpy()
    neg = neg_activations.float().cpu().numpy()

    X = np.vstack([pos, neg])

    # Adjust perplexity if needed
    n_samples = len(X)
    perplexity = min(perplexity, n_samples - 1)

    if n_samples < 5:
        return {"error": "not enough samples for t-SNE"}

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        tsne = TSNE(n_components=2, perplexity=perplexity, random_state=42)
        X_tsne = tsne.fit_transform(X)

    pos_tsne = X_tsne[:len(pos)]
    neg_tsne = X_tsne[len(pos):]

    return {
        "pos_projected": pos_tsne,
        "neg_projected": neg_tsne,
        "title": title,
    }


def plot_umap_projection(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    n_neighbors: int = 15,
    min_dist: float = 0.1,
    title: str = "UMAP Projection",
) -> Dict[str, Any]:
    """
    Project activations using UMAP.
    """
    try:
        import umap
    except ImportError:
        return {"error": "umap not available"}

    pos = pos_activations.float().cpu().numpy()
    neg = neg_activations.float().cpu().numpy()

    X = np.vstack([pos, neg])

    n_samples = len(X)
    n_neighbors = min(n_neighbors, n_samples - 1)

    if n_samples < 5:
        return {"error": "not enough samples for UMAP"}

    reducer = umap.UMAP(n_components=2, n_neighbors=n_neighbors, min_dist=min_dist, random_state=42)
    X_umap = reducer.fit_transform(X)

    pos_umap = X_umap[:len(pos)]
    neg_umap = X_umap[len(pos):]

    return {
        "pos_projected": pos_umap,
        "neg_projected": neg_umap,
        "title": title,
    }


def plot_cone_visualization(
    activations: torch.Tensor,
    title: str = "Cone Structure",
) -> Dict[str, Any]:
    """
    Visualize how cone-like the activations are.

    Projects to 2D and shows angular distribution from mean direction.
    """
    X = activations.float().cpu().numpy()

    # Normalize
    norms = np.linalg.norm(X, axis=1, keepdims=True)
    valid = norms.squeeze() > 1e-8
    X_norm = X[valid] / norms[valid]

    # Mean direction
    mean_dir = X_norm.mean(axis=0)
    mean_dir = mean_dir / (np.linalg.norm(mean_dir) + 1e-8)

    # Angles from mean
    cos_angles = X_norm @ mean_dir
    angles = np.degrees(np.arccos(np.clip(cos_angles, -1, 1)))

    # Project to 2D: mean direction and orthogonal
    from sklearn.decomposition import PCA
    pca = PCA(n_components=2)
    X_2d = pca.fit_transform(X_norm)

    return {
        "angles_from_mean": angles,
        "projected_2d": X_2d,
        "mean_angle": float(angles.mean()),
        "max_angle": float(angles.max()),
        "title": title,
    }


def plot_layer_comparison(
    activations_by_layer: Dict[int, Tuple[torch.Tensor, torch.Tensor]],
    metric_fn,
    metric_name: str = "metric",
    title: str = "Metric Across Layers",
) -> Dict[str, Any]:
    """
    Plot a metric across different layers.
    """
    layers = sorted(activations_by_layer.keys())
    values = []

    for layer in layers:
        pos, neg = activations_by_layer[layer]
        try:
            result = metric_fn(pos, neg)
            if isinstance(result, dict):
                # Try to extract a scalar
                val = result.get(metric_name, result.get("mean", 0))
            else:
                val = result
            values.append(float(val))
        except:
            values.append(0)

    return {
        "layers": layers,
        "values": values,
        "metric_name": metric_name,
        "title": title,
    }


def create_summary_figure(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
) -> Dict[str, Any]:
    """
    Create a summary of all visualizations.
    """
    return {
        "pca": plot_pca_projection(pos_activations, neg_activations),
        "diff_vectors": plot_diff_vectors(pos_activations, neg_activations),
        "norms": plot_norm_distribution(pos_activations, neg_activations),
        "alignments": plot_alignment_distribution(pos_activations, neg_activations),
        "eigenvalues": plot_eigenvalue_spectrum(pos_activations, neg_activations),
        "distances": plot_pairwise_distances(pos_activations, neg_activations),
    }


def render_matplotlib_figure(
    plot_data: Dict[str, Any],
    plot_type: str,
    figsize: Tuple[int, int] = (8, 6),
):
    """
    Render plot data to a matplotlib figure.

    Args:
        plot_data: Output from one of the plot_* functions
        plot_type: One of 'pca', 'diff_vectors', 'norms', 'alignments',
                   'eigenvalues', 'distances', 'tsne', 'umap', 'cone'
    """
    import matplotlib.pyplot as plt

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

    if "error" in plot_data:
        ax.text(0.5, 0.5, f"Error: {plot_data['error']}", ha='center', va='center')
        return fig

    title = plot_data.get("title", plot_type)

    if plot_type in ['pca', 'tsne', 'umap']:
        pos = plot_data["pos_projected"]
        neg = plot_data["neg_projected"]
        ax.scatter(pos[:, 0], pos[:, 1], c='blue', alpha=0.6, label='Positive')
        ax.scatter(neg[:, 0], neg[:, 1], c='red', alpha=0.6, label='Negative')
        ax.legend()
        ax.set_xlabel("Component 1")
        ax.set_ylabel("Component 2")

    elif plot_type == 'diff_vectors':
        diffs = plot_data["diffs_projected"]
        mean = plot_data["mean_diff_projected"]
        ax.scatter(diffs[:, 0], diffs[:, 1], c='green', alpha=0.6)
        ax.arrow(0, 0, mean[0], mean[1], head_width=0.1, head_length=0.05, fc='black', ec='black')
        ax.set_xlabel("PC1")
        ax.set_ylabel("PC2")

    elif plot_type == 'norms':
        ax.hist(plot_data["pos_norms"], bins=20, alpha=0.5, label='Positive', color='blue')
        ax.hist(plot_data["neg_norms"], bins=20, alpha=0.5, label='Negative', color='red')
        ax.legend()
        ax.set_xlabel("Norm")
        ax.set_ylabel("Count")

    elif plot_type == 'alignments':
        ax.hist(plot_data["alignments"], bins=20, color='green', alpha=0.7)
        ax.axvline(plot_data["mean"], color='black', linestyle='--', label=f'Mean: {plot_data["mean"]:.2f}')
        ax.legend()
        ax.set_xlabel("Alignment with mean direction")
        ax.set_ylabel("Count")

    elif plot_type == 'eigenvalues':
        ax.bar(range(len(plot_data["explained_variance_ratio"])), plot_data["explained_variance_ratio"])
        ax.set_xlabel("Principal Component")
        ax.set_ylabel("Explained Variance Ratio")

    elif plot_type == 'distances':
        ax.hist(plot_data["pos_within"], bins=30, alpha=0.5, label='Within Pos', color='blue')
        ax.hist(plot_data["neg_within"], bins=30, alpha=0.5, label='Within Neg', color='red')
        ax.hist(plot_data["between"], bins=30, alpha=0.5, label='Between', color='green')
        ax.legend()
        ax.set_xlabel("Distance")
        ax.set_ylabel("Count")

    elif plot_type == 'cone':
        ax.hist(plot_data["angles_from_mean"], bins=20, color='purple', alpha=0.7)
        ax.axvline(plot_data["mean_angle"], color='black', linestyle='--',
                   label=f'Mean: {plot_data["mean_angle"]:.1f}°')
        ax.legend()
        ax.set_xlabel("Angle from mean direction (degrees)")
        ax.set_ylabel("Count")

    ax.set_title(title)
    plt.tight_layout()

    return fig
