"""Tests for the TaskContextAnalyzer."""

import numpy as np
import pytest

from orbit.analyzer.task_context import (
    TaskContextAnalyzer,
    TaskContext,
    _dimension_complexity,
    _temporal_complexity,
    _coordination_complexity,
    _detect_bimanual,
)


class TestDimensionComplexity:
    """Test action dimension -> complexity mapping."""

    def test_2d_is_simple(self):
        score = _dimension_complexity(2)
        assert score < 0.20

    def test_7d_is_moderate(self):
        score = _dimension_complexity(7)
        assert 0.30 <= score <= 0.50

    def test_14d_is_complex(self):
        score = _dimension_complexity(14)
        assert score >= 0.60

    def test_30d_is_very_complex(self):
        score = _dimension_complexity(30)
        assert score >= 0.90

    def test_monotonically_increasing(self):
        scores = [_dimension_complexity(d) for d in [2, 7, 14, 30]]
        for i in range(len(scores) - 1):
            assert scores[i] < scores[i + 1]


class TestTemporalComplexity:
    """Test episode length -> temporal complexity."""

    def test_short_consistent_is_simple(self):
        lengths = [50] * 20
        score = _temporal_complexity(lengths)
        assert score < 0.30

    def test_long_variable_is_complex(self):
        rng = np.random.RandomState(42)
        lengths = (rng.randn(30) * 100 + 400).astype(int).tolist()
        score = _temporal_complexity(lengths)
        assert score > 0.40

    def test_empty_returns_default(self):
        assert _temporal_complexity([]) == 0.5


class TestCoordinationComplexity:
    """Test multi-joint coordination scoring."""

    def test_single_joint_movement(self):
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(5):
            ep = np.zeros((100, 6))
            ep[:, 0] = rng.randn(100) * 0.1  # Only dim 0 moves
            episodes.append(ep)
        score = _coordination_complexity(episodes)
        assert score < 0.30

    def test_all_joints_move(self):
        rng = np.random.RandomState(42)
        episodes = [rng.randn(100, 6) * 0.5 for _ in range(5)]
        score = _coordination_complexity(episodes)
        assert score > 0.40

    def test_empty_returns_default(self):
        assert _coordination_complexity([]) == 0.5


class TestBimanualDetection:
    """Test bimanual detection heuristic."""

    def test_6d_not_bimanual(self):
        assert not _detect_bimanual(6)

    def test_7d_not_bimanual(self):
        assert not _detect_bimanual(7)

    def test_14d_is_bimanual(self):
        assert _detect_bimanual(14)

    def test_12d_is_bimanual(self):
        assert _detect_bimanual(12)


class TestTaskContextAnalyzer:
    """Test the full analyzer."""

    def setup_method(self):
        self.analyzer = TaskContextAnalyzer()
        self.rng = np.random.RandomState(42)

    def test_simple_2d_task(self):
        episodes = [self.rng.randn(50, 2) * 0.01 for _ in range(20)]
        ctx = self.analyzer.analyze(episodes)
        assert isinstance(ctx, TaskContext)
        assert ctx.action_dim == 2
        assert not ctx.is_bimanual
        assert ctx.estimated_task_type in ("simple", "moderate")

    def test_complex_bimanual_task(self):
        episodes = [self.rng.randn(300, 14) * 0.5 for _ in range(10)]
        ctx = self.analyzer.analyze(episodes)
        assert ctx.action_dim == 14
        assert ctx.is_bimanual
        assert ctx.task_complexity_score > 0.4

    def test_empty_episodes(self):
        ctx = self.analyzer.analyze([])
        assert ctx.task_complexity_score == 0.5
        assert ctx.action_dim == 0

    def test_single_episode(self):
        episodes = [self.rng.randn(100, 6)]
        ctx = self.analyzer.analyze(episodes)
        assert ctx.action_dim == 6
        assert 0.0 <= ctx.task_complexity_score <= 1.0

    def test_all_scores_bounded(self):
        episodes = [self.rng.randn(100, 7) for _ in range(20)]
        ctx = self.analyzer.analyze(episodes)
        assert 0.0 <= ctx.task_complexity_score <= 1.0
        assert 0.0 <= ctx.dimension_complexity <= 1.0
        assert 0.0 <= ctx.temporal_complexity <= 1.0
        assert 0.0 <= ctx.coordination_complexity <= 1.0
        assert 0.0 <= ctx.object_interaction_complexity <= 1.0
        assert 0.0 <= ctx.trajectory_multimodality <= 1.0

    def test_metadata_override(self):
        episodes = [self.rng.randn(100, 6) for _ in range(10)]
        ctx = self.analyzer.analyze(episodes, metadata={"is_bimanual": True})
        assert ctx.is_bimanual

    def test_complexity_breakdown_present(self):
        episodes = [self.rng.randn(100, 7) for _ in range(10)]
        ctx = self.analyzer.analyze(episodes)
        assert "dimension" in ctx.complexity_breakdown
        assert "temporal" in ctx.complexity_breakdown
        assert "coordination" in ctx.complexity_breakdown
