"""Tests for distribution-free outlier detection using Modified Z-Score."""

import numpy as np
import pytest

from orbit.analyzer.outlier_detection import (
    DatasetOutlierReport,
    EpisodeOutlierResult,
    MZS_THRESHOLD,
    _compute_jerk,
    _compute_path_length,
    _compute_smoothness,
    _modified_z_scores,
    detect_outliers,
)


def _make_smooth_episode(length=100, dim=6, seed=0):
    """Generate a smooth sinusoidal trajectory."""
    rng = np.random.RandomState(seed)
    t = np.linspace(0, 2 * np.pi, length)
    actions = np.stack([np.sin(t + rng.uniform(0, np.pi)) for _ in range(dim)], axis=-1)
    return actions.astype(np.float32)


def _make_noisy_episode(length=100, dim=6, seed=99):
    """Generate a noisy random-walk trajectory (outlier-like)."""
    rng = np.random.RandomState(seed)
    actions = rng.randn(length, dim).astype(np.float32) * 5.0
    return actions


class TestModifiedZScore:
    def test_uniform_values_no_outliers(self):
        values = np.array([10.0, 10.1, 9.9, 10.0, 10.05])
        mzs = _modified_z_scores(values)
        assert all(abs(z) < MZS_THRESHOLD for z in mzs)

    def test_single_outlier_detected(self):
        values = np.array([10.0, 10.1, 9.9, 10.0, 50.0])
        mzs = _modified_z_scores(values)
        # The last value should have a high MZS
        assert abs(mzs[-1]) > MZS_THRESHOLD

    def test_all_identical_no_outliers(self):
        values = np.array([5.0, 5.0, 5.0, 5.0, 5.0])
        mzs = _modified_z_scores(values)
        assert all(z == 0 for z in mzs)


class TestMetricComputation:
    def test_jerk_smooth_trajectory(self):
        actions = _make_smooth_episode(length=100)
        jerk = _compute_jerk(actions)
        assert jerk >= 0

    def test_jerk_noisy_higher_than_smooth(self):
        smooth = _make_smooth_episode(length=100)
        noisy = _make_noisy_episode(length=100)
        assert _compute_jerk(noisy) > _compute_jerk(smooth)

    def test_path_length_positive(self):
        actions = _make_smooth_episode()
        pl = _compute_path_length(actions)
        assert pl > 0

    def test_path_length_zero_for_single_step(self):
        actions = np.array([[1.0, 2.0, 3.0]])
        assert _compute_path_length(actions) == 0.0

    def test_smoothness_bounded(self):
        actions = _make_smooth_episode()
        s = _compute_smoothness(actions)
        assert 0.0 <= s <= 1.0


class TestDetectOutliers:
    def test_empty_dataset(self):
        report = detect_outliers({})
        assert report.total_episodes == 0
        assert report.outlier_count == 0
        assert report.dataset_health == "healthy"

    def test_too_few_episodes_no_detection(self):
        """With < 5 episodes, no outlier detection is performed."""
        data = {
            i: {"actions": _make_smooth_episode(seed=i), "states": None, "length": 100}
            for i in range(3)
        }
        report = detect_outliers(data)
        assert report.total_episodes == 3
        assert report.outlier_count == 0
        assert all(not r.is_outlier for r in report.episode_results)

    def test_normal_episodes_not_problematic(self):
        """8 similar smooth episodes → not problematic."""
        data = {
            i: {"actions": _make_smooth_episode(seed=i), "states": None, "length": 100}
            for i in range(8)
        }
        report = detect_outliers(data)
        assert report.total_episodes == 8
        assert report.dataset_health in ("healthy", "noisy")
        assert report.outlier_fraction <= 0.25

    def test_one_outlier_detected(self):
        """5 normal + 1 very noisy episode → outlier detected."""
        data = {}
        for i in range(5):
            data[i] = {"actions": _make_smooth_episode(seed=i), "states": None, "length": 100}
        # Add one very different episode
        data[5] = {"actions": _make_noisy_episode(length=100), "states": None, "length": 100}

        report = detect_outliers(data)
        assert report.total_episodes == 6
        assert report.outlier_count >= 1
        # The noisy episode should be flagged
        outlier_indices = [r.episode_index for r in report.episode_results if r.is_outlier]
        assert 5 in outlier_indices

    def test_outlier_has_reasons(self):
        """Outlier episodes should have specific reasons."""
        data = {}
        for i in range(5):
            data[i] = {"actions": _make_smooth_episode(seed=i), "states": None, "length": 100}
        data[5] = {"actions": _make_noisy_episode(length=100), "states": None, "length": 100}

        report = detect_outliers(data)
        outlier_results = [r for r in report.episode_results if r.is_outlier]
        assert len(outlier_results) >= 1
        for r in outlier_results:
            assert len(r.outlier_reasons) > 0
            assert any("MZS=" in reason for reason in r.outlier_reasons)

    def test_very_short_outlier_detected(self):
        """An episode with very different length should be detected."""
        data = {}
        for i in range(5):
            data[i] = {"actions": _make_smooth_episode(length=100, seed=i), "states": None, "length": 100}
        # Very short episode
        data[5] = {"actions": _make_smooth_episode(length=5, seed=99), "states": None, "length": 5}

        report = detect_outliers(data)
        outlier_indices = [r.episode_index for r in report.episode_results if r.is_outlier]
        assert 5 in outlier_indices

    def test_metric_summaries_populated(self):
        data = {
            i: {"actions": _make_smooth_episode(seed=i), "states": None, "length": 100}
            for i in range(6)
        }
        report = detect_outliers(data)
        assert "smoothness" in report.metric_summaries
        assert "path_length" in report.metric_summaries
        for name, summary in report.metric_summaries.items():
            assert "median" in summary
            assert "mad" in summary

    def test_health_classification(self):
        """Many outliers → 'problematic' health."""
        data = {}
        # 3 normal + 3 very noisy
        for i in range(3):
            data[i] = {"actions": _make_smooth_episode(seed=i), "states": None, "length": 100}
        for i in range(3, 6):
            data[i] = {"actions": _make_noisy_episode(length=100, seed=i), "states": None, "length": 100}

        report = detect_outliers(data)
        # With 50% outliers, should be problematic
        assert report.dataset_health in ("noisy", "problematic")

    def test_quality_vector_populated(self):
        data = {
            i: {"actions": _make_smooth_episode(seed=i), "states": None, "length": 100}
            for i in range(6)
        }
        report = detect_outliers(data)
        for r in report.episode_results:
            assert "smoothness" in r.quality_vector
            assert "path_length" in r.quality_vector
            assert "duration" in r.quality_vector
