"""Tests for orbit.analyzer.embeddings."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.analyzer.embeddings import (
    EmbeddingAnalysis,
    EmbeddingResult,
    _generate_insights,
    _repo_cache_dir,
    _sample_episode_indices,
    _sample_frame_indices,
    analyze_embeddings,
)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_info(num_episodes: int = 50) -> DatasetInfo:
    return DatasetInfo(
        repo_id="test/dataset",
        num_episodes=num_episodes,
        num_frames=num_episodes * 250,
        fps=30.0,
        robot_type="aloha",
        camera_names=["cam_high"],
        video_keys=["observation.images.cam_high"],
        shapes={"action": [14]},
        duration_seconds=num_episodes * 250 / 30.0,
    )


def _make_embedding_result(
    n_episodes: int = 20,
    frames_per_ep: int = 4,
    dim: int = 768,
    n_clusters: int = 1,
    n_outliers: int = 0,
    seed: int = 42,
) -> EmbeddingResult:
    """Create synthetic EmbeddingResult with controllable clustering."""
    rng = np.random.RandomState(seed)
    embeddings = []
    episode_indices = []
    frame_types = []
    types = ["first", "mid_early", "mid_late", "last"]

    eps_per_cluster = max(1, (n_episodes - n_outliers) // max(1, n_clusters))

    for ep_idx in range(n_episodes):
        if ep_idx >= n_episodes - n_outliers:
            # Outlier episodes — far from any cluster
            center = rng.randn(dim) * 10 + 20
        else:
            cluster_id = ep_idx // eps_per_cluster
            center = rng.randn(dim) * 0.1 + cluster_id * 3
        for f in range(frames_per_ep):
            emb = center + rng.randn(dim) * 0.05
            embeddings.append(emb)
            episode_indices.append(ep_idx)
            frame_types.append(types[f % len(types)])

    return EmbeddingResult(
        embeddings=np.array(embeddings),
        episode_indices=episode_indices,
        frame_types=frame_types,
        metadata={"model_name": "test", "processing_time_seconds": 1.0},
    )


# ---------------------------------------------------------------------------
# _sample_frame_indices
# ---------------------------------------------------------------------------


class TestSampleFrameIndices:
    def test_empty_video(self):
        assert _sample_frame_indices(0) == []

    def test_single_frame(self):
        result = _sample_frame_indices(1)
        assert result == [(0, "first")]

    def test_four_frames(self):
        result = _sample_frame_indices(100)
        assert len(result) == 4
        types = [t for _, t in result]
        assert types == ["first", "mid_early", "mid_late", "last"]
        assert result[0][0] == 0
        assert result[-1][0] == 99

    def test_no_duplicate_indices(self):
        result = _sample_frame_indices(3)
        indices = [i for i, _ in result]
        assert len(indices) == len(set(indices))

    def test_two_frame_video(self):
        result = _sample_frame_indices(2)
        indices = [i for i, _ in result]
        assert 0 in indices
        assert 1 in indices


# ---------------------------------------------------------------------------
# _sample_episode_indices
# ---------------------------------------------------------------------------


class TestSampleEpisodeIndices:
    def test_fewer_than_max(self):
        assert _sample_episode_indices(10, 100) == list(range(10))

    def test_exact_max(self):
        assert _sample_episode_indices(100, 100) == list(range(100))

    def test_more_than_max(self):
        result = _sample_episode_indices(200, 50)
        assert len(result) == 50
        assert result[0] == 0
        # Should be evenly spaced
        assert all(isinstance(i, int) for i in result)


# ---------------------------------------------------------------------------
# analyze_embeddings — synthetic data
# ---------------------------------------------------------------------------


class TestAnalyzeEmbeddings:
    def test_empty_embeddings(self):
        result = EmbeddingResult(
            embeddings=np.empty((0, 768)),
            episode_indices=[],
            frame_types=[],
            metadata={},
        )
        analysis = analyze_embeddings(result)
        assert isinstance(analysis, EmbeddingAnalysis)
        assert analysis.n_clusters == 0
        assert analysis.diversity_score == 0.0
        assert len(analysis.insights) > 0

    def test_single_cluster(self):
        result = _make_embedding_result(n_episodes=20, n_clusters=1)
        analysis = analyze_embeddings(result)
        assert isinstance(analysis, EmbeddingAnalysis)
        assert analysis.n_clusters >= 1
        assert 0.0 <= analysis.diversity_score <= 1.0
        assert analysis.pca_2d.shape == (80, 2)  # 20 eps * 4 frames

    def test_multiple_clusters(self):
        result = _make_embedding_result(n_episodes=30, n_clusters=3, dim=128)
        analysis = analyze_embeddings(result)
        assert analysis.n_clusters >= 1
        assert 0.0 <= analysis.diversity_score <= 1.0

    def test_with_outliers(self):
        result = _make_embedding_result(
            n_episodes=25, n_clusters=2, n_outliers=3, dim=128
        )
        analysis = analyze_embeddings(result)
        assert isinstance(analysis.outlier_indices, list)
        # Outliers should be detected (episodes 22, 23, 24)
        # DBSCAN may or may not detect them depending on eps tuning
        assert 0.0 <= analysis.diversity_score <= 1.0

    def test_diversity_score_range(self):
        """Diversity score must always be in [0, 1]."""
        for n_clusters in [1, 2, 5]:
            result = _make_embedding_result(
                n_episodes=20, n_clusters=n_clusters, dim=64
            )
            analysis = analyze_embeddings(result)
            assert 0.0 <= analysis.diversity_score <= 1.0, (
                f"Diversity {analysis.diversity_score} out of range "
                f"for {n_clusters} clusters"
            )

    def test_pca_2d_shape(self):
        result = _make_embedding_result(n_episodes=10, frames_per_ep=4, dim=64)
        analysis = analyze_embeddings(result)
        assert analysis.pca_2d.shape == (40, 2)

    def test_cluster_sizes_sum(self):
        """Cluster sizes should sum to <= total (noise excluded)."""
        result = _make_embedding_result(n_episodes=20, n_clusters=2, dim=128)
        analysis = analyze_embeddings(result)
        total_clustered = sum(analysis.cluster_sizes)
        assert total_clustered <= result.embeddings.shape[0]

    def test_insights_generated(self):
        result = _make_embedding_result(n_episodes=20, n_clusters=2, dim=128)
        analysis = analyze_embeddings(result)
        assert isinstance(analysis.insights, list)
        assert len(analysis.insights) > 0

    def test_low_diversity_single_cluster(self):
        """Tight single cluster should have low diversity."""
        rng = np.random.RandomState(0)
        center = np.ones(64)
        n = 40
        embs = center + rng.randn(n, 64) * 0.001
        result = EmbeddingResult(
            embeddings=embs,
            episode_indices=[i // 4 for i in range(n)],
            frame_types=["first", "mid_early", "mid_late", "last"] * (n // 4),
            metadata={},
        )
        analysis = analyze_embeddings(result)
        assert analysis.diversity_score < 0.3


# ---------------------------------------------------------------------------
# _generate_insights
# ---------------------------------------------------------------------------


class TestGenerateInsights:
    def test_no_clusters(self):
        insights = _generate_insights(
            n_clusters=0,
            cluster_sizes=[],
            diversity_score=0.0,
            outlier_indices=[],
            n_samples=80,
            unique_episodes=list(range(20)),
            cluster_labels=np.full(80, -1),
        )
        assert any("no clear clusters" in i.lower() for i in insights)

    def test_single_cluster_high_diversity(self):
        insights = _generate_insights(
            n_clusters=1,
            cluster_sizes=[80],
            diversity_score=0.7,
            outlier_indices=[],
            n_samples=80,
            unique_episodes=list(range(20)),
            cluster_labels=np.zeros(80, dtype=int),
        )
        assert any("uniform coverage" in i.lower() for i in insights)

    def test_single_cluster_low_diversity(self):
        insights = _generate_insights(
            n_clusters=1,
            cluster_sizes=[80],
            diversity_score=0.2,
            outlier_indices=[],
            n_samples=80,
            unique_episodes=list(range(20)),
            cluster_labels=np.zeros(80, dtype=int),
        )
        assert any("more variety" in i.lower() for i in insights)

    def test_many_clusters(self):
        insights = _generate_insights(
            n_clusters=7,
            cluster_sizes=[10, 10, 10, 10, 10, 10, 10],
            diversity_score=0.8,
            outlier_indices=[],
            n_samples=70,
            unique_episodes=list(range(20)),
            cluster_labels=np.array([i // 10 for i in range(70)]),
        )
        assert any("inconsistent" in i.lower() for i in insights)

    def test_outlier_insight(self):
        insights = _generate_insights(
            n_clusters=2,
            cluster_sizes=[35, 35],
            diversity_score=0.5,
            outlier_indices=[12, 47],
            n_samples=80,
            unique_episodes=list(range(20)),
            cluster_labels=np.array([0] * 40 + [1] * 40),
        )
        assert any("outlier" in i.lower() for i in insights)
        assert any("12" in i for i in insights)


# ---------------------------------------------------------------------------
# Caching
# ---------------------------------------------------------------------------


class TestCaching:
    def test_repo_cache_dir(self):
        path = _repo_cache_dir("lerobot/pusht")
        assert isinstance(path, Path)
        assert ".cache" in str(path)
        assert "orbit" in str(path)

    def test_cache_dir_deterministic(self):
        a = _repo_cache_dir("lerobot/pusht")
        b = _repo_cache_dir("lerobot/pusht")
        assert a == b

    def test_different_repos_different_dirs(self):
        a = _repo_cache_dir("lerobot/pusht")
        b = _repo_cache_dir("lerobot/aloha")
        assert a != b

    def test_cache_roundtrip(self, tmp_path):
        """Test saving and loading embeddings from cache."""
        cache_file = tmp_path / "embeddings.npz"
        rng = np.random.RandomState(42)
        embeddings = rng.randn(20, 64)
        episode_indices = np.array([i // 4 for i in range(20)])
        frame_types = np.array(
            ["first", "mid_early", "mid_late", "last"] * 5
        )
        metadata = {"model_name": "test", "processing_time_seconds": 1.5}

        np.savez(
            cache_file,
            embeddings=embeddings,
            episode_indices=episode_indices,
            frame_types=frame_types,
            metadata=np.array(metadata),
        )

        loaded = np.load(cache_file, allow_pickle=True)
        np.testing.assert_array_almost_equal(loaded["embeddings"], embeddings)
        assert loaded["episode_indices"].tolist() == episode_indices.tolist()
        assert loaded["frame_types"].tolist() == frame_types.tolist()


# ---------------------------------------------------------------------------
# Integration test markers (skipped in unit test runs)
# ---------------------------------------------------------------------------


@pytest.mark.integration
class TestSigLIPIntegration:
    """Integration tests that require torch + transformers + real model.

    Run with: pytest -v -m integration
    """

    def test_compute_real_embeddings(self):
        """Smoke test with actual SigLIP model on random images."""
        import torch
        from transformers import AutoImageProcessor, AutoModel

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

        # Generate random images
        rng = np.random.RandomState(0)
        images = [rng.randint(0, 255, (384, 384, 3), dtype=np.uint8) for _ in range(4)]
        inputs = processor(images=images, return_tensors="pt")

        with torch.no_grad():
            outputs = model(**inputs)
            embs = outputs.pooler_output
            if embs is None:
                embs = outputs.last_hidden_state[:, 0]

        assert embs.shape[0] == 4
        assert embs.shape[1] > 0
