"""Tests for the problem-based grading catalog.

Each test verifies a specific problem detection and its penalty in isolation.
"""

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


def _base_kwargs(**overrides):
    """Default kwargs that produce a perfect score (grade A)."""
    kwargs = dict(
        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="act",
    )
    kwargs.update(overrides)
    return kwargs


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

    def test_dead_joints_penalty_proportional(self):
        """Dead joint penalty is ratio * 40."""
        # 1/14 dead → penalty = round(1/14 * 40) = round(2.86) = 3
        result = self.grader.grade(**_base_kwargs(
            signal_health=_default_signal_health(dead_dimensions=[3], total_joints=14),
        ))
        perfect = self.grader.grade(**_base_kwargs())
        assert result.score < perfect.score
        assert any("dead" in p.lower() for p in result.problems)

    def test_dead_joints_critical_above_60pct(self):
        """More than 60% dead joints → critical failure (F).

        Threshold is 60% (not 30%) because many real datasets have
        padded/unused action dimensions (e.g., 3/7 constant dims in
        berkeley_cable_routing).  Only a true majority indicates
        hardware failure.
        """
        result = self.grader.grade(**_base_kwargs(
            signal_health=_default_signal_health(dead_dimensions=[0, 1, 2, 3], total_joints=6),
        ))
        assert result.grade == "F"
        assert result.score == 0

    def test_clipping_penalty_capped(self):
        """Clipping penalty is min(8, 2 + n_clipping)."""
        result = self.grader.grade(**_base_kwargs(
            signal_health=_default_signal_health(clipping_joints=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
        ))
        perfect = self.grader.grade(**_base_kwargs())
        # Penalty capped at 8
        assert perfect.score - result.score == 8

    def test_episode_sufficiency_below_40pct(self):
        """Below 40% of ideal episodes → -20 penalty."""
        # ACT ideal=50, 40% = 20. 15 episodes = 30% of ideal.
        result = self.grader.grade(**_base_kwargs(episode_count=15))
        assert any("episode" in p.lower() for p in result.problems)

    def test_episode_sufficiency_at_ideal(self):
        """At or above ideal → no penalty, listed as strength."""
        result = self.grader.grade(**_base_kwargs(episode_count=100))
        assert any("sufficient" in s.lower() for s in result.strengths)


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

    def test_low_consistency_penalty(self):
        result = self.grader.grade(**_base_kwargs(episode_consistency=0.30))
        assert any("consistency" in p.lower() for p in result.problems)

    def test_high_consistency_strength(self):
        result = self.grader.grade(**_base_kwargs(episode_consistency=0.92))
        assert any("consistency" in s.lower() for s in result.strengths)


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

    def test_high_divergence_unimodal_penalty(self):
        """High divergence on ACT (unimodal) should penalize."""
        result = self.grader.grade(**_base_kwargs(
            action_divergence=0.45,
            policy_type="act",
        ))
        assert any("divergence" in p.lower() for p in result.problems)

    def test_moderate_divergence_multimodal_halved(self):
        """Moderate divergence on diffusion_policy gets halved penalty.

        Diffusion Policy handles multimodality by design (Chi et al. 2023),
        so the penalty is halved compared to unimodal policies.
        """
        uni = self.grader.grade(**_base_kwargs(
            action_divergence=0.25,
            policy_type="act",
            episode_count=100,
        ))
        multi = self.grader.grade(**_base_kwargs(
            action_divergence=0.25,
            policy_type="diffusion_policy",
            episode_count=200,
        ))
        # Multimodal should score higher (less penalty)
        assert multi.score > uni.score

    def test_low_divergence_strength(self):
        result = self.grader.grade(**_base_kwargs(action_divergence=0.01))
        assert any("divergence" in s.lower() for s in result.strengths)


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

    def test_proxy_go_low_loss_boost(self):
        """GO + val_loss < 0.1 → +5 point boost."""
        from unittest.mock import MagicMock
        proxy = MagicMock()
        proxy.go_no_go = "GO"
        proxy.validation_loss_final = 0.05
        proxy.reason = "Good"

        base = self.grader.grade(**_base_kwargs())
        boosted = self.grader.grade(**_base_kwargs(proxy_result=proxy))
        # Boost only applies if score < 100
        assert boosted.score >= base.score

    def test_proxy_nogo_penalty(self):
        """NO-GO → -10 penalty."""
        from unittest.mock import MagicMock
        proxy = MagicMock()
        proxy.go_no_go = "NO-GO"
        proxy.validation_loss_final = 0.9
        proxy.reason = "Not learnable"

        base = self.grader.grade(**_base_kwargs())
        penalized = self.grader.grade(**_base_kwargs(proxy_result=proxy))
        assert penalized.score < base.score

    def test_no_proxy_suggests_for_b_c(self):
        """When proxy not run and grade is B/C, suggest running it."""
        result = self.grader.grade(**_base_kwargs(
            episode_consistency=0.60,
            policy_fit_score=0.65,
        ))
        if result.grade in ("B", "C"):
            assert any("--proxy" in p for p in result.problems)


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

    def test_high_outlier_fraction_penalty(self):
        """Outlier fraction > 25% → -15 penalty."""
        from unittest.mock import MagicMock
        outlier = MagicMock()
        outlier.outlier_fraction = 0.30
        outlier.outlier_count = 15
        outlier.total_episodes = 50
        outlier.dataset_health = "problematic"

        base = self.grader.grade(**_base_kwargs())
        penalized = self.grader.grade(**_base_kwargs(outlier_report=outlier))
        assert penalized.score < base.score
        assert any("outlier" in p.lower() for p in penalized.problems)

    def test_low_outlier_fraction_strength(self):
        """Outlier fraction < 5% → listed as strength."""
        from unittest.mock import MagicMock
        outlier = MagicMock()
        outlier.outlier_fraction = 0.02
        outlier.outlier_count = 1
        outlier.total_episodes = 50
        outlier.dataset_health = "healthy"

        result = self.grader.grade(**_base_kwargs(outlier_report=outlier))
        assert any("outlier" in s.lower() for s in result.strengths)


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

    def test_grade_a_threshold(self):
        """Score >= 88 → Grade A."""
        result = self.grader.grade(**_base_kwargs())
        assert result.grade == "A"
        assert result.score >= 88

    def test_grade_b_range(self):
        """Score 80-87 → Grade B."""
        # Create mild issues to get into B range
        result = self.grader.grade(**_base_kwargs(
            signal_health=_default_signal_health(clipping_joints=[0, 1, 2, 3, 4, 5]),
            action_divergence=0.15,
        ))
        assert result.grade in ("A", "B")  # may still be A depending on exact penalties

    def test_grade_f_critical(self):
        """Critical failure → Grade F, score 0."""
        result = self.grader.grade(**_base_kwargs(episode_count=1))
        assert result.grade == "F"
        assert result.score == 0


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

    def test_vision_penalty_applied(self):
        """Vision quality penalties should reduce grade."""
        from unittest.mock import MagicMock
        vision = MagicMock()
        vision.total_penalty = -18
        vision.grade_adjustments = {
            "incomplete_tasks": -10,
            "inconsistent_tasks": -8,
        }

        base = self.grader.grade(**_base_kwargs())
        penalized = self.grader.grade(**_base_kwargs(vision_quality=vision))
        assert penalized.score == base.score - 18

    def test_no_vision_penalty_when_zero(self):
        """Zero penalty from vision should not change score."""
        from unittest.mock import MagicMock
        vision = MagicMock()
        vision.total_penalty = 0
        vision.grade_adjustments = {}

        base = self.grader.grade(**_base_kwargs())
        same = self.grader.grade(**_base_kwargs(vision_quality=vision))
        assert same.score == base.score
