"""Tests for orbit.analyzer.coverage."""

from __future__ import annotations

import numpy as np
import pandas as pd

from orbit.analyzer.coverage import FullCoverageReport, compute_coverage
from orbit.analyzer.dataset_loader import DatasetInfo


def _make_info(num_episodes=10, num_frames=200):
    return DatasetInfo(
        repo_id="test/dataset",
        num_episodes=num_episodes,
        num_frames=num_frames,
        fps=30.0,
        robot_type="test",
        camera_names=[],
        video_keys=[],
    )


class TestComputeCoverage:
    def test_empty_dataframe(self):
        df = pd.DataFrame()
        info = _make_info()
        result = compute_coverage(info, df)
        assert isinstance(result, FullCoverageReport)
        assert result.position_diversity == 0.0
        assert result.action_diversity == 0.0

    def test_basic_coverage(self):
        np.random.seed(42)
        n = 200
        df = pd.DataFrame({
            "episode_index": np.repeat(range(10), 20),
            "state": [np.random.randn(3).tolist() for _ in range(n)],
            "action": [np.random.randn(6).tolist() for _ in range(n)],
        })
        info = _make_info(num_episodes=10, num_frames=n)
        result = compute_coverage(info, df)

        assert 0 <= result.position_diversity <= 1
        assert 0 <= result.action_diversity <= 1
        assert 0 <= result.episode_consistency <= 1
        assert 0 <= result.temporal_coverage <= 1

    def test_uniform_actions_low_diversity(self):
        n = 100
        df = pd.DataFrame({
            "episode_index": np.repeat(range(5), 20),
            "state": [[0.0, 0.0, 0.0] for _ in range(n)],
            "action": [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0] for _ in range(n)],
        })
        info = _make_info(num_episodes=5, num_frames=n)
        result = compute_coverage(info, df)
        assert result.action_diversity == 0.0

    def test_consistent_episode_lengths(self):
        n = 100
        df = pd.DataFrame({
            "episode_index": np.repeat(range(5), 20),
            "state": [np.random.randn(3).tolist() for _ in range(n)],
            "action": [np.random.randn(6).tolist() for _ in range(n)],
        })
        info = _make_info(num_episodes=5, num_frames=n)
        result = compute_coverage(info, df)
        assert result.episode_consistency == 1.0
        assert result.outlier_episodes == []

    def test_no_state_column(self):
        n = 50
        df = pd.DataFrame({
            "episode_index": np.repeat(range(5), 10),
            "action": [np.random.randn(6).tolist() for _ in range(n)],
        })
        info = _make_info(num_episodes=5, num_frames=n)
        result = compute_coverage(info, df)
        assert result.position_diversity == 0.5  # default when missing

    def test_no_action_column(self):
        n = 50
        df = pd.DataFrame({
            "episode_index": np.repeat(range(5), 10),
            "state": [np.random.randn(3).tolist() for _ in range(n)],
        })
        info = _make_info(num_episodes=5, num_frames=n)
        result = compute_coverage(info, df)
        assert result.action_diversity == 0.5  # default when missing
