from __future__ import annotations

from dataclasses import dataclass
from typing import List, Tuple

import torch

from hippotorch.memory.store import MemoryStore


@dataclass
class PrototypeResult:
    num_candidates: int
    num_prototypes: int
    indices_pruned: List[int]


class PrototypeExtractor:
    def __init__(self, kmeans_iters: int = 10, keep_per_cluster: int = 1) -> None:
        self.kmeans_iters = int(max(1, kmeans_iters))
        self.keep_per_cluster = int(max(1, keep_per_cluster))

    @torch.no_grad()
    def extract_prototypes(
        self,
        memory: MemoryStore,
        encoder,
        *,
        num_prototypes: int = 2,
        reward_percentile: float = 0.5,
        prune_sources: bool = True,
    ) -> PrototypeResult:
        del encoder  # encoder is not needed for prototype extraction
        if len(memory) == 0:
            return PrototypeResult(0, 0, [])

        # Select candidate episodes by reward percentile
        rewards = torch.tensor(
            [memory._get_episode(i).total_reward for i in range(len(memory))],
            dtype=torch.float,
        )
        q = float(max(0.0, min(1.0, reward_percentile)))
        threshold = torch.quantile(rewards, q=q, interpolation="lower")
        candidate_indices: List[int] = [
            i for i, r in enumerate(rewards.tolist()) if r >= float(threshold)
        ]
        if not candidate_indices:
            return PrototypeResult(0, 0, [])

        # Build candidate key matrix
        device = memory.keys.device
        cand_keys = memory.keys[torch.tensor(candidate_indices, dtype=torch.long, device=device)]
        if cand_keys.numel() == 0:
            return PrototypeResult(0, 0, [])

        k = int(max(1, min(num_prototypes, cand_keys.shape[0])))
        centroids, assignments = _kmeans_full_batch(cand_keys, k=k, iters=self.kmeans_iters)

        # Medoid selection: pick closest real episode(s) to each centroid
        keep_indices: List[int] = []
        for cluster_id in range(k):
            member_mask = assignments == cluster_id
            if not torch.any(member_mask):
                continue
            members = cand_keys[member_mask]
            # distances to centroid
            dists = (
                torch.cdist(
                    members.unsqueeze(0),
                    centroids[cluster_id : cluster_id + 1].unsqueeze(0),
                )
                .squeeze(0)
                .squeeze(-1)
            )
            order = torch.argsort(dists)
            # map back to candidate_indices
            member_positions = torch.nonzero(member_mask, as_tuple=False).squeeze(-1)
            take = min(self.keep_per_cluster, members.shape[0])
            selected_local = member_positions[order[:take]].tolist()
            keep_indices.extend([candidate_indices[int(i)] for i in selected_local])

        # Ensure uniqueness and preserve original order among candidates
        keep_set = set(keep_indices)
        keep_indices = [idx for idx in candidate_indices if idx in keep_set]

        pruned: List[int] = []
        if prune_sources:
            drop = [idx for idx in candidate_indices if idx not in keep_set]
            if drop:
                memory.prune_indices(drop)
                pruned = drop

        return PrototypeResult(
            num_candidates=len(candidate_indices),
            num_prototypes=len(keep_indices),
            indices_pruned=pruned,
        )


def _kmeans_full_batch(x: torch.Tensor, *, k: int, iters: int) -> Tuple[torch.Tensor, torch.Tensor]:
    """Simple full-batch k-means on ``x`` using L2 distance.

    Args:
        x: Tensor shaped [N, D]
        k: Number of clusters (1..N)
        iters: Number of refinement iterations

    Returns:
        centroids: [k, D] tensor of cluster centers
        assignments: [N] long tensor of cluster ids per point
    """

    N = int(x.shape[0])
    if k <= 0 or N == 0:
        raise ValueError("invalid k or empty input for k-means")
    device = x.device
    # Deterministic farthest-first initialization for stable behavior
    chosen: List[int] = [0]
    while len(chosen) < k:
        current = x[torch.tensor(chosen, dtype=torch.long, device=device)]  # [m, D]
        # distances to nearest chosen centroid
        d2 = ((x.unsqueeze(1) - current.unsqueeze(0)) ** 2).sum(dim=-1)  # [N, m]
        min_d2 = d2.min(dim=1).values  # [N]
        # exclude already chosen
        min_d2[torch.tensor(chosen, dtype=torch.long, device=device)] = -1.0
        next_idx = int(torch.argmax(min_d2).item())
        if next_idx in chosen:
            break
        chosen.append(next_idx)
    centroids = x[torch.tensor(chosen[:k], dtype=torch.long, device=device)].clone()

    assignments = torch.zeros(N, dtype=torch.long, device=device)
    for _ in range(max(1, iters)):
        # Assign to nearest centroid
        # Compute pairwise distances: [N, k]
        # Using (a-b)^2 = a^2 + b^2 - 2ab for efficiency
        x_norm = (x**2).sum(dim=1, keepdim=True)  # [N, 1]
        c_norm = (centroids**2).sum(dim=1, keepdim=True).t()  # [1, k]
        d2 = x_norm + c_norm - 2 * (x @ centroids.t())  # [N, k]
        assignments = torch.argmin(d2, dim=1)

        # Update centroids
        new_centroids = torch.empty_like(centroids)
        for j in range(k):
            mask = assignments == j
            if torch.any(mask):
                new_centroids[j] = x[mask].mean(dim=0)
            else:
                # Reinitialize empty cluster to a random point
                new_centroids[j] = x[torch.randint(0, N, (1,), device=device)]
        centroids = new_centroids

    return centroids, assignments


__all__ = ["PrototypeExtractor", "PrototypeResult"]
