"""Tests for ported quality metrics in orbit.analyzer.quality."""

from __future__ import annotations

import numpy as np
import pytest

from orbit.analyzer.coverage import FullCoverageReport
from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.analyzer.quality import (
    compute_action_smoothness,
    compute_action_state_mi,
    compute_episode_completion,
    compute_observation_consistency,
    compute_quality_score,
    compute_spectral_smoothness,
    compute_temporal_autocorrelation,
)


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


def _make_info(num_episodes: int = 50) -> DatasetInfo:
    return DatasetInfo(
        repo_id="test/dataset",
        num_episodes=num_episodes,
        num_frames=num_episodes * 100,
        fps=30.0,
        robot_type="test",
        camera_names=[],
        video_keys=[],
    )


def _make_coverage(score: float = 0.8) -> FullCoverageReport:
    return FullCoverageReport(
        position_diversity=score,
        action_diversity=score,
        episode_consistency=score,
        temporal_coverage=score,
    )


# ===========================================================================
# Action Smoothness
# ===========================================================================


class TestActionSmoothness:
    def test_smooth_sine_wave(self):
        """Sine wave → high smoothness."""
        t = np.linspace(0, 2 * np.pi, 200)
        actions = np.column_stack([np.sin(t), np.cos(t)])
        score = compute_action_smoothness(actions)
        assert score > 0.7

    def test_random_noise(self):
        """Random noise → low smoothness."""
        rng = np.random.RandomState(42)
        actions = rng.randn(200, 4)
        score = compute_action_smoothness(actions)
        assert score < 0.6

    def test_smooth_beats_noisy(self):
        """Sine wave should be smoother than random noise."""
        t = np.linspace(0, 2 * np.pi, 200)
        smooth = np.column_stack([np.sin(t + i) for i in range(4)])
        noisy = np.random.RandomState(42).randn(200, 4)

        assert compute_action_smoothness(smooth) > compute_action_smoothness(noisy)

    def test_short_episode(self):
        """< 4 timesteps → default 0.5."""
        actions = np.random.randn(3, 4)
        assert compute_action_smoothness(actions) == 0.5

    def test_constant_actions(self):
        """Constant actions → high smoothness (no jerk)."""
        actions = np.ones((100, 4)) * 0.5
        score = compute_action_smoothness(actions)
        assert score >= 0.5


# ===========================================================================
# Spectral Smoothness
# ===========================================================================


class TestSpectralSmoothness:
    def test_smooth_sine_wave(self):
        """Sine wave concentrates power in low frequencies → high score."""
        t = np.linspace(0, 2 * np.pi, 200)
        actions = np.column_stack([np.sin(t), np.cos(t)])
        score = compute_spectral_smoothness(actions)
        assert score > 0.7

    def test_random_noise(self):
        """Random noise spreads power across frequencies → low score."""
        rng = np.random.RandomState(42)
        actions = rng.randn(200, 4)
        score = compute_spectral_smoothness(actions)
        assert score < 0.5

    def test_constant_actions(self):
        """Constant signal = all power in DC → high score."""
        actions = np.ones((100, 4)) * 0.5
        score = compute_spectral_smoothness(actions)
        assert score >= 0.9

    def test_short_episode(self):
        """< 16 timesteps → default 0.5."""
        actions = np.random.randn(10, 4)
        assert compute_spectral_smoothness(actions) == 0.5


# ===========================================================================
# Episode Completion
# ===========================================================================


class TestEpisodeCompletion:
    def test_converging_episode(self):
        """State variance decreases → high completion."""
        # Start with spread, end at a point
        rng = np.random.RandomState(42)
        states = np.zeros((100, 4))
        for t in range(100):
            scale = 1.0 - t / 100.0  # Decreasing variance
            states[t] = rng.randn(4) * scale
        score = compute_episode_completion(states)
        assert score > 0.5

    def test_diverging_episode(self):
        """State variance increases → low completion."""
        rng = np.random.RandomState(42)
        states = np.zeros((100, 4))
        for t in range(100):
            scale = t / 100.0  # Increasing variance
            states[t] = rng.randn(4) * scale
        score = compute_episode_completion(states)
        assert score < 0.5

    def test_short_episode(self):
        """< 10 timesteps → default 0.5."""
        states = np.random.randn(5, 4)
        assert compute_episode_completion(states) == 0.5


# ===========================================================================
# Observation Consistency
# ===========================================================================


class TestObservationConsistency:
    def test_same_initial_states(self):
        """All episodes start at same state → high consistency."""
        episodes = [np.vstack([np.zeros(4), np.random.randn(50, 4)]) for _ in range(10)]
        score = compute_observation_consistency(episodes)
        assert score > 0.8

    def test_random_initial_states(self):
        """Random initial states → lower consistency."""
        rng = np.random.RandomState(42)
        episodes = [np.vstack([rng.randn(1, 4) * 5, rng.randn(50, 4)]) for _ in range(10)]
        score = compute_observation_consistency(episodes)
        assert score < 0.5

    def test_consistent_beats_random(self):
        """Consistent initial states > random initial states."""
        consistent = [np.vstack([np.zeros(4), np.random.randn(50, 4)]) for _ in range(10)]
        rng = np.random.RandomState(42)
        random_init = [np.vstack([rng.randn(1, 4) * 5, rng.randn(50, 4)]) for _ in range(10)]

        assert compute_observation_consistency(consistent) > compute_observation_consistency(random_init)

    def test_single_episode(self):
        """1 episode → default 0.5."""
        assert compute_observation_consistency([np.random.randn(50, 4)]) == 0.5


# ===========================================================================
# Temporal Autocorrelation
# ===========================================================================


class TestTemporalAutocorrelation:
    def test_correlated_actions(self):
        """Smooth trajectory → high autocorrelation."""
        t = np.linspace(0, 4 * np.pi, 200)
        actions = np.column_stack([np.sin(t), np.cos(t)])
        score = compute_temporal_autocorrelation(actions)
        assert score > 0.7

    def test_random_actions(self):
        """Random noise → autocorrelation near 0.5."""
        rng = np.random.RandomState(42)
        actions = rng.randn(200, 4)
        score = compute_temporal_autocorrelation(actions)
        assert 0.3 <= score <= 0.7

    def test_short_episode(self):
        """< 3 timesteps → default 0.5."""
        assert compute_temporal_autocorrelation(np.random.randn(2, 4)) == 0.5


# ===========================================================================
# Mutual Information (correlation proxy)
# ===========================================================================


class TestActionStateMI:
    def test_correlated_pairs(self):
        """Actions = f(states) → higher MI than baseline."""
        rng = np.random.RandomState(42)
        states = rng.randn(500, 4)
        # Actions are linear functions of states
        actions = states @ rng.randn(4, 3) + rng.randn(3) * 0.1
        score = compute_action_state_mi(actions, states)
        assert score > 0.1

    def test_independent_pairs(self):
        """Independent actions and states → low MI."""
        rng = np.random.RandomState(42)
        actions = rng.randn(500, 4)
        states = rng.randn(500, 6)
        score = compute_action_state_mi(actions, states)
        assert score < 0.1

    def test_correlated_beats_independent(self):
        """Correlated > independent."""
        rng = np.random.RandomState(42)
        states = rng.randn(500, 4)
        correlated_actions = states @ rng.randn(4, 3)
        independent_actions = rng.randn(500, 3)

        corr_score = compute_action_state_mi(correlated_actions, states)
        indep_score = compute_action_state_mi(independent_actions, states)
        assert corr_score > indep_score

    def test_short_data(self):
        """< 5 samples → default 0.5."""
        assert compute_action_state_mi(np.random.randn(3, 4), np.random.randn(3, 4)) == 0.5


# ===========================================================================
# Enhanced Quality Scoring
# ===========================================================================


class TestEnhancedQualityScore:
    def test_enhanced_with_action_data(self):
        """When action_episodes is provided, enhanced scoring is used."""
        info = _make_info(50)
        coverage = _make_coverage(0.8)

        rng = np.random.RandomState(42)
        action_eps = [rng.randn(100, 6) for _ in range(10)]
        state_eps = [rng.randn(100, 14) for _ in range(10)]

        report = compute_quality_score(
            info, coverage,
            action_episodes=action_eps,
            state_episodes=state_eps,
        )

        assert "action_quality" in report.component_scores
        assert "episode_quality" in report.component_scores
        assert "signal_summary" in report.component_scores
        assert 0.0 <= report.overall_score <= 1.0

    def test_backward_compatible_without_action_data(self):
        """Without action_episodes, standard scoring is used."""
        info = _make_info(50)
        coverage = _make_coverage(0.8)

        report = compute_quality_score(info, coverage)

        assert "action_quality" not in report.component_scores
        assert "episode_quality" not in report.component_scores
        assert 0.0 <= report.overall_score <= 1.0

    def test_signal_diagnostics_summary_used(self):
        """signal_diagnostics_summary is included when provided."""
        info = _make_info(50)
        coverage = _make_coverage(0.8)
        action_eps = [np.random.randn(100, 6) for _ in range(10)]

        report = compute_quality_score(
            info, coverage,
            action_episodes=action_eps,
            signal_diagnostics_summary=0.9,
        )

        assert report.component_scores["signal_summary"] == 0.9
