"""Tests for the DemInf-style data curator."""

import numpy as np
import pytest

from orbit.analyzer.data_curator import CurationResult, curate_episodes


def _make_episodes(n: int, n_dims: int = 6, n_steps: int = 50) -> list[np.ndarray]:
    """Generate synthetic episodes with one clear outlier at index 0."""
    episodes = []
    for i in range(n):
        rng = np.random.RandomState(i + 1)
        if i == 0:
            # Outlier: much larger amplitude
            ep = rng.randn(n_steps, n_dims) * 10.0
        else:
            ep = rng.randn(n_steps, n_dims) * 0.5
        episodes.append(ep)
    return episodes


def _make_uniform_episodes(n: int) -> list[np.ndarray]:
    """Generate similar episodes (no outliers)."""
    return [np.random.RandomState(i).randn(50, 6) * 0.5 + 1.0 for i in range(n)]


class TestCurateEpisodes:
    def test_basic_output(self):
        episodes = _make_episodes(20)
        result = curate_episodes(episodes)

        assert isinstance(result, CurationResult)
        assert len(result.episode_scores) == 20
        assert len(result.ranked_episodes) == 20
        assert result.method in ("consistency", "mi")

    def test_outlier_ranked_low(self):
        episodes = _make_episodes(20)
        result = curate_episodes(episodes, method="consistency")

        # Episode 0 (outlier) should have lower score than average
        avg_score = np.mean(list(result.episode_scores.values()))
        assert result.episode_scores[0] < avg_score

    def test_budget_keeps_n(self):
        episodes = _make_episodes(50)
        result = curate_episodes(episodes, budget=25)

        assert len(result.recommended_keep) == 25
        assert len(result.recommended_remove) == 25

    def test_budget_curve(self):
        episodes = _make_episodes(50)
        result = curate_episodes(episodes)

        assert len(result.budget_curve) > 0
        for n_eps, quality in result.budget_curve:
            assert n_eps > 0
            assert 0.0 <= quality <= 1.0

    def test_consistency_method(self):
        episodes = _make_episodes(30)
        result = curate_episodes(episodes, method="consistency")
        assert result.method == "consistency"

    def test_auto_method_small_dataset(self):
        episodes = _make_episodes(30)
        result = curate_episodes(episodes, method="auto")
        assert result.method == "consistency"  # <200 episodes → consistency

    def test_empty_dataset(self):
        result = curate_episodes([])
        assert result.method == "none"
        assert len(result.ranked_episodes) == 0

    def test_single_episode(self):
        episodes = [np.random.randn(50, 6)]
        result = curate_episodes(episodes)
        assert len(result.ranked_episodes) == 1

    def test_uniform_dataset_no_removals(self):
        episodes = _make_uniform_episodes(20)
        result = curate_episodes(episodes)
        # Uniform data should have few/no removals
        assert len(result.recommended_remove) <= 2

    def test_scores_between_0_and_1(self):
        episodes = _make_episodes(30)
        result = curate_episodes(episodes, method="consistency")
        for score in result.episode_scores.values():
            assert 0.0 <= score <= 1.0

    def test_ranked_order_matches_scores(self):
        episodes = _make_episodes(20)
        result = curate_episodes(episodes)

        # Ranked order should be descending by score
        for i in range(len(result.ranked_episodes) - 1):
            idx_a = result.ranked_episodes[i]
            idx_b = result.ranked_episodes[i + 1]
            assert result.episode_scores[idx_a] >= result.episode_scores[idx_b] - 1e-8

    def test_with_state_episodes(self):
        actions = _make_episodes(20)
        states = [np.random.randn(50, 12) for _ in range(20)]
        result = curate_episodes(actions, state_episodes=states)
        assert len(result.episode_scores) == 20

    def test_1d_episodes(self):
        episodes = [np.random.randn(50) for _ in range(10)]
        result = curate_episodes(episodes, method="consistency")
        assert len(result.episode_scores) == 10
