"""Tests for orbit.analyzer.episode_ranker."""

from __future__ import annotations

import numpy as np
import pytest

from orbit.analyzer.episode_ranker import (
    EpisodeRankingResult,
    compute_episode_ranking,
)


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


def _smooth_episodes(
    num_episodes: int = 10,
    timesteps: int = 100,
    action_dim: int = 6,
) -> list[np.ndarray]:
    """Create smooth, consistent sine-wave episodes."""
    t = np.linspace(0, 2 * np.pi, timesteps)
    episodes = []
    for i in range(num_episodes):
        ep = np.column_stack([
            np.sin(t + i * 0.2 + d * 0.3) for d in range(action_dim)
        ])
        episodes.append(ep)
    return episodes


def _jerky_episodes(
    num_episodes: int = 10,
    timesteps: int = 100,
    action_dim: int = 6,
) -> list[np.ndarray]:
    """Create jerky, inconsistent episodes with random spikes."""
    rng = np.random.RandomState(42)
    episodes = []
    for _ in range(num_episodes):
        ep = rng.randn(timesteps, action_dim) * 5.0
        # Add random spikes
        for _ in range(10):
            t = rng.randint(0, timesteps)
            ep[t] *= 10.0
        episodes.append(ep)
    return episodes


def _mixed_episodes(
    n_good: int = 8,
    n_bad: int = 2,
    timesteps: int = 100,
    action_dim: int = 6,
) -> list[np.ndarray]:
    """Mix of good smooth episodes and bad jerky ones."""
    good = _smooth_episodes(n_good, timesteps, action_dim)
    bad = _jerky_episodes(n_bad, timesteps, action_dim)
    return good + bad


# ===========================================================================
# BASIC TESTS
# ===========================================================================


class TestEpisodeRankerBasic:
    """Basic functionality tests."""

    def test_smooth_episodes_score_high(self):
        """Smooth episodes should score well."""
        episodes = _smooth_episodes()
        result = compute_episode_ranking(episodes)

        assert result.episodes
        avg_score = np.mean([e.composite_score for e in result.episodes])
        assert avg_score > 0.4, f"Expected good scores for smooth episodes, got {avg_score}"

    def test_empty_input(self):
        """Empty input returns empty result."""
        result = compute_episode_ranking([])
        assert result.keep_count == 0
        assert result.remove_count == 0
        assert result.episodes == []

    def test_single_episode(self):
        """Single episode is handled."""
        rng = np.random.RandomState(42)
        episodes = [np.sin(np.linspace(0, 2 * np.pi, 100)).reshape(-1, 1)]
        result = compute_episode_ranking(episodes)
        assert len(result.episodes) == 1

    def test_mixed_episodes_bad_ranked_lower(self):
        """Bad episodes should rank lower than good ones."""
        episodes = _mixed_episodes(n_good=8, n_bad=2)
        result = compute_episode_ranking(episodes)

        # Episodes are sorted by score (best first)
        scores = [e.composite_score for e in result.episodes]
        # The worst episodes should be at the end
        assert scores[0] >= scores[-1]

    def test_with_state_episodes(self):
        """Works when state episodes are provided."""
        action_eps = _smooth_episodes(5, 100, 6)
        state_eps = _smooth_episodes(5, 100, 6)
        result = compute_episode_ranking(action_eps, state_episodes=state_eps)
        assert len(result.episodes) == 5


class TestEpisodeRankerOutput:
    """Test output structure."""

    def test_counts_add_up(self):
        """Keep + review + remove = total."""
        episodes = _mixed_episodes()
        result = compute_episode_ranking(episodes)
        total = len(episodes)
        assert result.keep_count + result.review_count + result.remove_count == total

    def test_remove_indices_consistent(self):
        """remove_indices matches episodes with 'remove' verdict."""
        episodes = _mixed_episodes()
        result = compute_episode_ranking(episodes)
        remove_from_list = sorted([e.episode_index for e in result.episodes if e.verdict == "remove"])
        assert result.remove_indices == remove_from_list

    def test_scores_in_range(self):
        """All scores are in [0, 1]."""
        episodes = _smooth_episodes()
        result = compute_episode_ranking(episodes)
        for ep in result.episodes:
            assert 0.0 <= ep.composite_score <= 1.0
            assert 0.0 <= ep.smoothness <= 1.0
            assert 0.0 <= ep.consistency <= 1.0
            assert 0.0 <= ep.length_score <= 1.0
            assert 0.0 <= ep.range_utilization <= 1.0

    def test_verdict_values(self):
        """Verdicts are valid strings."""
        episodes = _mixed_episodes()
        result = compute_episode_ranking(episodes)
        for ep in result.episodes:
            assert ep.verdict in ("keep", "review", "remove")

    def test_improvement_non_negative_when_removing(self):
        """Removing bad episodes should not decrease quality."""
        episodes = _mixed_episodes(n_good=8, n_bad=4)
        result = compute_episode_ranking(episodes)
        if result.remove_count > 0:
            assert result.estimated_improvement >= 0

    def test_summary_not_empty(self):
        """Summary is always provided."""
        episodes = _smooth_episodes()
        result = compute_episode_ranking(episodes)
        assert result.summary
        assert len(result.summary) > 5


class TestEpisodeRankerEdgeCases:
    """Edge case tests."""

    def test_constant_actions(self):
        """Episodes with constant actions."""
        episodes = [np.zeros((100, 6)) for _ in range(5)]
        result = compute_episode_ranking(episodes)
        assert isinstance(result, EpisodeRankingResult)

    def test_1d_actions(self):
        """Works with 1D action arrays."""
        rng = np.random.RandomState(42)
        episodes = [rng.randn(100) for _ in range(5)]
        result = compute_episode_ranking(episodes)
        assert len(result.episodes) == 5

    def test_variable_lengths(self):
        """Episodes with different lengths."""
        rng = np.random.RandomState(42)
        episodes = [rng.randn(length, 4) for length in [50, 100, 150, 80, 120]]
        result = compute_episode_ranking(episodes)
        assert len(result.episodes) == 5

    def test_very_short_episodes(self):
        """Very short episodes (< 4 frames)."""
        rng = np.random.RandomState(42)
        episodes = [rng.randn(2, 4) for _ in range(5)]
        result = compute_episode_ranking(episodes)
        assert isinstance(result, EpisodeRankingResult)
