"""
projections.py
---------------
Embedding and graph projection utilities for HIEROMEM.

This module provides helpers for projecting high-dimensional episode
embeddings into 2D or 3D for visualization and analysis. It is used by:

  • runners/narrative_eval.py
  • hieromem_visualizer/ (frontend visual graphs)
  • notebooks/01_trust_decay.ipynb

Supports:
  - PCA, t-SNE, and UMAP projections
  - Trust-weighted coloring
  - Scene graph extraction summary (optional)

Dependencies: numpy, sklearn; umap-learn (optional)

Author: Sai Rishi Kiran Mannava
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
from typing import Protocol

import numpy as np
import numpy.typing as npt

try:
    from sklearn.decomposition import PCA
    from sklearn.manifold import TSNE
except ImportError:
    PCA = TSNE = None

try:
    import umap
except ImportError:
    umap = None


# --------------------------------------------------------------------------- #
# Config Dataclass
# --------------------------------------------------------------------------- #
@dataclass
class ProjectionConfig:
    method: str = "umap"  # "pca" | "tsne" | "umap"
    n_components: int = 2
    random_state: int = 42
    perplexity: int = 30
    min_dist: float = 0.1
    n_neighbors: int = 10


class EpisodeLike(Protocol):
    e: npt.NDArray[np.float32]
    trust: float


# --------------------------------------------------------------------------- #
# Projection Utilities
# --------------------------------------------------------------------------- #
def embedding_matrix(
        episodes: Sequence[EpisodeLike]) -> npt.NDArray[np.float32]:
    """Stack all embedding vectors into a (N, d) numpy array."""
    if not episodes:
        return np.zeros((0, 0), dtype=np.float32)
    vecs = [np.asarray(ep.e, dtype=np.float32) for ep in episodes]
    return np.stack(vecs) if vecs else np.zeros((0, 0), dtype=np.float32)


def color_by_trust(
    episodes: Sequence[EpisodeLike],
    cmap: str = "viridis",
) -> npt.NDArray[np.float64]:
    """
    Generate trust-based color intensities (0–1).
    Returns a (N,) array normalized to [0,1].
    """
    if not episodes:
        return np.array([], dtype=np.float64)
    trusts = np.array([ep.trust for ep in episodes], dtype=np.float64)
    t_min, t_max = trusts.min(), trusts.max()
    if t_max - t_min < 1e-8:
        return np.ones_like(trusts, dtype=np.float64)
    return np.asarray((trusts - t_min) / (t_max - t_min), dtype=np.float64)


def project_embeddings(
    episodes: Sequence[EpisodeLike],
    config: ProjectionConfig | None = None,
) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.float64]]:
    """
    Reduce episode embeddings to 2D or 3D for visualization.

    Returns:
        coords : np.ndarray of shape (N, n_components)
        colors : np.ndarray of trust-normalized color values
    """
    resolved_config = config or ProjectionConfig()

    embeddings = embedding_matrix(episodes)
    if embeddings.size == 0:
        return (
            np.zeros((0, resolved_config.n_components), dtype=np.float32),
            np.array([], dtype=np.float64),
        )

    method = resolved_config.method.lower()

    if method == "pca":
        if PCA is None:
            raise ImportError("scikit-learn required for PCA projection.")
        reducer = PCA(
            n_components=resolved_config.n_components,
            random_state=resolved_config.random_state,
        )
        coords = reducer.fit_transform(embeddings)

    elif method == "tsne":
        if TSNE is None:
            raise ImportError("scikit-learn required for t-SNE projection.")
        reducer = TSNE(
            n_components=resolved_config.n_components,
            perplexity=resolved_config.perplexity,
            random_state=resolved_config.random_state,
        )
        coords = reducer.fit_transform(embeddings)

    elif method == "umap":
        if umap is None:
            raise ImportError(
                "umap-learn not installed. Run `pip install umap-learn`.")
        reducer = umap.UMAP(
            n_components=resolved_config.n_components,
            n_neighbors=resolved_config.n_neighbors,
            min_dist=resolved_config.min_dist,
            random_state=resolved_config.random_state,
        )
        coords = reducer.fit_transform(embeddings)

    else:
        raise ValueError(f"Unknown projection method: {method}")

    colors = color_by_trust(episodes)
    return np.asarray(coords, dtype=np.float32), colors


def trust_summary(episodes: Sequence[EpisodeLike]) -> dict[str, float]:
    """Return a quick numeric summary of trust values for plotting."""
    if not episodes:
        return {"mean": 0.0, "min": 0.0, "max": 0.0, "std": 0.0}
    trusts = np.array([ep.trust for ep in episodes])
    return {
        "mean": float(np.mean(trusts)),
        "min": float(np.min(trusts)),
        "max": float(np.max(trusts)),
        "std": float(np.std(trusts)),
    }


def plot_projection(
    coords: npt.NDArray[np.float32],
    colors: npt.NDArray[np.float64],
    title: str = "HIEROMEM Embedding Projection",
    save_path: str | None = None,
) -> None:
    """
    Optional quick visualization (requires matplotlib).
    Used only in notebooks or research runners.
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise ImportError(
            "matplotlib is required for plotting projections.") from exc

    plt.figure(figsize=(6, 5))
    sc = plt.scatter(coords[:, 0], coords[:, 1],
                     c=colors, cmap="viridis", s=50)
    plt.colorbar(sc, label="Trust")
    plt.title(title)
    plt.xlabel("Dim 1")
    plt.ylabel("Dim 2")
    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=200)
    plt.show()
