"""Tests for the Sakr et al. consistency metrics module."""

import numpy as np
import pytest

from orbit.analyzer.consistency_metrics import (
    ConsistencyReport,
    _compute_curvature,
    _compute_effort,
    _compute_jerk,
    _compute_joint_limit_proximity,
    _compute_orientation_path_length,
    _compute_path_length,
    _cv,
    compute_consistency_metrics,
)


# ---------------------------------------------------------------------------
# Helper: generate synthetic episodes
# ---------------------------------------------------------------------------


def _make_smooth_episode(n_steps: int = 50, n_dims: int = 6, seed: int = 0) -> np.ndarray:
    """Generate a smooth sinusoidal episode."""
    rng = np.random.RandomState(seed)
    t = np.linspace(0, 2 * np.pi, n_steps)
    base = np.column_stack([np.sin(t + i * 0.5) for i in range(n_dims)])
    noise = rng.randn(n_steps, n_dims) * 0.01
    return base + noise


def _make_jerky_episode(n_steps: int = 50, n_dims: int = 6, seed: int = 0) -> np.ndarray:
    """Generate a jerky/noisy episode."""
    rng = np.random.RandomState(seed)
    return rng.randn(n_steps, n_dims)


def _make_consistent_dataset(n_episodes: int = 10, n_steps: int = 50, n_dims: int = 6) -> list[np.ndarray]:
    """Generate a dataset of similar smooth episodes (low CV)."""
    return [_make_smooth_episode(n_steps, n_dims, seed=i) for i in range(n_episodes)]


def _make_inconsistent_dataset(n_episodes: int = 10, n_steps: int = 50, n_dims: int = 6) -> list[np.ndarray]:
    """Generate a dataset with varied episodes (high CV)."""
    episodes = []
    for i in range(n_episodes):
        # Vary amplitude and frequency across episodes
        rng = np.random.RandomState(i)
        amplitude = rng.uniform(0.5, 5.0)
        freq = rng.uniform(0.5, 3.0)
        t = np.linspace(0, freq * 2 * np.pi, n_steps)
        base = amplitude * np.column_stack([np.sin(t + j) for j in range(n_dims)])
        noise = rng.randn(n_steps, n_dims) * 0.1
        episodes.append(base + noise)
    return episodes


# ---------------------------------------------------------------------------
# Per-episode metric tests
# ---------------------------------------------------------------------------


class TestPathLength:
    def test_stationary_episode(self):
        ep = np.ones((50, 6))
        assert _compute_path_length(ep) == pytest.approx(0.0, abs=1e-8)

    def test_moving_episode(self):
        ep = np.zeros((50, 2))
        ep[:, 0] = np.linspace(0, 10, 50)
        pl = _compute_path_length(ep)
        assert pl > 0

    def test_single_step(self):
        ep = np.ones((1, 3))
        assert _compute_path_length(ep) == 0.0


class TestJerk:
    def test_smooth_low_jerk(self):
        ep = _make_smooth_episode(100, 3)
        jerk_val = _compute_jerk(ep)
        assert jerk_val >= 0

    def test_jerky_high_jerk(self):
        smooth = _make_smooth_episode(100, 3)
        jerky = _make_jerky_episode(100, 3)
        assert _compute_jerk(jerky) > _compute_jerk(smooth)

    def test_too_short(self):
        ep = np.ones((3, 2))
        assert _compute_jerk(ep) == 0.0


class TestCurvature:
    def test_straight_line_low_curvature(self):
        ep = np.zeros((50, 3))
        ep[:, 0] = np.linspace(0, 10, 50)
        curv = _compute_curvature(ep)
        assert curv < 1.0  # Nearly straight

    def test_circle_higher_curvature(self):
        t = np.linspace(0, 2 * np.pi, 100)
        ep = np.column_stack([np.cos(t), np.sin(t), np.zeros(100)])
        curv = _compute_curvature(ep)
        assert curv > 0

    def test_1d_returns_zero(self):
        ep = np.ones((50, 1))
        assert _compute_curvature(ep) == 0.0


class TestEffort:
    def test_zero_actions(self):
        ep = np.zeros((50, 6))
        assert _compute_effort(ep) == 0.0

    def test_large_actions_more_effort(self):
        small = np.ones((50, 6)) * 0.1
        large = np.ones((50, 6)) * 10.0
        assert _compute_effort(large) > _compute_effort(small)


class TestJointLimitProximity:
    def test_centered_trajectory(self):
        """Actions near center of range should score high."""
        ep = np.zeros((50, 3)) + 0.5
        ep[:, 0] = np.linspace(0, 1, 50)  # Give some range
        val = _compute_joint_limit_proximity(ep)
        assert val > 0

    def test_dead_dimensions_skipped(self):
        ep = np.zeros((50, 3))  # All zeros, zero range
        val = _compute_joint_limit_proximity(ep)
        assert val == 0.0


class TestOrientationPathLength:
    def test_no_orientation(self):
        ep = np.ones((50, 3))  # Only 3 dims, no orientation
        assert _compute_orientation_path_length(ep) == 0.0

    def test_with_orientation(self):
        ep = np.zeros((50, 7))
        ep[:, -3:] = np.linspace(0, np.pi, 50)[:, None]  # Rotating
        val = _compute_orientation_path_length(ep)
        assert val > 0


# ---------------------------------------------------------------------------
# CV computation
# ---------------------------------------------------------------------------


class TestCV:
    def test_identical_values(self):
        assert _cv([5.0, 5.0, 5.0]) == 0.0

    def test_varied_values(self):
        cv = _cv([1.0, 2.0, 3.0, 4.0, 5.0])
        assert cv > 0

    def test_single_value(self):
        assert _cv([42.0]) == 0.0

    def test_near_zero_mean(self):
        assert _cv([0.0, 0.0, 0.0]) == 0.0


# ---------------------------------------------------------------------------
# Main compute_consistency_metrics
# ---------------------------------------------------------------------------


class TestConsistencyMetrics:
    def test_consistent_dataset_high_score(self):
        episodes = _make_consistent_dataset(20)
        report = compute_consistency_metrics(episodes)

        assert isinstance(report, ConsistencyReport)
        assert report.overall_consistency > 0.5

    def test_inconsistent_dataset_low_score(self):
        consistent = _make_consistent_dataset(20)
        inconsistent = _make_inconsistent_dataset(20)

        r_con = compute_consistency_metrics(consistent)
        r_inc = compute_consistency_metrics(inconsistent)

        assert r_con.overall_consistency > r_inc.overall_consistency

    def test_overall_consistency_range(self):
        episodes = _make_consistent_dataset(50)
        report = compute_consistency_metrics(episodes)

        assert 0.0 <= report.overall_consistency <= 1.0

    def test_per_episode_scores(self):
        episodes = _make_consistent_dataset(10)
        report = compute_consistency_metrics(episodes)

        assert len(report.per_episode_scores) == 10
        for i, score in report.per_episode_scores.items():
            assert 0.0 <= score <= 1.0

    def test_single_episode(self):
        episodes = [_make_smooth_episode(50, 6)]
        report = compute_consistency_metrics(episodes)
        assert report.overall_consistency == 0.5  # Not enough data

    def test_two_episodes(self):
        episodes = _make_consistent_dataset(2)
        report = compute_consistency_metrics(episodes)
        assert report.overall_consistency > 0

    def test_with_state_episodes(self):
        actions = _make_consistent_dataset(10, n_dims=6)
        states = _make_consistent_dataset(10, n_dims=12)
        report = compute_consistency_metrics(actions, state_episodes=states)
        assert report.overall_consistency > 0

    def test_1d_episodes(self):
        episodes = [np.random.randn(50) for _ in range(10)]
        report = compute_consistency_metrics(episodes)
        assert isinstance(report, ConsistencyReport)

    def test_per_episode_metrics_stored(self):
        episodes = _make_consistent_dataset(10)
        report = compute_consistency_metrics(episodes)
        assert "path_length" in report.per_episode_metrics
        assert len(report.per_episode_metrics["path_length"]) == 10

    def test_many_episodes_high_consistency(self):
        episodes = _make_consistent_dataset(50)
        report = compute_consistency_metrics(episodes)
        assert report.overall_consistency > 0.5

    def test_few_episodes_still_works(self):
        episodes = _make_consistent_dataset(5)
        report = compute_consistency_metrics(episodes)
        assert 0.0 <= report.overall_consistency <= 1.0
