"""Tests for the PolicyFitAnalyzer."""

import numpy as np
import pytest

from orbit.analyzer.policy_fit import PolicyFitAnalyzer, PolicyFit, POLICY_PROFILES


class TestPolicyFitAnalyzer:
    """Test the policy fit analyzer."""

    def setup_method(self):
        self.analyzer = PolicyFitAnalyzer()
        self.rng = np.random.RandomState(42)

    def _make_episodes(self, n_eps=50, n_steps=100, action_dim=7):
        return [self.rng.randn(n_steps, action_dim) for _ in range(n_eps)]

    def test_basic_analysis(self):
        episodes = self._make_episodes()
        result = self.analyzer.analyze(episodes)
        assert isinstance(result, PolicyFit)
        assert 0.0 <= result.fit_score <= 1.0
        assert result.policy_type is not None

    def test_auto_policy_detection(self):
        episodes = self._make_episodes()
        result = self.analyzer.analyze(episodes, policy_type="auto")
        assert result.policy_type in POLICY_PROFILES

    def test_specific_policy(self):
        episodes = self._make_episodes()
        result = self.analyzer.analyze(episodes, policy_type="act")
        assert result.policy_type == "act"

    def test_episode_sufficiency_low(self):
        episodes = self._make_episodes(n_eps=5)
        result = self.analyzer.analyze(episodes, episode_count=5)
        assert result.episode_sufficiency < 0.7

    def test_episode_sufficiency_high(self):
        episodes = self._make_episodes(n_eps=200)
        result = self.analyzer.analyze(episodes, episode_count=200, policy_type="act")
        assert result.episode_sufficiency > 0.7

    def test_multimodality_warning_for_bc(self):
        episodes = self._make_episodes()
        result = self.analyzer.analyze(
            episodes,
            policy_type="bc",
            multimodality_score=0.8,
            consistency_score=0.4,
        )
        assert any("multimodal" in w.lower() or "diffusion" in w.lower()
                    for w in result.specific_warnings)

    def test_recommended_policies_present(self):
        episodes = self._make_episodes()
        result = self.analyzer.analyze(episodes, policy_type="auto")
        assert len(result.recommended_policies) > 0
        assert len(result.recommended_policies) <= 3

    def test_empty_episodes(self):
        result = self.analyzer.analyze([])
        assert result.fit_score == 0.5

    def test_unknown_policy_defaults(self):
        episodes = self._make_episodes()
        result = self.analyzer.analyze(episodes, policy_type="my_custom_policy")
        assert result.policy_type == "diffusion_policy"  # Falls back

    def test_all_scores_bounded(self):
        episodes = self._make_episodes()
        result = self.analyzer.analyze(episodes)
        assert 0.0 <= result.fit_score <= 1.0
        assert 0.0 <= result.episode_sufficiency <= 1.0
        assert 0.0 <= result.action_space_compatibility <= 1.0
        assert 0.0 <= result.consistency_compatibility <= 1.0

    def test_bimanual_boosts_act(self):
        episodes = self._make_episodes(action_dim=14)
        result_act = self.analyzer.analyze(
            episodes, policy_type="act",
            metadata={"is_bimanual": True},
        )
        result_bc = self.analyzer.analyze(
            episodes, policy_type="bc",
            metadata={"is_bimanual": True},
        )
        # ACT should generally fit bimanual data better
        # (not strictly enforced but should be in recommended)
        assert result_act.fit_score >= 0  # Basic sanity


class TestPolicyProfiles:
    """Test that policy profiles are well-formed."""

    def test_all_profiles_have_required_keys(self):
        required = [
            "full_name", "min_episodes_simple", "min_episodes_complex",
            "sweet_spot_low", "sweet_spot_high", "consistency_sensitivity",
            "multimodality_tolerance", "action_dim_range",
        ]
        for name, profile in POLICY_PROFILES.items():
            for key in required:
                assert key in profile, f"{name} missing {key}"

    def test_sweet_spot_above_minimum(self):
        for name, profile in POLICY_PROFILES.items():
            assert profile["sweet_spot_low"] >= profile["min_episodes_simple"], \
                f"{name}: sweet_spot_low < min_episodes_simple"

    def test_normalize_policy_name_expanded_aliases(self):
        from orbit.analyzer.policy_fit import normalize_policy_name
        assert normalize_policy_name('BC-Z') == 'bc'
        assert normalize_policy_name('RT-1-X') == 'smolvla'  # RT-1 is a VLA, not ACT
        assert normalize_policy_name('LSTM') == 'bc_rnn'
        assert normalize_policy_name('Octo-Base') == 'diffusion_policy'
        assert normalize_policy_name('π₀') == 'diffusion_policy'
        assert normalize_policy_name('CQL') == 'bc'
        assert normalize_policy_name('TD3+BC') == 'bc'
        assert normalize_policy_name('HULC') == 'bc'
        assert normalize_policy_name('3D Diffusion Policy') == 'dp3'
        assert normalize_policy_name('OpenVLA-7B') == 'smolvla'

    def test_normalize_unknown_logs_warning(self, caplog):
        import logging
        from orbit.analyzer.policy_fit import normalize_policy_name
        with caplog.at_level(logging.WARNING):
            result = normalize_policy_name("totally_unknown_policy_xyz")
        assert result == "diffusion_policy"
        assert "totally_unknown_policy_xyz" in caplog.text
