"""Tests for baseline agents."""

import gymnasium as gym
import numpy as np
import pytest

import oncosim  # noqa: F401
from oncosim.agents.random_agent import RandomAgent, evaluate_random
from oncosim.agents.heuristic_agent import (
    BeamSelectionHeuristic,
    DoseFractionationHeuristic,
    AdaptiveRTHeuristic,
    evaluate_heuristic,
)


class TestRandomAgent:
    def test_beam_selection(self):
        env = gym.make("oncosim/BeamSelection-v0")
        agent = RandomAgent(env)
        obs, _ = env.reset(seed=42)
        action = agent.act(obs)
        assert env.action_space.contains(action)
        env.close()

    def test_dose_fractionation(self):
        env = gym.make("oncosim/DoseFractionation-v0")
        agent = RandomAgent(env)
        obs, _ = env.reset(seed=42)
        action = agent.act(obs)
        assert env.action_space.contains(action)
        env.close()

    def test_adaptive_rt(self):
        env = gym.make("oncosim/AdaptiveRT-v0")
        agent = RandomAgent(env)
        obs, _ = env.reset(seed=42)
        action = agent.act(obs)
        assert env.action_space.contains(action)
        env.close()

    def test_evaluate_random(self):
        env = gym.make("oncosim/BeamSelection-v0")
        results = evaluate_random(env, n_episodes=5)
        assert "mean_reward" in results
        assert np.isfinite(results["mean_reward"])
        env.close()


class TestHeuristicAgents:
    def test_beam_equispaced(self):
        env = gym.make("oncosim/BeamSelection-v0")
        agent = BeamSelectionHeuristic(env, num_beams=5)
        actions = [agent.act(None) for _ in range(5)]
        # All actions should be valid angle indices
        assert all(0 <= a < 36 for a in actions)
        # Should be different angles (equispaced)
        assert len(set(actions)) == 5
        env.close()

    def test_dose_standard(self):
        agent = DoseFractionationHeuristic()
        action = agent.act(None)
        assert action[0] == pytest.approx(2.0)

    def test_adaptive_conservative(self):
        agent = AdaptiveRTHeuristic()
        action = agent.act(None)
        assert action[0] == 0  # No replan
        assert action[1] == 2  # Factor 1.0

    def test_evaluate_heuristic(self):
        env = gym.make("oncosim/DoseFractionation-v0")
        agent = DoseFractionationHeuristic()
        results = evaluate_heuristic(env, agent, n_episodes=5)
        assert "mean_reward" in results
        assert np.isfinite(results["mean_reward"])
        env.close()

    def test_heuristic_reset(self):
        env = gym.make("oncosim/BeamSelection-v0")
        agent = BeamSelectionHeuristic(env)
        agent.act(None)
        agent.reset()
        # After reset, should start from beginning
        first_action = agent.act(None)
        assert 0 <= first_action < 36
        env.close()

    def test_multiple_episodes(self):
        env = gym.make("oncosim/AdaptiveRT-v0")
        agent = AdaptiveRTHeuristic()
        for _ in range(3):
            obs, _ = env.reset(seed=42)
            agent.reset()
            done = False
            while not done:
                action = agent.act(obs)
                obs, _, terminated, truncated, _ = env.step(action)
                done = terminated or truncated
        env.close()
