"""Integration tests for AnestheSim."""

import gymnasium as gym
import numpy as np
import pytest

import anesthesim  # noqa: F401
from anesthesim.models.marsh_pk import MarshPKModel
from anesthesim.models.minto_pk import MintoPKModel
from anesthesim.models.pd_models import HillPDModel, GrecoInteractionModel
from anesthesim.models.patient import generate_patient_cohort


class TestFullPipeline:
    """End-to-end tests running complete episodes."""

    def test_propofol_full_episode(self):
        env = gym.make("anesthesim/PropofolControl-v0", difficulty="easy")
        obs, _ = env.reset(seed=42)
        total_reward = 0.0
        for _ in range(200):
            action = np.array([100.0], dtype=np.float32)
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            if terminated or truncated:
                break
        assert total_reward != 0.0
        env.close()

    def test_remifentanil_full_episode(self):
        env = gym.make("anesthesim/RemifentanilDosing-v0", difficulty="easy")
        obs, _ = env.reset(seed=42)
        total_reward = 0.0
        for _ in range(200):
            action = np.array([0.15], dtype=np.float32)
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            if terminated or truncated:
                break
        assert total_reward != 0.0
        env.close()

    def test_combined_full_episode(self):
        env = gym.make("anesthesim/CombinedAnesthesia-v0", difficulty="easy")
        obs, _ = env.reset(seed=42)
        total_reward = 0.0
        for _ in range(200):
            action = np.array([100.0, 0.15], dtype=np.float32)
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            if terminated or truncated:
                break
        assert total_reward != 0.0
        env.close()


class TestModelEnvironmentIntegration:
    """Test that models integrate correctly with environments."""

    def test_marsh_pk_in_env_context(self):
        pk = MarshPKModel(weight_kg=70.0)
        pd = HillPDModel()
        for _ in range(100):
            pk.step(120.0)
        bis = pd.predict_BIS(pk.Ce)
        assert 0.0 <= bis <= 97.0

    def test_minto_pk_in_env_context(self):
        pk = MintoPKModel()
        for _ in range(100):
            pk.step(0.2)
        assert pk.Ce > 0.0

    def test_combined_models_interaction(self):
        prop_pk = MarshPKModel()
        remi_pk = MintoPKModel()
        pd = GrecoInteractionModel()
        for _ in range(100):
            prop_pk.step(100.0)
            remi_pk.step(0.15)
        bis = pd.predict_BIS(prop_pk.Ce, remi_pk.Ce)
        assert 0.0 <= bis <= 97.0


class TestPatientVariability:
    """Test patient variability in environments."""

    def test_different_seeds_different_episodes(self):
        env = gym.make("anesthesim/PropofolControl-v0", difficulty="medium")
        rewards = []
        for seed in [1, 2, 3]:
            env.reset(seed=seed)
            total = 0.0
            for _ in range(50):
                _, r, t, tr, _ = env.step(np.array([100.0]))
                total += r
                if t or tr:
                    break
            rewards.append(total)
        # Not all should be identical
        assert len(set(f"{r:.2f}" for r in rewards)) > 1
        env.close()

    def test_patient_cohort_in_env(self):
        patients = generate_patient_cohort(n_patients=6)
        for p in patients:
            env = gym.make("anesthesim/PropofolControl-v0", weight_kg=p.weight_kg)
            obs, _ = env.reset(seed=42)
            assert obs.shape == (6,)
            env.close()


class TestReproducibility:
    """Test deterministic behavior with same seed."""

    def test_propofol_deterministic(self):
        env = gym.make("anesthesim/PropofolControl-v0", difficulty="easy")
        results = []
        for _ in range(2):
            env.reset(seed=42)
            ep_reward = 0.0
            for _ in range(50):
                _, r, t, tr, _ = env.step(np.array([100.0]))
                ep_reward += r
                if t or tr:
                    break
            results.append(ep_reward)
        assert results[0] == pytest.approx(results[1])
        env.close()

    def test_combined_deterministic(self):
        env = gym.make("anesthesim/CombinedAnesthesia-v0", difficulty="easy")
        results = []
        for _ in range(2):
            env.reset(seed=42)
            ep_reward = 0.0
            for _ in range(50):
                _, r, t, tr, _ = env.step(np.array([100.0, 0.15]))
                ep_reward += r
                if t or tr:
                    break
            results.append(ep_reward)
        assert results[0] == pytest.approx(results[1])
        env.close()
