"""Tests for multi-framework suggestion generation."""

from unittest.mock import MagicMock

import numpy as np
import pandas as pd
import pytest

from orbit.suggester import generate_multi_framework_suggestion


def _make_episode_data(n_episodes=50, ep_length=100, action_dim=6):
    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):
    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 TestMultiFrameworkSuggestion:
    def test_single_framework(self):
        info = _make_info()
        ep_data = _make_episode_data()

        base, configs = generate_multi_framework_suggestion(
            "test/dataset", info, ep_data, framework="lerobot"
        )
        assert len(configs) == 1
        assert configs[0][0] == "lerobot"
        assert "lerobot-train" in configs[0][1].command

    def test_all_frameworks(self):
        info = _make_info()
        ep_data = _make_episode_data()

        base, configs = generate_multi_framework_suggestion(
            "test/dataset", info, ep_data, framework="all"
        )
        # Should have configs for all frameworks
        assert len(configs) >= 5  # lerobot, openvla, openpi, groot, osmo, custom
        framework_names = [c[0] for c in configs]
        assert "lerobot" in framework_names
        assert "openvla" in framework_names
        assert "openpi" in framework_names
        assert "groot" in framework_names
        assert "osmo" in framework_names

    def test_openvla_framework(self):
        info = _make_info()
        ep_data = _make_episode_data()

        base, configs = generate_multi_framework_suggestion(
            "test/dataset", info, ep_data, framework="openvla"
        )
        assert len(configs) == 1
        assert configs[0][0] == "openvla"
        assert "torchrun" in configs[0][1].command

    def test_osmo_framework(self):
        info = _make_info()
        ep_data = _make_episode_data()

        base, configs = generate_multi_framework_suggestion(
            "test/dataset", info, ep_data, framework="osmo"
        )
        assert len(configs) == 1
        assert configs[0][0] == "osmo"
        assert "workflow:" in configs[0][1].command

    def test_base_suggestion_returned(self):
        info = _make_info()
        ep_data = _make_episode_data()

        base, configs = generate_multi_framework_suggestion(
            "test/dataset", info, ep_data, policy="act", framework="lerobot"
        )
        assert base.policy_type == "act"
        assert base.dataset_repo == "test/dataset"
