"""Tests for the meta-learner success prediction system."""

import json
import os
import tempfile
from pathlib import Path
from unittest.mock import patch

import numpy as np
import pytest


class TestMetaFeatures:
    """Test feature extraction and conversion."""

    def test_to_array_has_correct_length(self):
        from orbit.analyzer.meta_learner import FEATURE_NAMES, MetaFeatures
        f = MetaFeatures()
        arr = f.to_array()
        assert len(arr) == len(FEATURE_NAMES)

    def test_to_dict_keys_match_feature_names(self):
        from orbit.analyzer.meta_learner import FEATURE_NAMES, MetaFeatures
        f = MetaFeatures(episode_count=50, action_dim=7)
        d = f.to_dict()
        assert set(d.keys()) == set(FEATURE_NAMES)
        assert d["episode_count"] == 50.0
        assert d["action_dim"] == 7.0

    def test_extract_features_from_analysis_json(self):
        from orbit.analyzer.meta_learner import extract_features
        analysis = {
            "repo_id": "lerobot/test_dataset",
            "num_episodes": 100,
            "num_frames": 50000,
            "camera_names": ["top"],
            "coverage": {"episode_consistency": 0.85},
            "quality_score": 0.78,
            "task_context": {"action_dim": 7, "is_bimanual": False},
            "policy_fit": {"fit_score": 0.88},
            "signal_diagnostics": {
                "dead_dimensions": [],
                "joint_diagnostics": [
                    {"is_clipping": False} for _ in range(7)
                ],
            },
            "action_divergence": {"score": 0.12},
            "episode_ranking": {
                "keep_count": 90, "review_count": 7, "remove_count": 3
            },
            "consistency_diagnostics": {
                "overall_consistency": 0.72,
                "metrics": {
                    "path_length_cv": 0.35,
                    "joint_path_length_cv": 0.30,
                    "cartesian_jerk_cv": 0.45,
                    "joint_jerk_cv": 0.40,
                    "curvature_cv": 0.25,
                    "effort_cv": 0.28,
                    "joint_limit_cv": 0.20,
                    "orientation_cv": 0.15,
                },
            },
        }
        f = extract_features(analysis)
        assert f.episode_count == 100.0
        assert f.action_dim == 7.0
        assert f.action_divergence == 0.12
        assert f.consistency_overall == 0.72
        assert f.has_images == 1.0
        assert f.is_simulation == 0.0
        assert f.dead_joint_ratio == 0.0
        assert abs(f.remove_fraction - 0.03) < 0.01

    def test_extract_features_simulation_detection(self):
        from orbit.analyzer.meta_learner import extract_features
        sim = extract_features({"repo_id": "lerobot/pusht_sim"})
        assert sim.is_simulation == 1.0
        real = extract_features({"repo_id": "user/my_franka_data"})
        assert real.is_simulation == 0.0

    def test_extract_features_handles_missing_fields(self):
        from orbit.analyzer.meta_learner import extract_features
        f = extract_features({})
        # Should not crash, uses defaults
        assert f.episode_count == 0.0
        assert f.quality_score == 0.5


class TestBootstrap:
    """Test bootstrap data from published papers."""

    def test_bootstrap_creates_training_data(self, tmp_path):
        from orbit.analyzer import meta_learner as ml
        with patch.object(ml, "TRAINING_DATA_PATH", tmp_path / "data.jsonl"):
            n = ml.bootstrap_training_data()
            assert n >= 40  # at least 40 published results
            assert (tmp_path / "data.jsonl").exists()

    def test_bootstrap_data_is_valid_jsonl(self, tmp_path):
        from orbit.analyzer import meta_learner as ml
        data_path = tmp_path / "data.jsonl"
        with patch.object(ml, "TRAINING_DATA_PATH", data_path):
            ml.bootstrap_training_data()
            with open(data_path) as f:
                for line in f:
                    d = json.loads(line)
                    assert "features" in d
                    assert "policy" in d
                    assert "success_rate" in d
                    assert 0.0 <= d["success_rate"] <= 1.0

    def test_bootstrap_covers_multiple_policies(self, tmp_path):
        from orbit.analyzer import meta_learner as ml
        data_path = tmp_path / "data.jsonl"
        with patch.object(ml, "TRAINING_DATA_PATH", data_path):
            ml.bootstrap_training_data()
            with patch.object(ml, "TRAINING_DATA_PATH", data_path):
                examples = ml.load_training_data()
            policies = set(e.policy for e in examples)
            assert "act" in policies
            assert "bc" in policies
            assert "diffusion_policy" in policies
            assert "bc_rnn" in policies

    def test_bootstrap_covers_diverse_success_rates(self, tmp_path):
        from orbit.analyzer import meta_learner as ml
        data_path = tmp_path / "data.jsonl"
        with patch.object(ml, "TRAINING_DATA_PATH", data_path):
            ml.bootstrap_training_data()
            with patch.object(ml, "TRAINING_DATA_PATH", data_path):
                examples = ml.load_training_data()
            rates = [e.success_rate for e in examples]
            assert min(rates) <= 0.10  # has failures
            assert max(rates) >= 0.90  # has successes
            assert 0.40 <= np.median(rates) <= 0.85  # reasonable median


class TestTraining:
    """Test model training."""

    def test_train_requires_min_examples(self):
        from orbit.analyzer.meta_learner import TrainingExample, train_model
        few = [TrainingExample(
            features={"episode_count": 50}, policy="act",
            success_rate=0.8, source="test"
        ) for _ in range(5)]
        with pytest.raises(ValueError, match="at least 10"):
            train_model(few)

    def test_train_produces_model_file(self, tmp_path):
        from orbit.analyzer import meta_learner as ml
        data_path = tmp_path / "data.jsonl"
        model_path = tmp_path / "model.pkl"
        metadata_path = tmp_path / "meta.json"
        with (
            patch.object(ml, "TRAINING_DATA_PATH", data_path),
            patch.object(ml, "MODEL_PATH", model_path),
            patch.object(ml, "METADATA_PATH", metadata_path),
            patch.object(ml, "META_LEARNER_DIR", tmp_path),
        ):
            ml.bootstrap_training_data()
            examples = ml.load_training_data()
            meta = ml.train_model(examples)
            assert model_path.exists()
            assert metadata_path.exists()
            assert meta["n_samples"] >= 40
            assert meta["r2_score"] > 0.3  # should fit reasonably well

    def test_train_r2_is_reasonable(self, tmp_path):
        from orbit.analyzer import meta_learner as ml
        data_path = tmp_path / "data.jsonl"
        model_path = tmp_path / "model.pkl"
        metadata_path = tmp_path / "meta.json"
        with (
            patch.object(ml, "TRAINING_DATA_PATH", data_path),
            patch.object(ml, "MODEL_PATH", model_path),
            patch.object(ml, "METADATA_PATH", metadata_path),
            patch.object(ml, "META_LEARNER_DIR", tmp_path),
        ):
            ml.bootstrap_training_data()
            meta = ml.train_model()
            # R² should be decent on training data with our features
            assert meta["r2_score"] > 0.5, f"R² too low: {meta['r2_score']}"


class TestPrediction:
    """Test success rate predictions."""

    @pytest.fixture(autouse=True)
    def setup_model(self, tmp_path):
        """Bootstrap and train model in temp dir for all prediction tests."""
        from orbit.analyzer import meta_learner as ml
        self.tmp = tmp_path
        self.patches = [
            patch.object(ml, "TRAINING_DATA_PATH", tmp_path / "data.jsonl"),
            patch.object(ml, "MODEL_PATH", tmp_path / "model.pkl"),
            patch.object(ml, "METADATA_PATH", tmp_path / "meta.json"),
            patch.object(ml, "META_LEARNER_DIR", tmp_path),
        ]
        for p in self.patches:
            p.start()
        ml.bootstrap_training_data()
        ml.train_model()
        yield
        for p in self.patches:
            p.stop()

    def test_clean_data_predicts_high(self):
        from orbit.analyzer.meta_learner import MetaFeatures, predict_success
        clean = MetaFeatures(
            episode_count=200, action_dim=7, action_divergence=0.05,
            consistency_overall=0.85, quality_score=0.90, is_simulation=1,
            episode_consistency=0.90, policy_fit_score=0.92, mean_episode_length=400,
        )
        pred = predict_success(clean, "diffusion_policy")
        assert pred.model_available
        assert pred.predicted_success_rate >= 0.70, \
            f"Clean data should predict high, got {pred.predicted_success_rate}"

    def test_noisy_data_predicts_low(self):
        from orbit.analyzer.meta_learner import MetaFeatures, predict_success
        noisy = MetaFeatures(
            episode_count=200, action_dim=7, action_divergence=0.28,
            consistency_overall=0.40, quality_score=0.55, is_simulation=1,
            episode_consistency=0.50, policy_fit_score=0.50, mean_episode_length=500,
        )
        pred = predict_success(noisy, "bc")
        assert pred.predicted_success_rate <= 0.50, \
            f"Noisy data with BC should predict low, got {pred.predicted_success_rate}"

    def test_prediction_has_confidence_interval(self):
        from orbit.analyzer.meta_learner import MetaFeatures, predict_success
        f = MetaFeatures(episode_count=50, action_dim=14)
        pred = predict_success(f, "act")
        assert pred.confidence_low <= pred.predicted_success_rate
        assert pred.predicted_success_rate <= pred.confidence_high
        assert pred.confidence_width >= 0
        assert pred.confidence_level in ("HIGH", "MEDIUM", "LOW")

    def test_prediction_bounded_0_to_1(self):
        from orbit.analyzer.meta_learner import MetaFeatures, predict_success
        # Extreme good
        f = MetaFeatures(
            episode_count=10000, action_divergence=0.0, consistency_overall=1.0,
            quality_score=1.0, policy_fit_score=1.0, episode_consistency=1.0,
        )
        pred = predict_success(f, "diffusion_policy")
        assert 0.0 <= pred.predicted_success_rate <= 1.0
        assert 0.0 <= pred.confidence_low <= 1.0
        assert 0.0 <= pred.confidence_high <= 1.0

        # Extreme bad
        f2 = MetaFeatures(
            episode_count=1, action_divergence=1.0, consistency_overall=0.0,
            quality_score=0.0, policy_fit_score=0.0,
        )
        pred2 = predict_success(f2, "bc")
        assert 0.0 <= pred2.predicted_success_rate <= 1.0

    def test_no_model_returns_unavailable(self, tmp_path):
        from orbit.analyzer import meta_learner as ml
        from orbit.analyzer.meta_learner import MetaFeatures, predict_success
        # Point to empty dir
        with patch.object(ml, "MODEL_PATH", tmp_path / "nonexistent.pkl"):
            pred = predict_success(MetaFeatures(), "act")
            assert not pred.model_available
            assert pred.confidence_level == "LOW"

    def test_clean_aloha_matches_published_range(self):
        """ALOHA Transfer Cube human: published 82%. Prediction should be in range."""
        from orbit.analyzer.meta_learner import MetaFeatures, predict_success
        aloha = MetaFeatures(
            episode_count=50, action_dim=14, action_divergence=0.14,
            consistency_overall=0.65, quality_score=0.75, is_bimanual=1,
            has_images=1, episode_consistency=0.72, policy_fit_score=0.88,
            mean_episode_length=400,
        )
        pred = predict_success(aloha, "act")
        # Should be within reasonable range of 82%
        assert pred.confidence_low <= 0.82 <= pred.confidence_high, \
            f"Published 82% should be in CI [{pred.confidence_low}, {pred.confidence_high}]"


class TestAddTrainingExample:
    """Test adding new training data."""

    def test_add_example_grows_dataset(self, tmp_path):
        from orbit.analyzer import meta_learner as ml
        data_path = tmp_path / "data.jsonl"
        model_path = tmp_path / "model.pkl"
        metadata_path = tmp_path / "meta.json"
        with (
            patch.object(ml, "TRAINING_DATA_PATH", data_path),
            patch.object(ml, "MODEL_PATH", model_path),
            patch.object(ml, "METADATA_PATH", metadata_path),
            patch.object(ml, "META_LEARNER_DIR", tmp_path),
        ):
            ml.bootstrap_training_data()
            n_before = len(ml.load_training_data())

            new = ml.TrainingExample(
                features={"episode_count": 100, "action_divergence": 0.1},
                policy="act", success_rate=0.85, source="test_report",
            )
            n_after = ml.add_training_example(new)
            assert n_after == n_before + 1

    def test_add_from_report_works(self, tmp_path):
        from orbit.analyzer import meta_learner as ml
        data_path = tmp_path / "data.jsonl"
        with (
            patch.object(ml, "TRAINING_DATA_PATH", data_path),
            patch.object(ml, "MODEL_PATH", tmp_path / "model.pkl"),
            patch.object(ml, "METADATA_PATH", tmp_path / "meta.json"),
            patch.object(ml, "META_LEARNER_DIR", tmp_path),
        ):
            ml.bootstrap_training_data()
            analysis = {
                "num_episodes": 50, "quality_score": 0.78,
                "coverage": {"episode_consistency": 0.80},
                "task_context": {"action_dim": 7},
            }
            n = ml.add_from_report_and_analysis(analysis, "act", 0.82)
            assert n >= 41
