"""Tests for the DatasetReadinessGrader."""

import pytest

from orbit.analyzer.success_predictor import DatasetReadinessGrader, ReadinessResult


def _default_signal_health(**overrides):
    base = {
        "dead_dimensions": [],
        "blockers": [],
        "warnings": [],
        "total_joints": 6,
        "clipping_joints": [],
    }
    base.update(overrides)
    return base


def _default_episode_ranking(**overrides):
    base = {"total": 50, "remove_count": 0, "review_count": 5}
    base.update(overrides)
    return base


class TestDatasetReadinessGrader:
    """Test the readiness grading logic."""

    def setup_method(self):
        self.grader = DatasetReadinessGrader()

    def test_perfect_metrics_produce_grade_a(self):
        result = self.grader.grade(
            quality_score=85,
            episode_consistency=0.92,
            signal_health=_default_signal_health(),
            policy_fit_score=0.85,
            episode_ranking=_default_episode_ranking(),
            action_divergence=0.02,
            episode_count=100,
            policy_type="diffusion_policy",
        )
        assert isinstance(result, ReadinessResult)
        assert result.grade == "A"
        assert result.score >= 90

    def test_critical_dead_joints_produce_grade_f(self):
        result = self.grader.grade(
            quality_score=80,
            episode_consistency=0.90,
            signal_health=_default_signal_health(
                dead_dimensions=[0, 1, 2, 3, 4],
                total_joints=7,
            ),
            policy_fit_score=0.80,
            episode_ranking=_default_episode_ranking(),
            action_divergence=0.03,
            episode_count=100,
            policy_type="act",
        )
        # >60% dead joints = critical failure (raised from 30% because
        # many real datasets have padded/unused action dimensions)
        assert result.grade == "F"
        assert result.score == 0

    def test_too_few_episodes_produce_grade_f(self):
        result = self.grader.grade(
            quality_score=90,
            episode_consistency=0.95,
            signal_health=_default_signal_health(),
            policy_fit_score=0.90,
            episode_ranking=_default_episode_ranking(total=5, remove_count=0),
            action_divergence=0.01,
            episode_count=5,
            policy_type="act",
        )
        assert result.grade == "F"
        assert "Not enough" in result.problems[0]

    def test_majority_unusable_episodes_produce_grade_f(self):
        result = self.grader.grade(
            quality_score=60,
            episode_consistency=0.50,
            signal_health=_default_signal_health(),
            policy_fit_score=0.60,
            episode_ranking=_default_episode_ranking(total=50, remove_count=30),
            action_divergence=0.10,
            episode_count=50,
            policy_type="diffusion_policy",
        )
        assert result.grade == "F"
        assert result.score == 0

    def test_severe_temporal_misalignment_produces_grade_f(self):
        result = self.grader.grade(
            quality_score=80,
            episode_consistency=0.85,
            signal_health=_default_signal_health(),
            policy_fit_score=0.80,
            episode_ranking=_default_episode_ranking(),
            action_divergence=0.03,
            episode_count=100,
            policy_type="act",
            temporal_alignment={"misaligned_pct": 0.60},
        )
        assert result.grade == "F"

    def test_grade_is_deterministic(self):
        """Same inputs must always produce the same grade."""
        kwargs = dict(
            quality_score=72,
            episode_consistency=0.75,
            signal_health=_default_signal_health(dead_dimensions=[3]),
            policy_fit_score=0.65,
            episode_ranking=_default_episode_ranking(remove_count=3),
            action_divergence=0.08,
            episode_count=60,
            policy_type="act",
        )
        r1 = self.grader.grade(**kwargs)
        r2 = self.grader.grade(**kwargs)
        assert r1.grade == r2.grade
        assert r1.score == r2.score
        assert r1.strengths == r2.strengths
        assert r1.problems == r2.problems

    def test_zero_episodes(self):
        result = self.grader.grade(
            quality_score=0,
            episode_consistency=0,
            signal_health=_default_signal_health(),
            policy_fit_score=0,
            episode_ranking=_default_episode_ranking(total=0),
            action_divergence=0,
            episode_count=0,
            policy_type="act",
        )
        assert result.grade == "F"

    def test_one_episode_below_minimum(self):
        result = self.grader.grade(
            quality_score=80,
            episode_consistency=0.90,
            signal_health=_default_signal_health(),
            policy_fit_score=0.80,
            episode_ranking=_default_episode_ranking(total=1),
            action_divergence=0.0,
            episode_count=1,
            policy_type="act",
        )
        assert result.grade == "F"  # ACT needs >= 15

    def test_exactly_at_minimum_episodes(self):
        """At exactly the minimum, should not be a critical failure."""
        result = self.grader.grade(
            quality_score=80,
            episode_consistency=0.85,
            signal_health=_default_signal_health(),
            policy_fit_score=0.80,
            episode_ranking=_default_episode_ranking(total=15),
            action_divergence=0.03,
            episode_count=15,  # ACT minimum
            policy_type="act",
        )
        assert result.grade != "F"  # should pass critical check
        assert result.score > 0

    def test_moderate_issues_produce_below_a(self):
        """Multiple moderate issues should produce a grade below A."""
        sig = _default_signal_health(dead_dimensions=[3])
        sig["total_joints"] = 14  # 1/14 = 7%, well below critical 30%
        sig["clipping_joints"] = [1, 2, 5]
        result = self.grader.grade(
            quality_score=70,
            episode_consistency=0.70,
            signal_health=sig,
            policy_fit_score=0.55,
            episode_ranking=_default_episode_ranking(remove_count=8, review_count=10),
            action_divergence=0.15,
            episode_count=60,
            policy_type="act",
        )
        assert result.grade in ("B", "C", "D")
        assert result.score < 90

    def test_strengths_populated_when_metrics_good(self):
        result = self.grader.grade(
            quality_score=90,
            episode_consistency=0.92,
            signal_health=_default_signal_health(),
            policy_fit_score=0.85,
            episode_ranking=_default_episode_ranking(),
            action_divergence=0.02,
            episode_count=200,
            policy_type="diffusion_policy",
        )
        assert len(result.strengths) > 0

    def test_problems_populated_when_metrics_bad(self):
        result = self.grader.grade(
            quality_score=50,
            episode_consistency=0.40,
            signal_health=_default_signal_health(dead_dimensions=[0]),
            policy_fit_score=0.40,
            episode_ranking=_default_episode_ranking(remove_count=12),
            action_divergence=0.18,
            episode_count=30,
            policy_type="diffusion_policy",
        )
        assert len(result.problems) > 0

    def test_top_action_always_present(self):
        result = self.grader.grade(
            quality_score=70,
            episode_consistency=0.70,
            signal_health=_default_signal_health(),
            policy_fit_score=0.70,
            episode_ranking=_default_episode_ranking(remove_count=2),
            action_divergence=0.05,
            episode_count=50,
            policy_type="act",
        )
        assert result.top_action != ""

    def test_result_has_correct_metadata(self):
        result = self.grader.grade(
            quality_score=80,
            episode_consistency=0.85,
            signal_health=_default_signal_health(),
            policy_fit_score=0.80,
            episode_ranking=_default_episode_ranking(),
            action_divergence=0.03,
            episode_count=100,
            policy_type="act",
        )
        assert result.episode_count == 100
        assert result.policy_type == "act"

    def test_policy_name_normalization(self):
        """Unknown policy names should be normalized."""
        result = self.grader.grade(
            quality_score=80,
            episode_consistency=0.85,
            signal_health=_default_signal_health(),
            policy_fit_score=0.80,
            episode_ranking=_default_episode_ranking(),
            action_divergence=0.03,
            episode_count=100,
            policy_type="action_chunking_transformer",
        )
        assert result.policy_type == "act"

    def test_high_divergence_penalizes_score(self):
        base = self.grader.grade(
            quality_score=80,
            episode_consistency=0.85,
            signal_health=_default_signal_health(),
            policy_fit_score=0.80,
            episode_ranking=_default_episode_ranking(),
            action_divergence=0.02,
            episode_count=100,
            policy_type="act",
        )
        divergent = self.grader.grade(
            quality_score=80,
            episode_consistency=0.85,
            signal_health=_default_signal_health(),
            policy_fit_score=0.80,
            episode_ranking=_default_episode_ranking(),
            action_divergence=0.20,
            episode_count=100,
            policy_type="act",
        )
        assert base.score > divergent.score
