"""Tests for orbit.analyzer.temporal_alignment."""

from __future__ import annotations

import numpy as np
import pytest

from orbit.analyzer.temporal_alignment import (
    TemporalAlignmentResult,
    compute_temporal_alignment,
)


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


def _aligned_data(
    num_episodes: int = 5,
    timesteps: int = 200,
    action_dim: int = 6,
) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """Create perfectly aligned state-action pairs.

    Actions directly match state velocity (d_state/dt = action).
    """
    rng = np.random.RandomState(42)
    action_eps = []
    state_eps = []
    for _ in range(num_episodes):
        # State is a smooth trajectory; actions = state velocity
        t = np.linspace(0, 4 * np.pi, timesteps + 1).reshape(-1, 1)
        states_full = np.sin(t + rng.randn(1, action_dim) * 0.5) * 2.0
        states = states_full[:-1]
        actions = np.diff(states_full, axis=0)  # velocity = d(state)/dt
        action_eps.append(actions)
        state_eps.append(states)
    return action_eps, state_eps


def _misaligned_data(
    num_episodes: int = 5,
    timesteps: int = 200,
    action_dim: int = 6,
    lag: int = 2,
) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """Create misaligned state-action pairs with known lag.

    Actions are shifted by `lag` frames relative to states.
    """
    rng = np.random.RandomState(42)
    action_eps = []
    state_eps = []
    for _ in range(num_episodes):
        actions = np.sin(
            np.linspace(0, 4 * np.pi, timesteps + lag).reshape(-1, 1)
            + rng.randn(1, action_dim) * 0.5
        ) * 0.5
        states = np.cumsum(actions, axis=0)
        # Shift: state is delayed by `lag` frames
        action_eps.append(actions[:timesteps])
        state_eps.append(states[lag:lag + timesteps])
    return action_eps, state_eps


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


class TestTemporalAlignmentBasic:
    """Basic functionality tests."""

    def test_aligned_data_passes(self):
        """Aligned data should be detected as aligned."""
        actions, states = _aligned_data()
        result = compute_temporal_alignment(actions, state_episodes=states)

        assert result.aligned is True
        assert result.severity in ("ok", "mild")
        assert abs(result.estimated_lag) <= 1

    def test_empty_input(self):
        """Empty input returns aligned (no issues)."""
        result = compute_temporal_alignment([])
        assert result.aligned is True
        assert result.estimated_lag == 0

    def test_actions_only(self):
        """Works with only action episodes (no states)."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100, 6) for _ in range(5)]
        result = compute_temporal_alignment(actions)
        # Without states, lag detection is skipped
        assert result.estimated_lag == 0
        assert isinstance(result, TemporalAlignmentResult)

    def test_with_timestamps_stable_fps(self):
        """Stable timestamps give low jitter."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100, 6) for _ in range(3)]
        # 30 fps timestamps with very small jitter
        timestamps = [
            np.arange(100) / 30.0 + rng.randn(100) * 0.0001
            for _ in range(3)
        ]
        result = compute_temporal_alignment(actions, timestamps=timestamps)
        assert result.fps_jitter < 0.05


class TestTemporalAlignmentOutput:
    """Test output structure."""

    def test_severity_values(self):
        """Severity is a valid string."""
        actions, states = _aligned_data()
        result = compute_temporal_alignment(actions, state_episodes=states)
        assert result.severity in ("ok", "mild", "moderate", "severe")

    def test_recommendation_not_empty(self):
        """Recommendation is always provided."""
        actions, states = _aligned_data()
        result = compute_temporal_alignment(actions, state_episodes=states)
        assert result.recommendation
        assert len(result.recommendation) > 5


class TestConstantDimensions:
    """Test detection of constant state dimensions."""

    def test_detects_constant_dims(self):
        """Constant state dimensions are detected."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100, 4) for _ in range(5)]
        # States where dim 2 is always 0
        states = []
        for _ in range(5):
            s = rng.randn(100, 4)
            s[:, 2] = 0.0  # constant
            states.append(s)

        result = compute_temporal_alignment(actions, state_episodes=states)
        assert 2 in result.constant_dimensions

    def test_no_false_positives(self):
        """Varying state dimensions are not flagged."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100, 4) for _ in range(5)]
        states = [rng.randn(100, 4) for _ in range(5)]

        result = compute_temporal_alignment(actions, state_episodes=states)
        assert len(result.constant_dimensions) == 0


class TestFPSJitter:
    """Test FPS consistency detection."""

    def test_stable_fps(self):
        """Stable FPS gives low jitter."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100, 4) for _ in range(3)]
        timestamps = [np.arange(100) / 30.0 for _ in range(3)]

        result = compute_temporal_alignment(actions, timestamps=timestamps)
        assert result.fps_jitter < 0.01

    def test_jittery_fps(self):
        """Jittery FPS is detected."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100, 4) for _ in range(3)]
        # Very jittery timestamps
        timestamps = [
            np.cumsum(rng.exponential(1.0 / 30, 100))
            for _ in range(3)
        ]

        result = compute_temporal_alignment(actions, timestamps=timestamps)
        assert result.fps_jitter > 0.1

    def test_missing_frames_detected(self):
        """Missing frames (gaps in timestamps) are detected."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100, 4) for _ in range(3)]
        # Regular timestamps with some gaps
        timestamps = []
        for _ in range(3):
            ts = np.arange(100) / 30.0
            # Insert a gap at frame 50
            ts[50:] += 0.5  # half-second gap
            timestamps.append(ts)

        result = compute_temporal_alignment(actions, timestamps=timestamps)
        assert result.missing_frames > 0


class TestEdgeCases:
    """Edge case tests."""

    def test_single_episode(self):
        """Works with single episode."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100, 4)]
        states = [rng.randn(100, 4)]
        result = compute_temporal_alignment(actions, state_episodes=states)
        assert isinstance(result, TemporalAlignmentResult)

    def test_short_episodes(self):
        """Works with very short episodes."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(5, 4) for _ in range(3)]
        states = [rng.randn(5, 4) for _ in range(3)]
        result = compute_temporal_alignment(actions, state_episodes=states)
        assert isinstance(result, TemporalAlignmentResult)

    def test_mismatched_lengths(self):
        """Handles state/action length mismatches."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100, 4) for _ in range(3)]
        states = [rng.randn(80, 4) for _ in range(3)]  # shorter
        result = compute_temporal_alignment(actions, state_episodes=states)
        assert isinstance(result, TemporalAlignmentResult)

    def test_1d_arrays(self):
        """Works with 1D arrays."""
        rng = np.random.RandomState(42)
        actions = [rng.randn(100) for _ in range(3)]
        states = [rng.randn(100) for _ in range(3)]
        result = compute_temporal_alignment(actions, state_episodes=states)
        assert isinstance(result, TemporalAlignmentResult)
