"""Tests for baseline agents."""

import numpy as np
import pytest

import anesthesim  # noqa: F401
from anesthesim.agents.random_agent import evaluate_random
from anesthesim.agents.heuristic import (
    evaluate_heuristic,
    _propofol_tci_action,
    _remifentanil_tci_action,
    _combined_tci_action,
)


class TestRandomAgent:
    def test_propofol_returns_results(self):
        result = evaluate_random("anesthesim/PropofolControl-v0", n_episodes=3, seed=42)
        assert "mean_reward" in result
        assert "std_reward" in result
        assert len(result["episode_rewards"]) == 3

    def test_remifentanil_returns_results(self):
        result = evaluate_random("anesthesim/RemifentanilDosing-v0", n_episodes=3, seed=42)
        assert len(result["episode_rewards"]) == 3

    def test_combined_returns_results(self):
        result = evaluate_random("anesthesim/CombinedAnesthesia-v0", n_episodes=3, seed=42)
        assert len(result["episode_rewards"]) == 3

    def test_returns_finite_rewards(self):
        r = evaluate_random("anesthesim/PropofolControl-v0", n_episodes=3, seed=42)
        assert all(np.isfinite(x) for x in r["episode_rewards"])


class TestHeuristicAgent:
    def test_propofol_returns_results(self):
        result = evaluate_heuristic("anesthesim/PropofolControl-v0", n_episodes=3, seed=42)
        assert "mean_reward" in result
        assert len(result["episode_rewards"]) == 3

    def test_remifentanil_returns_results(self):
        result = evaluate_heuristic("anesthesim/RemifentanilDosing-v0", n_episodes=3, seed=42)
        assert len(result["episode_rewards"]) == 3

    def test_combined_returns_results(self):
        result = evaluate_heuristic("anesthesim/CombinedAnesthesia-v0", n_episodes=3, seed=42)
        assert len(result["episode_rewards"]) == 3

    def test_heuristic_beats_random(self):
        h = evaluate_heuristic("anesthesim/PropofolControl-v0", n_episodes=10, seed=42)
        r = evaluate_random("anesthesim/PropofolControl-v0", n_episodes=10, seed=42)
        assert h["mean_reward"] > r["mean_reward"]

    def test_unknown_env_raises(self):
        with pytest.raises((ValueError, Exception)):
            evaluate_heuristic("unknown/Env-v0", n_episodes=1)


class TestHeuristicActions:
    def test_propofol_tci_shape(self):
        obs = np.array([2.0, 1.5, 60.0, 10.0, 0.5, 0.3], dtype=np.float32)
        action = _propofol_tci_action(obs)
        assert action.shape == (1,)
        assert 0.0 <= action[0] <= 200.0

    def test_propofol_tci_increases_when_light(self):
        obs_light = np.array([1.0, 0.5, 80.0, 30.0, 0.3, 0.2], dtype=np.float32)
        obs_deep = np.array([5.0, 4.0, 30.0, -20.0, 0.6, 0.5], dtype=np.float32)
        rate_light = _propofol_tci_action(obs_light)[0]
        rate_deep = _propofol_tci_action(obs_deep)[0]
        assert rate_light > rate_deep

    def test_remifentanil_tci_shape(self):
        obs = np.array([2.0, 3.0, 4.0, -1.0, 0.3, 0.2], dtype=np.float32)
        action = _remifentanil_tci_action(obs)
        assert action.shape == (1,)
        assert 0.0 <= action[0] <= 0.5

    def test_combined_tci_shape(self):
        obs = np.array([2.0, 1.5, 3.0, 2.5, 55.0, 5.0, 0.4, 0.3], dtype=np.float32)
        action = _combined_tci_action(obs)
        assert action.shape == (2,)
        assert 0.0 <= action[0] <= 200.0
        assert 0.0 <= action[1] <= 0.5
