"""Tests for the orbit suggester module."""

from unittest.mock import MagicMock

import numpy as np
import pandas as pd
import pytest

from orbit.suggester import TrainingSuggestion, generate_suggestion, _pick_batch_size, _pick_chunk_size


def _make_episode_data(n_episodes=50, ep_length=100, action_dim=6):
    """Create synthetic episode data for testing."""
    rows = []
    for ep_idx in range(n_episodes):
        for frame_idx in range(ep_length):
            t = frame_idx / ep_length
            action = [np.sin(t + i * 0.5) * 0.5 for i in range(action_dim)]
            state = [np.cos(t + i * 0.3) * 0.3 for i in range(action_dim)]
            rows.append({
                "episode_index": ep_idx,
                "frame_index": frame_idx,
                "action": action,
                "state": state,
            })
    return pd.DataFrame(rows)


def _make_info(n_episodes=50, fps=30.0, action_dim=6):
    """Create a mock DatasetInfo."""
    info = MagicMock()
    info.repo_id = "test/dataset"
    info.num_episodes = n_episodes
    info.num_frames = n_episodes * 100
    info.fps = fps
    info.robot_type = "so100"
    info.camera_names = ["observation.image"]
    info.video_keys = ["observation.image"]
    info.shapes = {
        "action": [action_dim],
        "observation.state": [action_dim],
        "observation.image": [3, 480, 640],
    }
    info.duration_seconds = n_episodes * 100 / fps
    return info


class TestGenerateSuggestion:
    """Test the suggestion generator."""

    def test_basic_suggestion(self):
        info = _make_info()
        ep_data = _make_episode_data()
        suggestion = generate_suggestion("test/dataset", info, ep_data)
        assert isinstance(suggestion, TrainingSuggestion)
        assert suggestion.dataset_repo == "test/dataset"
        assert suggestion.policy_type in ("act", "diffusion_policy", "smolvla", "dp3", "bc", "bc_rnn")
        assert suggestion.batch_size >= 2

    def test_policy_auto_selection(self):
        info = _make_info(n_episodes=50)
        ep_data = _make_episode_data(n_episodes=50)
        suggestion = generate_suggestion("test/dataset", info, ep_data)
        # With 50 episodes, should pick ACT or similar
        assert suggestion.policy_type is not None

    def test_explicit_policy_selection(self):
        info = _make_info()
        ep_data = _make_episode_data()
        suggestion = generate_suggestion("test/dataset", info, ep_data, policy="act")
        assert suggestion.policy_type == "act"

    def test_batch_size_never_one(self):
        """Batch size must never be 1."""
        assert _pick_batch_size(1, "480p", 4) >= 2
        assert _pick_batch_size(2, "1080p", 4) >= 2
        assert _pick_batch_size(1, "480p", 1) >= 2

    def test_chunk_size_respects_episode_length(self):
        # Short episodes
        cs = _pick_chunk_size("act", 50)
        assert cs <= 50
        assert cs >= 10

        # Long episodes
        cs = _pick_chunk_size("act", 500)
        assert cs <= 100  # capped at 100 for ACT

        # Diffusion policy
        cs = _pick_chunk_size("diffusion_policy", 200)
        assert cs <= 20

    def test_train_command_is_valid_syntax(self):
        info = _make_info()
        ep_data = _make_episode_data()
        suggestion = generate_suggestion("test/dataset", info, ep_data, policy="act")
        cmd = suggestion.train_command
        assert "lerobot-train" in cmd
        assert "--dataset.repo_id=test/dataset" in cmd
        assert "--policy.type=act" in cmd
        assert "--batch_size=" in cmd
        assert "--steps=" in cmd

    def test_warnings_for_small_dataset(self):
        info = _make_info(n_episodes=10)
        ep_data = _make_episode_data(n_episodes=10)
        suggestion = generate_suggestion("test/dataset", info, ep_data, policy="act")
        # Should warn about small dataset
        warning_text = " ".join(suggestion.warnings)
        assert "small" in warning_text.lower() or "episode" in warning_text.lower()

    def test_no_cameras_suggestion(self):
        info = _make_info()
        info.camera_names = []
        info.shapes = {"action": [6], "observation.state": [6]}
        ep_data = _make_episode_data()
        suggestion = generate_suggestion("test/dataset", info, ep_data)
        assert suggestion.camera_count == 0

    def test_steps_scale_with_episodes(self):
        info_small = _make_info(n_episodes=30)
        info_large = _make_info(n_episodes=300)
        ep_small = _make_episode_data(n_episodes=30)
        ep_large = _make_episode_data(n_episodes=300)
        s_small = generate_suggestion("test/small", info_small, ep_small, policy="act")
        s_large = generate_suggestion("test/large", info_large, ep_large, policy="act")
        assert s_large.steps >= s_small.steps

    def test_tips_present(self):
        info = _make_info()
        ep_data = _make_episode_data()
        suggestion = generate_suggestion("test/dataset", info, ep_data)
        assert len(suggestion.tips) > 0
