"""Per-concept visualizations for RepScan."""

import numpy as np
import torch
from typing import Dict, Any, List, Optional
import io
import base64


def create_concept_overview_figure(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    cluster_labels: np.ndarray,
    concept_names: Optional[List[str]] = None,
) -> str:
    """
    Create overview figure showing all concepts with different colors.

    Shows difference vectors colored by concept assignment.

    Returns:
        Base64-encoded PNG string
    """
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    from sklearn.decomposition import PCA

    n_pairs = min(len(pos_activations), len(neg_activations))
    pos = pos_activations[:n_pairs].float().cpu().numpy()
    neg = neg_activations[:n_pairs].float().cpu().numpy()
    diff = pos - neg

    labels = cluster_labels[:n_pairs]
    n_concepts = len(np.unique(labels))

    if concept_names is None:
        concept_names = [f"Concept {i+1}" for i in range(n_concepts)]

    # PCA on difference vectors
    pca = PCA(n_components=2)
    diff_2d = pca.fit_transform(diff)

    # Create figure
    fig, ax = plt.subplots(figsize=(10, 8))

    colors = plt.cm.tab10(np.linspace(0, 1, n_concepts))

    for i in range(n_concepts):
        mask = labels == i
        ax.scatter(
            diff_2d[mask, 0], diff_2d[mask, 1],
            c=[colors[i]], label=concept_names[i], alpha=0.7, s=50
        )

    ax.set_xlabel(f"PC1 ({pca.explained_variance_ratio_[0]:.1%} var)")
    ax.set_ylabel(f"PC2 ({pca.explained_variance_ratio_[1]:.1%} var)")
    ax.set_title("Concept Overview: Difference Vectors by Concept")
    ax.legend(loc='best')
    ax.grid(True, alpha=0.3)

    # Convert to base64
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
    buf.seek(0)
    plt.close(fig)

    return base64.b64encode(buf.read()).decode('utf-8')


def create_per_concept_figure(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    concept_id: int,
    concept_name: str = None,
    coherence: float = None,
) -> str:
    """
    Create visualization for a single concept showing pos vs neg separation.

    Shows PCA projection with pos/neg colored differently.

    Returns:
        Base64-encoded PNG string
    """
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    from sklearn.decomposition import PCA
    from sklearn.linear_model import LogisticRegression

    pos = pos_activations.float().cpu().numpy()
    neg = neg_activations.float().cpu().numpy()
    n_pairs = min(len(pos), len(neg))

    if n_pairs < 3:
        # Not enough data
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.text(0.5, 0.5, "Not enough data", ha='center', va='center')
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=150)
        buf.seek(0)
        plt.close(fig)
        return base64.b64encode(buf.read()).decode('utf-8')

    # Combine for PCA
    X = np.vstack([pos[:n_pairs], neg[:n_pairs]])
    y = np.array([1] * n_pairs + [0] * n_pairs)

    pca = PCA(n_components=2)
    X_2d = pca.fit_transform(X)

    pos_2d = X_2d[:n_pairs]
    neg_2d = X_2d[n_pairs:]

    # Fit linear classifier to get decision boundary
    clf = LogisticRegression(max_iter=1000)
    clf.fit(X_2d, y)
    accuracy = clf.score(X_2d, y)

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

    ax.scatter(pos_2d[:, 0], pos_2d[:, 1], c='green', label='Positive', alpha=0.7, s=50)
    ax.scatter(neg_2d[:, 0], neg_2d[:, 1], c='red', label='Negative', alpha=0.7, s=50)

    # Draw decision boundary
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()

    xx = np.linspace(xlim[0], xlim[1], 100)
    yy = np.linspace(ylim[0], ylim[1], 100)
    XX, YY = np.meshgrid(xx, yy)
    Z = clf.predict_proba(np.c_[XX.ravel(), YY.ravel()])[:, 1]
    Z = Z.reshape(XX.shape)

    ax.contour(XX, YY, Z, levels=[0.5], colors='black', linestyles='--', linewidths=2)

    title = f"Concept {concept_id}"
    if concept_name:
        title += f": {concept_name}"
    title += f"\nLinear accuracy: {accuracy:.1%}"
    if coherence is not None:
        title += f" | Coherence: {coherence:.2f}"
    title += f" | N={n_pairs} pairs"

    ax.set_xlabel(f"PC1 ({pca.explained_variance_ratio_[0]:.1%} var)")
    ax.set_ylabel(f"PC2 ({pca.explained_variance_ratio_[1]:.1%} var)")
    ax.set_title(title)
    ax.legend(loc='best')
    ax.grid(True, alpha=0.3)

    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
    buf.seek(0)
    plt.close(fig)

    return base64.b64encode(buf.read()).decode('utf-8')


def create_all_concept_figures(
    pos_activations: torch.Tensor,
    neg_activations: torch.Tensor,
    cluster_labels: np.ndarray,
    concepts: List[Dict[str, Any]],
) -> Dict[str, str]:
    """
    Create all concept visualizations.

    Returns:
        Dict with:
            - overview: Base64 overview figure
            - concept_1, concept_2, ...: Base64 per-concept figures
    """
    n_pairs = min(len(pos_activations), len(neg_activations))
    labels = cluster_labels[:n_pairs]

    concept_names = [c.get("name", f"Concept {c['id']}") for c in concepts]

    result = {}

    # Overview
    result["overview"] = create_concept_overview_figure(
        pos_activations, neg_activations, labels, concept_names
    )

    # Per-concept figures
    for concept in concepts:
        concept_id = concept["id"]
        idx = concept_id - 1  # 0-indexed
        mask = labels == idx

        if mask.sum() < 3:
            continue

        pos_concept = pos_activations[:n_pairs][mask]
        neg_concept = neg_activations[:n_pairs][mask]

        result[f"concept_{concept_id}"] = create_per_concept_figure(
            pos_concept,
            neg_concept,
            concept_id=concept_id,
            concept_name=concept.get("name"),
            coherence=concept.get("coherence"),
        )

    return result
