"""Tests for random and heuristic baseline agents."""

import numpy as np
import pytest

import nephrosim  # noqa: F401
from nephrosim.agents.random_agent import evaluate_random
from nephrosim.agents.heuristic import (
    evaluate_heuristic,
    HEURISTIC_MAP,
    _heuristic_action_urea,
    _heuristic_action_uf,
    _heuristic_action_phosphate,
    _heuristic_action_full,
)


class TestRandomAgent:
    def test_evaluate_returns_dict(self):
        result = evaluate_random("nephrosim/UreaClearing-v0", n_episodes=3)
        assert "mean_reward" in result
        assert "std_reward" in result
        assert "episode_rewards" in result

    def test_correct_episode_count(self):
        result = evaluate_random("nephrosim/UreaClearing-v0", n_episodes=5)
        assert len(result["episode_rewards"]) == 5

    def test_mean_is_correct(self):
        result = evaluate_random("nephrosim/UreaClearing-v0", n_episodes=3, seed=42)
        expected = np.mean(result["episode_rewards"])
        assert abs(result["mean_reward"] - expected) < 0.01


class TestHeuristicActions:
    def test_urea_action_shape(self):
        obs = np.zeros(7, dtype=np.float32)
        action = _heuristic_action_urea(obs)
        assert action.shape == (2,)
        assert np.allclose(action, [0.0, 0.0])

    def test_uf_action_low_bp_reduces(self):
        obs = np.array([90, 60, 80, 500, 1.0, 3.0, 0.5, 5.0], dtype=np.float32)
        action = _heuristic_action_uf(obs)
        assert action[0] < 0  # should reduce UF

    def test_uf_action_normal_bp_maintains(self):
        obs = np.array([130, 80, 75, 500, 1.5, 3.0, 0.5, 5.0], dtype=np.float32)
        action = _heuristic_action_uf(obs)
        assert action[0] <= 0.11  # maintains or slight increase

    def test_phosphate_high_gives_more_binder(self):
        obs_high = np.array([7.0, 66.5, 1000, 0, 1, 1], dtype=np.float32)
        obs_low = np.array([3.0, 28.5, 1000, 0, 1, 1], dtype=np.float32)
        action_high = _heuristic_action_phosphate(obs_high)
        action_low = _heuristic_action_phosphate(obs_low)
        assert action_high[0] > action_low[0]

    def test_full_low_bp_reduces_uf(self):
        obs = np.zeros(11, dtype=np.float32)
        obs[3] = 95  # systolic BP
        action = _heuristic_action_full(obs)
        assert action[2] < 0  # UF adjustment negative


class TestHeuristicAgent:
    def test_evaluate_returns_dict(self):
        result = evaluate_heuristic("nephrosim/UreaClearing-v0", n_episodes=3)
        assert "mean_reward" in result

    def test_all_envs_have_heuristic(self):
        env_ids = [
            "nephrosim/UreaClearing-v0",
            "nephrosim/UltrafiltrationControl-v0",
            "nephrosim/PhosphateMgmt-v0",
            "nephrosim/FullDialysisSession-v0",
        ]
        for env_id in env_ids:
            assert env_id in HEURISTIC_MAP

    def test_unknown_env_raises(self):
        with pytest.raises(ValueError, match="No heuristic"):
            evaluate_heuristic("nephrosim/Nonexistent-v0")

    def test_heuristic_beats_random_on_urea(self):
        """Heuristic should outperform random on urea clearing (stable flows)."""
        random_r = evaluate_random(
            "nephrosim/UreaClearing-v0", n_episodes=20, seed=42
        )
        heuristic_r = evaluate_heuristic(
            "nephrosim/UreaClearing-v0", n_episodes=20, seed=42
        )
        # Heuristic maintains stable flows, avoids extreme-setting penalties
        assert heuristic_r["mean_reward"] > random_r["mean_reward"]
