"""Tests for the proxy trainer module."""

import numpy as np
import pytest

torch = pytest.importorskip("torch")

from orbit.analyzer.proxy_trainer import ProxyTrainingResult, run_proxy_training


def _make_learnable_data(n_episodes: int = 30, n_steps: int = 50, n_dims: int = 6) -> tuple:
    """Generate state→action data with a clear learnable mapping."""
    actions = []
    states = []
    for i in range(n_episodes):
        rng = np.random.RandomState(i)
        s = rng.randn(n_steps, n_dims).astype(np.float32)
        # Actions are a simple linear function of states + small noise
        a = (s @ np.eye(n_dims) * 0.5 + rng.randn(n_steps, n_dims) * 0.05).astype(np.float32)
        states.append(s)
        actions.append(a)
    return actions, states


def _make_noise_data(n_episodes: int = 30, n_steps: int = 50, n_dims: int = 6) -> tuple:
    """Generate pure random data — not learnable."""
    actions = []
    states = []
    for i in range(n_episodes):
        rng = np.random.RandomState(i)
        s = rng.randn(n_steps, n_dims).astype(np.float32)
        a = np.random.RandomState(i + 1000).randn(n_steps, n_dims).astype(np.float32)
        states.append(s)
        actions.append(a)
    return actions, states


class TestProxyTrainer:
    def test_learnable_data_gives_go(self):
        actions, states = _make_learnable_data()
        result = run_proxy_training(
            actions, state_episodes=states,
            max_steps=200, device="cpu",
        )

        assert isinstance(result, ProxyTrainingResult)
        assert result.go_no_go == "GO"
        assert result.validation_loss_final < result.validation_loss_curve[0]

    def test_noise_data_gives_no_go(self):
        actions, states = _make_noise_data()
        result = run_proxy_training(
            actions, state_episodes=states,
            max_steps=200, device="cpu",
        )

        assert result.go_no_go in ("NO-GO", "CAUTION")

    def test_without_states_uses_lagged_actions(self):
        actions, _ = _make_learnable_data()
        result = run_proxy_training(
            actions, state_episodes=None,
            max_steps=100, device="cpu",
        )

        assert isinstance(result, ProxyTrainingResult)
        assert result.steps_run == 100

    def test_training_curves_recorded(self):
        actions, states = _make_learnable_data()
        result = run_proxy_training(
            actions, state_episodes=states,
            max_steps=100, checkpoint_every=25, device="cpu",
        )

        assert len(result.training_loss_curve) >= 2
        assert len(result.validation_loss_curve) >= 2

    def test_overfitting_ratio_computed(self):
        actions, states = _make_learnable_data()
        result = run_proxy_training(
            actions, state_episodes=states,
            max_steps=100, device="cpu",
        )

        assert result.overfitting_ratio > 0

    def test_timing_recorded(self):
        actions, states = _make_learnable_data(n_episodes=10)
        result = run_proxy_training(
            actions, state_episodes=states,
            max_steps=50, device="cpu",
        )

        assert result.training_time_seconds >= 0
        assert result.training_time_seconds < 60  # should be fast for tiny data

    def test_dimensions_recorded(self):
        actions, states = _make_learnable_data(n_dims=4)
        result = run_proxy_training(
            actions, state_episodes=states,
            max_steps=50, device="cpu",
        )

        assert result.state_dim == 4
        assert result.action_dim == 4

    def test_too_few_episodes_raises(self):
        actions = [np.random.randn(50, 6).astype(np.float32)]
        with pytest.raises(ValueError, match="Not enough episodes"):
            run_proxy_training(actions, max_steps=10, device="cpu")

    def test_1d_episodes(self):
        actions = [np.random.randn(50).astype(np.float32) for _ in range(20)]
        result = run_proxy_training(actions, max_steps=50, device="cpu")
        assert result.action_dim == 1
