"""Exhaustive tests for orbit.analyzer.signal_diagnostics."""

from __future__ import annotations

import numpy as np
import pytest

from orbit.analyzer.signal_diagnostics import (
    DEFAULT_DEAD_JOINT_RATIO,
    DEFAULT_MIN_TRANSITIONS_PER_EPISODE,
    DEFAULT_ZERO_VELOCITY_THRESHOLD,
    DirectionalBalance,
    EpisodeDiagnostic,
    GripperAnalysis,
    JointDiagnostic,
    SignalDiagnosticsReport,
    compute_signal_diagnostics,
)


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


def _make_episodes(
    num_episodes: int = 10,
    timesteps: int = 100,
    action_dim: int = 6,
    seed: int = 42,
) -> list[np.ndarray]:
    """Create random episodes with uniform actions."""
    rng = np.random.RandomState(seed)
    return [rng.randn(timesteps, action_dim) for _ in range(num_episodes)]


def _sine_episodes(
    num_episodes: int = 10,
    timesteps: int = 100,
    action_dim: int = 6,
    freq: float = 0.1,
) -> list[np.ndarray]:
    """Create smooth sine-wave episodes."""
    t = np.linspace(0, 2 * np.pi * freq * timesteps, timesteps)
    episodes = []
    for i in range(num_episodes):
        ep = np.column_stack([
            np.sin(t + i * 0.5 + d * 0.3) for d in range(action_dim)
        ])
        episodes.append(ep)
    return episodes


# ===========================================================================
# JOINT DIAGNOSTICS
# ===========================================================================


class TestJointDiagnostics:
    """Tests for per-joint health diagnostics."""

    def test_healthy_joint(self):
        """Uniform random actions → high range utilization, not dead, not clipping."""
        episodes = _make_episodes(num_episodes=20, timesteps=200, action_dim=4)
        report = compute_signal_diagnostics(episodes)

        for jd in report.joint_diagnostics:
            assert jd.range_utilization > 0.5, f"Joint {jd.joint_index} low range util"
            assert jd.is_dead is False
            assert jd.is_clipping is False

    def test_dead_joint(self):
        """One dimension all zeros → is_dead=True."""
        episodes = _make_episodes(num_episodes=10, timesteps=100, action_dim=4)
        for ep in episodes:
            ep[:, 2] = 0.0  # Kill joint 2

        report = compute_signal_diagnostics(episodes)
        jd2 = report.joint_diagnostics[2]

        assert jd2.is_dead is True
        assert jd2.zero_velocity_ratio > 0.95
        assert 2 in report.dead_dimensions

    def test_partially_dead_joint(self):
        """85% zeros but nonzero range → is_dead=False (not constant).

        A joint with high zero-velocity ratio but a measurable range is
        "mostly stationary" — common in scripted demos and bimanual setups
        where a base joint is held.  Only truly constant joints (range ≈ 0)
        are flagged as dead hardware.
        """
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(10):
            ep = rng.randn(100, 4)
            # Make 85% of joint 1 zero — but 15% still have values
            mask = rng.random(100) < 0.85
            ep[mask, 1] = 0.0
            episodes.append(ep)

        report = compute_signal_diagnostics(episodes)
        jd1 = report.joint_diagnostics[1]

        # High ZVR but nonzero range → NOT dead (just mostly stationary)
        assert jd1.is_dead is False
        assert jd1.zero_velocity_ratio > DEFAULT_DEAD_JOINT_RATIO

    def test_not_quite_dead_joint(self):
        """70% zeros → is_dead=False."""
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(10):
            ep = rng.randn(100, 4)
            mask = rng.random(100) < 0.70
            ep[mask, 1] = 0.0
            episodes.append(ep)

        report = compute_signal_diagnostics(episodes)
        jd1 = report.joint_diagnostics[1]

        assert jd1.is_dead is False
        assert jd1.zero_velocity_ratio < DEFAULT_DEAD_JOINT_RATIO

    def test_clipping_joint(self):
        """Majority of values at extremes → is_clipping=True."""
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(10):
            ep = rng.uniform(-1.0, 1.0, size=(200, 4))
            # Force 50% of joint 0 to the max value
            ep[:100, 0] = 1.0
            episodes.append(ep)

        report = compute_signal_diagnostics(episodes)
        jd0 = report.joint_diagnostics[0]

        assert jd0.is_clipping is True

    def test_directional_bias(self):
        """All positive actions on one joint → high directional bias."""
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(10):
            ep = rng.randn(100, 4)
            ep[:, 0] = np.abs(ep[:, 0]) + 0.1  # All positive
            episodes.append(ep)

        report = compute_signal_diagnostics(episodes)
        jd0 = report.joint_diagnostics[0]

        assert jd0.directional_bias > 0.9

    def test_balanced_direction(self):
        """Equal positive/negative → directional_bias ≈ 0.5."""
        episodes = _make_episodes(num_episodes=20, timesteps=200, action_dim=4)
        report = compute_signal_diagnostics(episodes)

        for jd in report.joint_diagnostics:
            assert 0.4 <= jd.directional_bias <= 0.65, (
                f"Joint {jd.joint_index} bias {jd.directional_bias} not balanced"
            )

    def test_smoothness_smooth(self):
        """Sine wave actions → high smoothness."""
        episodes = _sine_episodes(num_episodes=10, timesteps=200, action_dim=4)
        report = compute_signal_diagnostics(episodes)

        for jd in report.joint_diagnostics:
            assert jd.smoothness > 0.5, f"Joint {jd.joint_index} smoothness too low"

    def test_smoothness_jerky(self):
        """Random noise → lower smoothness than sine waves."""
        smooth_eps = _sine_episodes(num_episodes=10, timesteps=200, action_dim=4)
        noisy_eps = _make_episodes(num_episodes=10, timesteps=200, action_dim=4)

        smooth_report = compute_signal_diagnostics(smooth_eps)
        noisy_report = compute_signal_diagnostics(noisy_eps)

        smooth_avg = np.mean([jd.smoothness for jd in smooth_report.joint_diagnostics])
        noisy_avg = np.mean([jd.smoothness for jd in noisy_report.joint_diagnostics])

        assert smooth_avg > noisy_avg, "Sine wave should be smoother than random noise"

    def test_range_utilization(self):
        """Actions only using 20% of known joint range → low range_utilization."""
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(10):
            # Actions between 0 and 0.2, but joint range is 0 to 1.0
            ep = rng.uniform(0.0, 0.2, size=(100, 4))
            episodes.append(ep)

        joint_limits = [(0.0, 1.0)] * 4
        report = compute_signal_diagnostics(episodes, joint_limits=joint_limits)

        for jd in report.joint_diagnostics:
            assert jd.range_utilization < 0.35, (
                f"Joint {jd.joint_index} range util {jd.range_utilization} too high"
            )


# ===========================================================================
# GRIPPER ANALYSIS
# ===========================================================================


class TestGripperAnalysis:
    """Tests for gripper detection and analysis."""

    def test_gripper_autodetect(self):
        """Last dim is binary (0/1) → correctly identified as gripper."""
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(10):
            ep = rng.randn(100, 6)
            ep[:, 5] = rng.choice([0.0, 1.0], size=100)
            episodes.append(ep)

        report = compute_signal_diagnostics(episodes)
        assert report.gripper_analysis.gripper_dim == 5

    def test_gripper_autodetect_non_binary(self):
        """No binary dim → gripper_dim may be None."""
        episodes = _make_episodes(num_episodes=10, timesteps=100, action_dim=6)
        report = compute_signal_diagnostics(episodes)

        # With random continuous data, gripper should not be detected
        # (the auto-detection looks for binary-like dims)
        ga = report.gripper_analysis
        # Either None or the detection admits it's not confident
        # The key is it shouldn't crash
        assert isinstance(ga, GripperAnalysis)

    def test_gripper_sufficient_grasps(self):
        """2+ transitions per episode → has_sufficient_grasps=True."""
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(10):
            ep = rng.randn(100, 6)
            # Create binary gripper with multiple transitions
            gripper = np.zeros(100)
            gripper[20:40] = 1.0
            gripper[60:80] = 1.0
            ep[:, 5] = gripper
            episodes.append(ep)

        report = compute_signal_diagnostics(episodes)
        ga = report.gripper_analysis

        assert ga.gripper_dim == 5
        assert ga.transitions_per_episode >= 2.0
        assert ga.has_sufficient_grasps is True

    def test_gripper_insufficient_grasps(self):
        """< 0.5 transitions per episode → has_sufficient_grasps=False."""
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(10):
            ep = rng.randn(100, 6)
            # Gripper always open (no transitions)
            ep[:, 5] = np.ones(100)
            episodes.append(ep)

        report = compute_signal_diagnostics(episodes)
        ga = report.gripper_analysis

        if ga.gripper_dim is not None:
            assert ga.transitions_per_episode < DEFAULT_MIN_TRANSITIONS_PER_EPISODE
            assert ga.has_sufficient_grasps is False

    def test_gripper_always_open(self):
        """95% open → issues mention this."""
        rng = np.random.RandomState(42)
        episodes = []
        for _ in range(10):
            ep = rng.randn(100, 6)
            # 95% ones (open), 5% zeros
            gripper = np.ones(100)
            gripper[:5] = 0.0
            ep[:, 5] = gripper
            episodes.append(ep)

        report = compute_signal_diagnostics(episodes)
        ga = report.gripper_analysis

        if ga.gripper_dim == 5:
            assert ga.open_ratio > 0.8 or ga.close_ratio > 0.8


# ===========================================================================
# EPISODE DIAGNOSTICS
# ===========================================================================


class TestEpisodeDiagnostics:
    """Tests for per-episode quality diagnostics."""

    def test_good_episode(self):
        """Smooth, full length, active joints → quality_score > 0.5, recommendation=keep."""
        episodes = _sine_episodes(num_episodes=10, timesteps=100, action_dim=6)
        report = compute_signal_diagnostics(episodes)

        for ed in report.episode_diagnostics:
            assert ed.quality_score > 0.5
            assert ed.recommendation == "keep"

    def test_truncated_episode(self):
        """Episode much shorter than median → is_truncated=True."""
        episodes = _make_episodes(num_episodes=9, timesteps=100, action_dim=4)
        # Add one very short episode
        episodes.append(np.random.randn(10, 4))  # 10 frames vs median ~100

        report = compute_signal_diagnostics(episodes)
        # The short episode is at index 9
        ed_short = report.episode_diagnostics[9]

        assert ed_short.is_truncated is True
        assert ed_short.recommendation in ("remove", "review")

    def test_jerky_episode(self):
        """Episode with sudden jerk spikes → detected."""
        rng = np.random.RandomState(42)
        episodes = _sine_episodes(num_episodes=9, timesteps=100, action_dim=4)

        # Create a jerky episode
        jerky = np.zeros((100, 4))
        for t in range(100):
            jerky[t] = np.sin(t * 0.1) * np.ones(4)
        # Insert sudden jumps
        for t in [20, 40, 60, 70, 85]:
            jerky[t] = jerky[t - 1] + rng.randn(4) * 50
        episodes.append(jerky)

        report = compute_signal_diagnostics(episodes)
        ed_jerky = report.episode_diagnostics[9]

        assert ed_jerky.jerk_spike_count > 0

    def test_episode_with_dead_joints(self):
        """2 joints don't move in this episode → dead_joints populated."""
        rng = np.random.RandomState(42)
        episodes = _make_episodes(num_episodes=9, timesteps=100, action_dim=6)

        # Episode with 2 dead joints
        dead_ep = rng.randn(100, 6)
        dead_ep[:, 2] = 0.0
        dead_ep[:, 4] = 0.0
        episodes.append(dead_ep)

        report = compute_signal_diagnostics(episodes)
        ed_dead = report.episode_diagnostics[9]

        assert 2 in ed_dead.dead_joints_in_episode
        assert 4 in ed_dead.dead_joints_in_episode


# ===========================================================================
# FULL REPORT
# ===========================================================================


class TestFullReport:
    """Tests for the full SignalDiagnosticsReport."""

    def test_healthy_dataset(self):
        """50 episodes, diverse actions → no blockers."""
        episodes = _make_episodes(num_episodes=50, timesteps=100, action_dim=6)
        report = compute_signal_diagnostics(episodes)

        assert isinstance(report, SignalDiagnosticsReport)
        assert len(report.blockers) == 0
        assert len(report.passing) > 0

    def test_dataset_with_dead_dim(self):
        """One dead dimension → blocker list includes it."""
        episodes = _make_episodes(num_episodes=20, timesteps=100, action_dim=6)
        for ep in episodes:
            ep[:, 3] = 0.0  # Kill joint 3

        report = compute_signal_diagnostics(episodes)

        assert 3 in report.dead_dimensions
        assert any("Dead" in b or "dead" in b.lower() for b in report.blockers)

    def test_dataset_needing_removal(self):
        """Many bad episodes → blocker about removing episodes."""
        rng = np.random.RandomState(42)
        episodes = []
        # 7 good episodes
        for _ in range(7):
            episodes.append(rng.randn(100, 4))
        # 3 truncated episodes (30%)
        for _ in range(3):
            episodes.append(rng.randn(5, 4))

        report = compute_signal_diagnostics(episodes)

        assert len(report.episodes_to_remove) > 0
        # 30% flagged should trigger a blocker
        assert any("removal" in b.lower() or "remove" in b.lower() for b in report.blockers)

    def test_prescription_ordering(self):
        """Dead joint + directional bias + bad episodes → prescription lists dead joint first."""
        rng = np.random.RandomState(42)
        episodes = []
        # Create dataset with multiple issues
        for _ in range(7):
            ep = rng.randn(100, 6)
            ep[:, 3] = 0.0  # Dead joint
            ep[:, 0] = np.abs(ep[:, 0]) + 0.5  # Directional bias
            episodes.append(ep)
        # Add truncated episodes
        for _ in range(3):
            ep = rng.randn(5, 6)
            ep[:, 3] = 0.0
            episodes.append(ep)

        report = compute_signal_diagnostics(episodes)

        assert len(report.prescription) > 0
        # Dead joint should be mentioned early in prescription
        assert any("joint_3" in p.lower() or "joint 3" in p.lower()
                    for p in report.prescription[:2])

    def test_empty_dataset(self):
        """0 episodes → handles gracefully."""
        report = compute_signal_diagnostics([])

        assert isinstance(report, SignalDiagnosticsReport)
        assert len(report.blockers) > 0
        assert len(report.joint_diagnostics) == 0

    def test_single_episode(self):
        """1 episode → handles gracefully."""
        ep = np.random.randn(50, 4)
        report = compute_signal_diagnostics([ep])

        assert isinstance(report, SignalDiagnosticsReport)
        assert len(report.episode_diagnostics) == 1

    def test_variable_length_episodes(self):
        """Episodes of different lengths → handles correctly."""
        rng = np.random.RandomState(42)
        episodes = [
            rng.randn(50, 4),
            rng.randn(100, 4),
            rng.randn(75, 4),
            rng.randn(120, 4),
            rng.randn(80, 4),
        ]
        report = compute_signal_diagnostics(episodes)

        assert len(report.episode_diagnostics) == 5
        for i, ed in enumerate(report.episode_diagnostics):
            assert ed.episode_index == i
            assert ed.length == len(episodes[i])

    def test_1d_actions(self):
        """Single action dimension → no crash."""
        episodes = [np.random.randn(50, 1) for _ in range(5)]
        report = compute_signal_diagnostics(episodes)

        assert len(report.joint_diagnostics) == 1

    def test_high_dim_actions(self):
        """20 action dimensions → handles correctly."""
        episodes = [np.random.randn(50, 20) for _ in range(5)]
        report = compute_signal_diagnostics(episodes)

        assert len(report.joint_diagnostics) == 20

    def test_3d_array_input(self):
        """3D numpy array input with episode_lengths → works correctly."""
        rng = np.random.RandomState(42)
        actions = rng.randn(5, 100, 4)
        lengths = [80, 100, 60, 90, 70]

        report = compute_signal_diagnostics(actions, episode_lengths=lengths)

        assert len(report.episode_diagnostics) == 5
        assert report.episode_diagnostics[0].length == 80
        assert report.episode_diagnostics[2].length == 60


# ===========================================================================
# EDGE CASES
# ===========================================================================


class TestEdgeCases:
    """Edge case tests for signal diagnostics."""

    def test_all_zeros(self):
        """Entire dataset is zeros → every joint dead."""
        episodes = [np.zeros((100, 6)) for _ in range(10)]
        report = compute_signal_diagnostics(episodes)

        assert len(report.dead_dimensions) == 6
        for jd in report.joint_diagnostics:
            assert jd.is_dead is True
        assert len(report.blockers) > 0

    def test_constant_actions(self):
        """All actions same nonzero value → detected as problematic."""
        episodes = [np.full((100, 4), 0.5) for _ in range(10)]
        report = compute_signal_diagnostics(episodes)

        # Constant value should show low diversity / zero velocity depending on threshold
        # Range utilization should be 0 (no range used)
        for jd in report.joint_diagnostics:
            assert jd.range_utilization == 0.0 or jd.is_dead is False

    def test_nan_handling(self):
        """Some NaN values in actions → doesn't crash."""
        rng = np.random.RandomState(42)
        episodes = [rng.randn(100, 4) for _ in range(5)]
        # Inject NaN
        episodes[0][10:15, :] = np.nan
        episodes[2][50, 2] = np.nan

        report = compute_signal_diagnostics(episodes)
        assert isinstance(report, SignalDiagnosticsReport)
        assert len(report.episode_diagnostics) == 5

    def test_inf_handling(self):
        """inf values → doesn't crash."""
        rng = np.random.RandomState(42)
        episodes = [rng.randn(100, 4) for _ in range(5)]
        episodes[1][30, 0] = np.inf
        episodes[3][60, 2] = -np.inf

        report = compute_signal_diagnostics(episodes)
        assert isinstance(report, SignalDiagnosticsReport)

    def test_very_short_episodes(self):
        """Episodes with only 1-3 timesteps → doesn't crash."""
        episodes = [
            np.random.randn(1, 4),
            np.random.randn(2, 4),
            np.random.randn(3, 4),
        ]
        report = compute_signal_diagnostics(episodes)
        assert len(report.episode_diagnostics) == 3

    def test_empty_episode_in_list(self):
        """Empty array in episode list → doesn't crash."""
        episodes = [
            np.random.randn(50, 4),
            np.zeros((0, 4)),
            np.random.randn(50, 4),
        ]
        report = compute_signal_diagnostics(episodes)
        assert isinstance(report, SignalDiagnosticsReport)

    def test_joint_names_and_limits(self):
        """Custom joint names and limits are reflected in output."""
        episodes = _make_episodes(num_episodes=5, timesteps=50, action_dim=3)
        names = ["shoulder", "elbow", "wrist"]
        limits = [(-3.14, 3.14), (-2.0, 2.0), (-1.5, 1.5)]

        report = compute_signal_diagnostics(
            episodes, joint_names=names, joint_limits=limits
        )

        assert report.joint_diagnostics[0].joint_name == "shoulder"
        assert report.joint_diagnostics[1].joint_name == "elbow"
        assert report.joint_diagnostics[2].joint_name == "wrist"
        assert report.joint_diagnostics[0].range_limit == (-3.14, 3.14)
