"""Tests for orbit.analyzer.action_divergence."""

from __future__ import annotations

import numpy as np
import pytest

from orbit.analyzer.action_divergence import (
    ActionDivergenceResult,
    compute_action_divergence,
)


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


def _consistent_episodes(
    num_episodes: int = 10,
    timesteps: int = 100,
    action_dim: int = 6,
) -> list[np.ndarray]:
    """Create consistent episodes: identical trajectory with tiny noise.

    All episodes follow the exact same path so that when states are
    similar, the actions taken are also similar (low divergence).
    """
    rng = np.random.RandomState(42)
    base = np.sin(np.linspace(0, 2 * np.pi, timesteps)).reshape(-1, 1)
    base = np.tile(base, (1, action_dim))
    # Very small noise to avoid identical data but keep consistent behavior
    return [base.copy() + rng.randn(timesteps, action_dim) * 0.001 for _ in range(num_episodes)]


def _inconsistent_episodes(
    num_episodes: int = 10,
    timesteps: int = 100,
    action_dim: int = 6,
) -> list[np.ndarray]:
    """Create inconsistent episodes: same states but random actions."""
    rng = np.random.RandomState(42)
    return [rng.randn(timesteps, action_dim) * 2.0 for _ in range(num_episodes)]


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


class TestActionDivergenceBasic:
    """Basic functionality tests."""

    def test_consistent_data_low_divergence(self):
        """Consistent demonstrations should have low divergence.

        Creates episodes where actions are deterministic functions of state
        plus tiny noise. Every time the robot is in state X, it takes the
        same action — this is what "consistent" means for divergence.
        """
        rng = np.random.RandomState(42)
        action_eps = []
        state_eps = []
        for _ in range(10):
            # States follow a trajectory
            t = np.linspace(0, 2 * np.pi, 100)
            states = np.column_stack([np.sin(t + d) for d in range(6)])
            # Actions are deterministic function of state + tiny noise
            # (action = cos of state, i.e. the derivative)
            actions = np.column_stack([np.cos(t + d) for d in range(6)])
            actions += rng.randn(100, 6) * 0.001
            action_eps.append(actions)
            state_eps.append(states)

        result = compute_action_divergence(action_eps, state_episodes=state_eps)

        assert result.score is not None
        assert result.score < 0.5, f"Expected low divergence, got {result.score}"

    def test_inconsistent_data_high_divergence(self):
        """Random actions should have high divergence."""
        episodes = _inconsistent_episodes(num_episodes=10, timesteps=100)
        result = compute_action_divergence(episodes)

        assert result.score is not None
        assert result.score > 0.2, f"Expected higher divergence, got {result.score}"

    def test_empty_episodes(self):
        """Empty input returns None score."""
        result = compute_action_divergence([])
        assert result.score is None
        assert result.severity == "unknown"

    def test_single_episode(self):
        """Single episode returns None (need >=2 for divergence)."""
        rng = np.random.RandomState(42)
        episodes = [rng.randn(100, 6)]
        result = compute_action_divergence(episodes)
        assert result.score is None
        assert "1 episode" in (result.note or "")

    def test_with_state_episodes(self):
        """Works when state episodes are provided."""
        rng = np.random.RandomState(42)
        action_eps = [rng.randn(100, 6) for _ in range(5)]
        state_eps = [rng.randn(100, 6) for _ in range(5)]
        result = compute_action_divergence(action_eps, state_episodes=state_eps)
        assert result.score is not None

    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_action_divergence(episodes)
        assert result.score is not None

    def test_very_short_episodes(self):
        """Very short episodes still work."""
        rng = np.random.RandomState(42)
        episodes = [rng.randn(3, 4) for _ in range(5)]
        result = compute_action_divergence(episodes)
        # May return None if not enough frames
        assert isinstance(result, ActionDivergenceResult)


class TestActionDivergenceOutput:
    """Test output structure and constraints."""

    def test_score_range(self):
        """Score is in [0, 1]."""
        episodes = _consistent_episodes()
        result = compute_action_divergence(episodes)
        assert result.score is not None
        assert 0.0 <= result.score <= 1.0

    def test_penalty_range(self):
        """Penalty is in [0, 0.20]."""
        episodes = _inconsistent_episodes()
        result = compute_action_divergence(episodes)
        assert 0.0 <= result.estimated_success_penalty <= 0.20

    def test_worst_clusters_limited(self):
        """At most 5 worst clusters."""
        episodes = _inconsistent_episodes(num_episodes=20)
        result = compute_action_divergence(episodes)
        assert len(result.worst_clusters) <= 5

    def test_severity_values(self):
        """Severity is a valid string."""
        episodes = _consistent_episodes()
        result = compute_action_divergence(episodes)
        assert result.severity in ("low", "moderate", "high", "critical", "unknown")

    def test_recommendation_not_empty(self):
        """Recommendation is always provided."""
        episodes = _consistent_episodes()
        result = compute_action_divergence(episodes)
        assert result.recommendation
        assert len(result.recommendation) > 10


class TestDivergencePenaltyCorrelation:
    """Test that divergence correlates with penalty."""

    def test_low_divergence_low_penalty(self):
        """Low divergence should have low penalty."""
        episodes = _consistent_episodes()
        result = compute_action_divergence(episodes)
        if result.score is not None and result.score < 0.15:
            assert result.estimated_success_penalty < 0.05

    def test_penalty_increases_with_score(self):
        """Higher divergence scores should generally have higher penalties."""
        consistent = compute_action_divergence(_consistent_episodes())
        inconsistent = compute_action_divergence(_inconsistent_episodes())

        if consistent.score is not None and inconsistent.score is not None:
            if inconsistent.score > consistent.score:
                assert inconsistent.estimated_success_penalty >= consistent.estimated_success_penalty
