"""Tests for orbit.analyzer.quality."""

from __future__ import annotations

import pytest

from orbit.analyzer.coverage import FullCoverageReport
from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.analyzer.quality import (
    QualityReport,
    ReadyAssessment,
    compute_quality_score,
    ready_to_train_assessment,
)

# ---------------------------------------------------------------------------
# Fixtures / helpers
# ---------------------------------------------------------------------------


def _make_info(num_episodes: int = 50, num_frames: int = 12500) -> DatasetInfo:
    """Create a DatasetInfo with configurable episode count."""
    return DatasetInfo(
        repo_id="test/dataset",
        num_episodes=num_episodes,
        num_frames=num_frames,
        fps=30.0,
        robot_type="aloha",
        camera_names=["cam_high"],
        video_keys=["observation.images.cam_high"],
        shapes={"action": [14]},
        duration_seconds=num_frames / 30.0,
    )


def _make_coverage(
    position: float = 1.0,
    action: float = 1.0,
    consistency: float = 1.0,
    temporal: float = 1.0,
) -> FullCoverageReport:
    """Create a FullCoverageReport with configurable scores."""
    return FullCoverageReport(
        position_diversity=position,
        action_diversity=action,
        episode_consistency=consistency,
        temporal_coverage=temporal,
    )


# ---------------------------------------------------------------------------
# compute_quality_score
# ---------------------------------------------------------------------------


class TestComputeQualityScore:
    """Tests for compute_quality_score."""

    def test_perfect_coverage_score(self):
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(1.0, 1.0, 1.0, 1.0)
        report = compute_quality_score(info, coverage)

        assert isinstance(report, QualityReport)
        assert report.overall_score == pytest.approx(1.0)
        assert report.ready_to_train is True
        assert all(s == pytest.approx(1.0) for s in report.component_scores.values())

    def test_zero_coverage_score(self):
        info = _make_info(num_episodes=0)
        coverage = _make_coverage(0.0, 0.0, 0.0, 0.0)
        report = compute_quality_score(info, coverage)

        assert report.overall_score == pytest.approx(0.0)
        assert report.ready_to_train is False

    def test_mixed_scores(self):
        info = _make_info(num_episodes=25)
        coverage = _make_coverage(0.8, 0.4, 0.9, 0.6)
        report = compute_quality_score(info, coverage)

        # episode_count = min(1.0, 25/50) = 0.5
        # weighted = 0.20*0.5 + 0.10*0.8 + 0.15*0.4 + 0.35*0.9 + 0.20*0.6
        expected = 0.20 * 0.5 + 0.10 * 0.8 + 0.15 * 0.4 + 0.35 * 0.9 + 0.20 * 0.6
        assert report.overall_score == pytest.approx(expected)

    def test_score_bounded_zero_to_one(self):
        # Even with more episodes than expected, score should cap at 1.0
        info = _make_info(num_episodes=200)
        coverage = _make_coverage(1.0, 1.0, 1.0, 1.0)
        report = compute_quality_score(info, coverage)

        assert 0.0 <= report.overall_score <= 1.0

    def test_estimated_success_rate_bounds(self):
        """Success rate bounds must always be within [0, 1]."""
        for ep_count in [0, 5, 25, 50, 100]:
            for score_val in [0.0, 0.3, 0.5, 0.8, 1.0]:
                info = _make_info(num_episodes=ep_count)
                coverage = _make_coverage(score_val, score_val, score_val, score_val)
                report = compute_quality_score(info, coverage)
                low, high = report.estimated_success_rate
                assert 0.0 <= low <= 1.0, f"Low bound {low} out of range"
                assert 0.0 <= high <= 1.0, f"High bound {high} out of range"
                assert low <= high, f"Low {low} > High {high}"


class TestConfidenceLevels:
    """Tests for confidence level based on episode count."""

    def test_high_confidence(self):
        info = _make_info(num_episodes=51)
        coverage = _make_coverage()
        report = compute_quality_score(info, coverage)
        assert report.confidence == "high"

    def test_high_confidence_boundary(self):
        info = _make_info(num_episodes=100)
        coverage = _make_coverage()
        report = compute_quality_score(info, coverage)
        assert report.confidence == "high"

    def test_medium_confidence(self):
        info = _make_info(num_episodes=35)
        coverage = _make_coverage()
        report = compute_quality_score(info, coverage)
        assert report.confidence == "medium"

    def test_medium_confidence_lower_boundary(self):
        info = _make_info(num_episodes=20)
        coverage = _make_coverage()
        report = compute_quality_score(info, coverage)
        assert report.confidence == "medium"

    def test_low_confidence(self):
        info = _make_info(num_episodes=19)
        coverage = _make_coverage()
        report = compute_quality_score(info, coverage)
        assert report.confidence == "low"

    def test_low_confidence_minimal(self):
        info = _make_info(num_episodes=1)
        coverage = _make_coverage()
        report = compute_quality_score(info, coverage)
        assert report.confidence == "low"


# ---------------------------------------------------------------------------
# ready_to_train_assessment
# ---------------------------------------------------------------------------


class TestReadyToTrainAssessment:
    """Tests for ready_to_train_assessment."""

    def test_yes_verdict(self):
        """Score >= 0.65 with no component < 0.3 -> YES."""
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(0.8, 0.7, 0.9, 0.7)
        quality = compute_quality_score(info, coverage)
        assessment = ready_to_train_assessment(quality)

        assert isinstance(assessment, ReadyAssessment)
        assert assessment.verdict == "YES"
        assert len(assessment.blocking_issues) == 0

    def test_not_yet_verdict_overall_in_range(self):
        """Score between 0.45 and 0.65 -> NOT YET."""
        info = _make_info(num_episodes=25)
        coverage = _make_coverage(0.6, 0.5, 0.7, 0.5)
        quality = compute_quality_score(info, coverage)
        # Verify score is in the NOT YET range
        assert 0.45 <= quality.overall_score < 0.65
        assessment = ready_to_train_assessment(quality)
        assert assessment.verdict == "NOT YET"

    def test_not_yet_verdict_critical_component(self):
        """Any component < 0.3 -> NOT YET even if overall is high."""
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(0.9, 0.2, 0.9, 0.9)
        quality = compute_quality_score(info, coverage)
        assessment = ready_to_train_assessment(quality)

        assert assessment.verdict == "NOT YET"
        assert len(assessment.blocking_issues) > 0

    def test_no_verdict(self):
        """Score < 0.45 and no component < 0.3 -> NO."""
        # Need overall < 0.45 but no individual component below 0.3
        # so the "NOT YET" (any component < 0.3) branch doesn't fire.
        info = _make_info(num_episodes=15)  # episode_count = 0.3
        coverage = _make_coverage(0.35, 0.35, 0.35, 0.35)
        quality = compute_quality_score(info, coverage)
        assert quality.overall_score < 0.45
        assert all(s >= 0.3 for s in quality.component_scores.values())
        assessment = ready_to_train_assessment(quality)
        assert assessment.verdict == "NO"

    def test_blocking_issues_populated(self):
        """Components below 0.3 should appear as blocking issues."""
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(0.1, 0.2, 0.9, 0.9)
        quality = compute_quality_score(info, coverage)
        assessment = ready_to_train_assessment(quality)

        assert len(assessment.blocking_issues) >= 2
        assert any("Position" in issue for issue in assessment.blocking_issues)
        assert any("Action" in issue for issue in assessment.blocking_issues)

    def test_nice_to_haves_populated(self):
        """Components between 0.3 and 0.5 should appear as nice-to-haves."""
        info = _make_info(num_episodes=50)
        coverage = _make_coverage(0.4, 0.4, 0.9, 0.9)
        quality = compute_quality_score(info, coverage)
        assessment = ready_to_train_assessment(quality)

        assert len(assessment.nice_to_haves) >= 2
