"""Visual embedding analysis for LeRobot datasets using SigLIP."""

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

from orbit.analyzer.dataset_loader import DatasetInfo


@dataclass
class EmbeddingResult:
    """Result of computing frame embeddings from a dataset."""

    embeddings: np.ndarray  # N x D (SigLIP embedding dim)
    episode_indices: list[int]  # which episode each embedding belongs to
    frame_types: list[str]  # "first", "mid_early", "mid_late", "last"
    metadata: dict = field(default_factory=dict)  # model name, timing, etc.


@dataclass
class EmbeddingAnalysis:
    """Analysis of visual embeddings."""

    n_clusters: int
    cluster_labels: np.ndarray
    cluster_sizes: list[int]
    pca_2d: np.ndarray  # N x 2 for visualization
    diversity_score: float  # 0-1
    outlier_indices: list[int]  # episode indices of outliers
    insights: list[str]  # human-readable findings


def _repo_cache_dir(repo_id: str) -> Path:
    """Return cache directory for a given repo ID."""
    repo_hash = hashlib.sha256(repo_id.encode()).hexdigest()[:16]
    return Path.home() / ".cache" / "orbit" / repo_hash


def _get_device() -> str:
    """Detect best available device."""
    try:
        import torch

        if torch.cuda.is_available():
            return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
    except ImportError:
        pass
    return "cpu"


def _sample_frame_indices(num_frames: int, max_frames: int = 4) -> list[tuple[int, str]]:
    """Return (frame_index, frame_type) pairs for an episode.

    Samples first, 1/3, 2/3, and last frames.
    """
    if num_frames <= 0:
        return []
    if num_frames == 1:
        return [(0, "first")]

    indices: list[tuple[int, str]] = []
    positions = [
        (0, "first"),
        (num_frames // 3, "mid_early"),
        (2 * num_frames // 3, "mid_late"),
        (num_frames - 1, "last"),
    ]
    seen = set()
    for idx, label in positions:
        if idx not in seen and len(indices) < max_frames:
            indices.append((idx, label))
            seen.add(idx)
    return indices


def _sample_episode_indices(total: int, max_episodes: int) -> list[int]:
    """Select evenly-spaced episode indices."""
    if total <= max_episodes:
        return list(range(total))
    step = total / max_episodes
    return [int(i * step) for i in range(max_episodes)]


def compute_frame_embeddings(
    repo_id: str,
    info: DatasetInfo,
    max_frames_per_episode: int = 4,
    max_episodes: int = 100,
    device: str | None = None,
    no_cache: bool = False,
) -> EmbeddingResult:
    """Compute SigLIP embeddings for sampled frames from a dataset.

    Args:
        repo_id: HuggingFace repo ID.
        info: Dataset metadata.
        max_frames_per_episode: Number of frames to sample per episode.
        max_episodes: Maximum episodes to process.
        device: Force device (cpu/cuda/mps). Auto-detects if None.
        no_cache: If True, bypass cached embeddings.

    Returns:
        EmbeddingResult with embeddings and metadata.
    """
    import torch
    from transformers import AutoImageProcessor, AutoModel

    cache_dir = _repo_cache_dir(repo_id)
    cache_file = cache_dir / "embeddings.npz"

    # Check cache
    if not no_cache and cache_file.exists():
        try:
            cached = np.load(cache_file, allow_pickle=True)
            meta = cached["metadata"].item() if "metadata" in cached else {}
            from rich.console import Console

            Console().print("[dim]Using cached embeddings[/dim]")
            return EmbeddingResult(
                embeddings=cached["embeddings"],
                episode_indices=cached["episode_indices"].tolist(),
                frame_types=cached["frame_types"].tolist(),
                metadata=meta,
            )
        except Exception as e:
            import logging as _log
            _log.getLogger(__name__).debug("Embedding cache corrupt, recomputing: %s", e)

    if device is None:
        device = _get_device()

    start_time = time.time()

    # Select episodes to process
    episode_indices = _sample_episode_indices(info.num_episodes, max_episodes)

    # Load model
    model_name = "google/siglip-base-patch16-384"
    processor = AutoImageProcessor.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).vision_model
    model = model.to(device)
    model.eval()

    from rich.progress import Progress

    all_embeddings = []
    all_episode_indices = []
    all_frame_types = []

    with Progress() as progress:
        task = progress.add_task(
            f"Computing embeddings ({device})...",
            total=len(episode_indices),
        )

        for ep_idx in episode_indices:
            frames = _load_episode_frames(repo_id, info, ep_idx, max_frames_per_episode)

            if not frames:
                progress.advance(task)
                continue

            # Process frames in batches
            images = [f[0] for f in frames]
            frame_types = [f[1] for f in frames]

            for batch_start in range(0, len(images), 32):
                batch_images = images[batch_start : batch_start + 32]
                batch_types = frame_types[batch_start : batch_start + 32]

                inputs = processor(images=batch_images, return_tensors="pt").to(device)

                with torch.no_grad():
                    outputs = model(**inputs)
                    # Use CLS token / pooled output
                    embs = outputs.pooler_output
                    if embs is None:
                        embs = outputs.last_hidden_state[:, 0]
                    embs = embs.cpu().numpy()

                all_embeddings.append(embs)
                all_episode_indices.extend([ep_idx] * len(batch_images))
                all_frame_types.extend(batch_types)

            progress.advance(task)

    elapsed = time.time() - start_time

    if all_embeddings:
        embeddings = np.concatenate(all_embeddings, axis=0)
    else:
        embeddings = np.empty((0, 768))

    metadata = {
        "model_name": model_name,
        "device": device,
        "processing_time_seconds": elapsed,
        "num_episodes_sampled": len(episode_indices),
        "num_frames": len(all_episode_indices),
        "max_frames_per_episode": max_frames_per_episode,
    }

    result = EmbeddingResult(
        embeddings=embeddings,
        episode_indices=all_episode_indices,
        frame_types=all_frame_types,
        metadata=metadata,
    )

    # Cache results
    if not no_cache:
        cache_dir.mkdir(parents=True, exist_ok=True)
        np.savez(
            cache_file,
            embeddings=embeddings,
            episode_indices=np.array(all_episode_indices),
            frame_types=np.array(all_frame_types),
            metadata=np.array(metadata),
        )

    return result


def _load_episode_frames(
    repo_id: str,
    info: DatasetInfo,
    episode_idx: int,
    max_frames: int,
) -> list[tuple[np.ndarray, str]]:
    """Load and sample frames from a single episode's video.

    Returns list of (image_array, frame_type) tuples.
    Images are resized to 384x384 for SigLIP.

    Supports both HuggingFace repos and local dataset directories.
    """
    from pathlib import Path

    frames: list[tuple[np.ndarray, str]] = []

    # Check if source is a local directory
    local_root = Path(repo_id).expanduser()
    is_local = local_root.is_dir()

    # Try each video key to find a video file
    for video_key in info.video_keys:
        # LeRobot v2 video path convention
        filename = f"videos/{video_key}/episode_{episode_idx:06d}.mp4"

        if is_local:
            video_path = local_root / filename
            if not video_path.exists():
                continue
            video_path = str(video_path)
        else:
            from huggingface_hub import hf_hub_download
            from huggingface_hub.utils import EntryNotFoundError

            try:
                video_path = hf_hub_download(
                    repo_id=repo_id,
                    filename=filename,
                    repo_type="dataset",
                )
            except (EntryNotFoundError, Exception):
                continue

        frames = _extract_frames_from_video(video_path, max_frames)
        if frames:
            break

    # For local dirs without video_keys, try to find any video files
    if not frames and is_local:
        video_dir = local_root / "videos"
        if video_dir.exists():
            pattern = f"**/episode_{episode_idx:06d}.mp4"
            for video_file in video_dir.glob(pattern):
                frames = _extract_frames_from_video(str(video_file), max_frames)
                if frames:
                    break

    return frames


def _extract_frames_from_video(
    video_path: str,
    max_frames: int,
) -> list[tuple[np.ndarray, str]]:
    """Extract specific frames from a video file.

    Tries decord first, falls back to OpenCV.
    """
    frames: list[tuple[np.ndarray, str]] = []

    try:
        frames = _extract_with_decord(video_path, max_frames)
    except Exception:
        try:
            frames = _extract_with_opencv(video_path, max_frames)
        except Exception:
            pass

    return frames


def _extract_with_decord(video_path: str, max_frames: int) -> list[tuple[np.ndarray, str]]:
    """Extract frames using decord."""
    from decord import VideoReader
    from PIL import Image

    vr = VideoReader(video_path)
    total = len(vr)
    sample_indices = _sample_frame_indices(total, max_frames)

    frames = []
    for frame_idx, frame_type in sample_indices:
        frame = vr[frame_idx].asnumpy()
        img = Image.fromarray(frame).resize((384, 384))
        frames.append((np.array(img), frame_type))
    return frames


def _extract_with_opencv(video_path: str, max_frames: int) -> list[tuple[np.ndarray, str]]:
    """Extract frames using OpenCV."""
    import cv2
    from PIL import Image

    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    sample_indices = _sample_frame_indices(total, max_frames)

    frames = []
    for frame_idx, frame_type in sample_indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
        ret, frame = cap.read()
        if ret:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb).resize((384, 384))
            frames.append((np.array(img), frame_type))
    cap.release()
    return frames


def analyze_embeddings(result: EmbeddingResult) -> EmbeddingAnalysis:
    """Analyze visual embeddings for clustering, diversity, and outliers.

    Args:
        result: EmbeddingResult from compute_frame_embeddings.

    Returns:
        EmbeddingAnalysis with clusters, diversity score, and insights.
    """
    from sklearn.cluster import DBSCAN
    from sklearn.decomposition import PCA
    from sklearn.metrics.pairwise import cosine_distances
    from sklearn.neighbors import NearestNeighbors

    embeddings = result.embeddings
    n_samples = embeddings.shape[0]

    if n_samples == 0:
        return EmbeddingAnalysis(
            n_clusters=0,
            cluster_labels=np.array([]),
            cluster_sizes=[],
            pca_2d=np.empty((0, 2)),
            diversity_score=0.0,
            outlier_indices=[],
            insights=["No embeddings to analyze."],
        )

    # --- Dimensionality reduction ---
    pca_dims = min(50, n_samples, embeddings.shape[1])
    pca_50 = PCA(n_components=pca_dims)
    embeddings_50 = pca_50.fit_transform(embeddings)

    pca_2d_model = PCA(n_components=min(2, pca_dims))
    pca_2d = pca_2d_model.fit_transform(embeddings_50)

    # --- Clustering with DBSCAN ---
    # Auto-tune eps using k-distance heuristic
    k = min(5, n_samples - 1)
    if k < 1:
        k = 1
    nn = NearestNeighbors(n_neighbors=k)
    nn.fit(embeddings_50)
    distances, _ = nn.kneighbors(embeddings_50)
    k_distances = np.sort(distances[:, -1])

    # Use the "knee" of the k-distance curve
    # Simple heuristic: use the distance at 90th percentile
    eps = float(np.percentile(k_distances, 90))
    if eps < 1e-8:
        eps = 1.0

    dbscan = DBSCAN(eps=eps, min_samples=max(2, k))
    cluster_labels = dbscan.fit_predict(embeddings_50)

    # Compute cluster stats (excluding noise label -1)
    unique_labels = set(cluster_labels)
    unique_labels.discard(-1)
    n_clusters = len(unique_labels)
    cluster_sizes = [int(np.sum(cluster_labels == label)) for label in sorted(unique_labels)]

    # --- Diversity score ---
    # Compute per-episode centroids and mean pairwise cosine distance
    unique_episodes = sorted(set(result.episode_indices))
    episode_centroids = []
    for ep_idx in unique_episodes:
        mask = [i for i, e in enumerate(result.episode_indices) if e == ep_idx]
        if mask:
            centroid = embeddings[mask].mean(axis=0)
            episode_centroids.append(centroid)

    if len(episode_centroids) >= 2:
        centroids_arr = np.array(episode_centroids)
        cos_dists = cosine_distances(centroids_arr)
        # Extract upper triangle
        triu_indices = np.triu_indices_from(cos_dists, k=1)
        pairwise_dists = cos_dists[triu_indices]
        # Cosine distance is in [0, 2], normalize to [0, 1]
        diversity_score = float(np.mean(pairwise_dists) / 2.0)
        diversity_score = max(0.0, min(1.0, diversity_score))
    else:
        diversity_score = 0.0

    # --- Outlier detection ---
    # Episodes where all frames are > 2 std devs from nearest cluster centroid
    outlier_indices: list[int] = []

    if n_clusters > 0 and len(unique_episodes) > 1:
        # Compute cluster centroids in PCA-50 space
        cluster_centroids = []
        for label in sorted(unique_labels):
            mask = cluster_labels == label
            cluster_centroids.append(embeddings_50[mask].mean(axis=0))
        cluster_centroids_arr = np.array(cluster_centroids)

        # Compute distances from each point to nearest cluster centroid
        from sklearn.metrics.pairwise import euclidean_distances

        all_dists = euclidean_distances(embeddings_50, cluster_centroids_arr)
        min_dists = all_dists.min(axis=1)

        mean_dist = np.mean(min_dists)
        std_dist = np.std(min_dists)
        threshold = mean_dist + 2 * std_dist

        for ep_idx in unique_episodes:
            ep_mask = [i for i, e in enumerate(result.episode_indices) if e == ep_idx]
            if ep_mask and all(min_dists[i] > threshold for i in ep_mask):
                outlier_indices.append(ep_idx)

    # --- Generate insights ---
    insights = _generate_insights(
        n_clusters,
        cluster_sizes,
        diversity_score,
        outlier_indices,
        n_samples,
        unique_episodes,
        cluster_labels,
    )

    return EmbeddingAnalysis(
        n_clusters=n_clusters,
        cluster_labels=cluster_labels,
        cluster_sizes=cluster_sizes,
        pca_2d=pca_2d,
        diversity_score=diversity_score,
        outlier_indices=outlier_indices,
        insights=insights,
    )


def _generate_insights(
    n_clusters: int,
    cluster_sizes: list[int],
    diversity_score: float,
    outlier_indices: list[int],
    n_samples: int,
    unique_episodes: list[int],
    cluster_labels: np.ndarray,
) -> list[str]:
    """Generate human-readable insights from embedding analysis."""
    insights: list[str] = []
    n_episodes = len(unique_episodes)
    noise_count = int(np.sum(cluster_labels == -1)) if len(cluster_labels) > 0 else 0

    if n_clusters == 0:
        insights.append("No clear clusters detected — your data may be very uniform or too noisy")
    elif n_clusters == 1:
        if diversity_score >= 0.5:
            insights.append(
                f"Your episodes form 1 cluster with good internal diversity "
                f"({diversity_score:.2f}) — uniform coverage"
            )
        else:
            insights.append(
                f"Your episodes form 1 cluster with low diversity "
                f"({diversity_score:.2f}) — consider adding more variety"
            )
    elif n_clusters <= 4:
        insights.append(f"Your data has {n_clusters} distinct visual strategies — good variety")
    else:
        insights.append(
            f"Your data has {n_clusters} clusters — may indicate inconsistent recording conditions"
        )

    if outlier_indices:
        ep_str = ", ".join(str(e) for e in outlier_indices[:10])
        insights.append(f"Episodes {ep_str} appear to be outliers — review for recording errors")

    # Small cluster warning
    for i, size in enumerate(cluster_sizes):
        # Compute episode count in this cluster
        if size < 5 and n_episodes > 10:
            insights.append(
                f"Cluster {i + 1} has only {size} frames — "
                f"collect more similar episodes for balance"
            )

    if noise_count > n_samples * 0.2:
        insights.append(
            f"{noise_count} frames ({noise_count / n_samples:.0%}) are noise points "
            f"— data may lack coherent structure"
        )

    return insights
