from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, List, Optional

import torch

from hippotorch.utils.diagnostics import pca_keys


@dataclass
class ProjectionResult:
    coords: torch.Tensor  # [N, dims]
    rewards: List[float]
    contexts: List[int]


class MemoryProjector:
    def __init__(
        self,
        *,
        sample_size: int = 100,
        dims: int = 2,
        method: str = "pca",
        context_fn: Optional[Callable[[object], int]] = None,
    ) -> None:
        self.sample_size = int(sample_size)
        self.dims = int(dims)
        self.method = method
        self.context_fn = context_fn

    @torch.no_grad()
    def project(self, memory) -> ProjectionResult:
        n = min(self.sample_size, len(memory))
        if n == 0:
            return ProjectionResult(
                coords=torch.zeros(0, self.dims), rewards=[], contexts=[]
            )
        indices = torch.randperm(len(memory))[:n].tolist()
        keys = memory.keys[indices]
        coords = self._project_keys(keys)
        rewards = [float(memory.episodes[i].total_reward) for i in indices]
        if self.context_fn is None:
            contexts = [0 for _ in indices]
        else:
            contexts = [int(self.context_fn(memory.episodes[i])) for i in indices]
        return ProjectionResult(coords=coords, rewards=rewards, contexts=contexts)

    def _project_keys(self, keys: torch.Tensor) -> torch.Tensor:
        if self.method == "pca":
            _, coords = pca_keys(keys, k=self.dims)
            return coords
        if self.method == "umap":
            try:  # pragma: no cover - optional dep
                import umap
            except Exception as exc:  # pragma: no cover - optional dep
                raise ImportError(
                    "UMAP projection requires umap-learn. Install hippotorch[umap]."
                ) from exc
            reducer = umap.UMAP(
                n_components=self.dims,
                metric="cosine",
                random_state=0,
            )
            coords = reducer.fit_transform(keys.cpu().numpy())
            return torch.from_numpy(coords).to(keys.device)
        raise ValueError(f"Unknown projector method: {self.method!r}")


__all__ = ["MemoryProjector", "ProjectionResult"]
