"""Tests for the real-time collection coach."""

import numpy as np
import pytest

from orbit.coach import (
    CoachFeedback,
    CoachState,
    _compute_baseline,
    _compute_episode_metrics,
    _detect_fatigue,
    _score_episode,
)


def _smooth_episode(length=100, action_dim=6):
    """Generate a smooth sine-wave episode."""
    t = np.linspace(0, 2 * np.pi, length)
    return np.column_stack([np.sin(t + i * 0.5) * 0.5 for i in range(action_dim)])


def _jerky_episode(length=100, action_dim=6):
    """Generate a jerky, noisy episode."""
    t = np.linspace(0, 2 * np.pi, length)
    base = np.column_stack([np.sin(t + i * 0.5) * 0.5 for i in range(action_dim)])
    noise = np.random.RandomState(42).randn(*base.shape) * 2.0
    return base + noise


class TestEpisodeMetrics:
    def test_smooth_episode_low_jerk(self):
        ep = _smooth_episode()
        metrics = _compute_episode_metrics(ep)
        assert "jerk" in metrics
        assert "path_length" in metrics
        assert "length" in metrics
        assert metrics["jerk"] >= 0

    def test_jerky_episode_high_jerk(self):
        smooth = _smooth_episode()
        jerky = _jerky_episode()
        m_smooth = _compute_episode_metrics(smooth)
        m_jerky = _compute_episode_metrics(jerky)
        assert m_jerky["jerk"] > m_smooth["jerk"]

    def test_short_episode_handles_gracefully(self):
        ep = np.array([[0.1, 0.2], [0.3, 0.4]])  # Only 2 frames
        metrics = _compute_episode_metrics(ep)
        assert metrics == {}  # Too short for meaningful metrics

    def test_1d_episode(self):
        ep = np.array([0.1, 0.2, 0.3, 0.4])  # 1D
        metrics = _compute_episode_metrics(ep)
        assert metrics == {}


class TestBaseline:
    def test_baseline_from_smooth_episodes(self):
        episodes = [_smooth_episode() for _ in range(10)]
        baseline = _compute_baseline(episodes)
        assert "jerk" in baseline
        assert "path_length" in baseline
        assert baseline["jerk"]["median"] > 0
        assert baseline["jerk"]["mad"] > 0

    def test_baseline_empty_list(self):
        baseline = _compute_baseline([])
        assert baseline == {}


class TestScoreEpisode:
    def test_good_episode_scores_good(self):
        episodes = [_smooth_episode() for _ in range(10)]
        baseline = _compute_baseline(episodes)
        new_ep = _smooth_episode()
        feedback = _score_episode(new_ep, baseline)
        assert isinstance(feedback, CoachFeedback)
        assert feedback.verdict == "good"
        assert feedback.consistency_score > 0.5

    def test_jerky_episode_flagged(self):
        episodes = [_smooth_episode() for _ in range(10)]
        baseline = _compute_baseline(episodes)
        jerky = _jerky_episode()
        feedback = _score_episode(jerky, baseline)
        assert feedback.verdict in ("review", "redo")
        assert len(feedback.issues) > 0

    def test_consistency_score_range(self):
        episodes = [_smooth_episode() for _ in range(10)]
        baseline = _compute_baseline(episodes)
        feedback = _score_episode(_smooth_episode(), baseline)
        assert 0.0 <= feedback.consistency_score <= 1.0


class TestFatigueDetection:
    def test_no_fatigue_with_few_episodes(self):
        state = CoachState(target_episodes=50, policy="act")
        state.recent_jerk_trend = [0.1, 0.1, 0.1]
        assert _detect_fatigue(state) is None

    def test_fatigue_detected_with_increasing_jerk(self):
        state = CoachState(target_episodes=50, policy="act")
        # First 5 episodes: low jerk, last 5: high jerk
        state.recent_jerk_trend = [0.1, 0.1, 0.1, 0.1, 0.1, 0.5, 0.6, 0.7, 0.8, 0.9]
        msg = _detect_fatigue(state)
        assert msg is not None
        assert "break" in msg.lower() or "fatigue" in msg.lower()

    def test_no_fatigue_with_consistent_jerk(self):
        state = CoachState(target_episodes=50, policy="act")
        state.recent_jerk_trend = [0.1] * 10
        assert _detect_fatigue(state) is None
