"""Tests for the ICLR 2025 scaling law advisor."""

import numpy as np
import pytest

from orbit.analyzer.scaling_advisor import ScalingAdvice, compute_scaling_advice


def _make_episodes(n: int, n_dims: int = 6, n_steps: int = 50) -> list[np.ndarray]:
    """Generate synthetic episodes."""
    return [np.random.RandomState(i).randn(n_steps, n_dims) for i in range(n)]


def _make_diverse_episodes(n: int, n_envs: int = 1) -> list[np.ndarray]:
    """Generate episodes with very distinct initial states (simulating environments)."""
    episodes = []
    per_env = n // max(1, n_envs)
    for env in range(n_envs):
        # Large offset per environment to guarantee detectability
        offset = np.array([env * 100.0] * 6)
        for i in range(per_env):
            rng = np.random.RandomState(env * 1000 + i)
            ep = rng.randn(50, 6) * 0.1 + offset
            episodes.append(ep)
    return episodes


class TestScalingAdvice:
    def test_basic_output(self):
        episodes = _make_episodes(50)
        advice = compute_scaling_advice(50, episodes)

        assert isinstance(advice, ScalingAdvice)
        assert advice.current_episodes == 50
        assert advice.estimated_diversity >= 1
        assert advice.bottleneck in ("diversity", "quantity", "quality")
        assert advice.advice != ""

    def test_low_episode_count_recommends_more(self):
        episodes = _make_episodes(10)
        advice = compute_scaling_advice(10, episodes)

        # With only 10 episodes, should recommend more
        assert advice.recommended_episodes >= 10

    def test_diverse_dataset(self):
        # Many distinct environments
        episodes = _make_diverse_episodes(200, n_envs=20)
        advice = compute_scaling_advice(200, episodes)

        # Should detect diversity
        assert advice.estimated_diversity >= 2

    def test_single_environment(self):
        # All episodes start from same state
        episodes = _make_diverse_episodes(100, n_envs=1)
        advice = compute_scaling_advice(100, episodes)

        # Low diversity should be detected
        assert advice.estimated_diversity <= 5

    def test_expected_success_bounds(self):
        episodes = _make_episodes(50)
        advice = compute_scaling_advice(50, episodes)

        assert 0.0 <= advice.expected_success_at_current <= 1.0
        assert 0.0 <= advice.expected_success_at_recommended <= 1.0

    def test_recommended_improves_over_current(self):
        episodes = _make_episodes(30)
        advice = compute_scaling_advice(30, episodes)

        # Recommendations should predict improvement
        assert advice.expected_success_at_recommended >= advice.expected_success_at_current - 0.01

    def test_with_state_episodes(self):
        actions = _make_episodes(50, n_dims=6)
        states = _make_episodes(50, n_dims=12)
        advice = compute_scaling_advice(50, actions, state_episodes=states)
        assert isinstance(advice, ScalingAdvice)

    def test_small_dataset(self):
        episodes = _make_episodes(3)
        advice = compute_scaling_advice(3, episodes)
        assert advice.estimated_diversity >= 1
