"""Heuristic clinical baseline agents for AnestheSim environments.

Implements simplified Target-Controlled Infusion (TCI) protocols as
clinical baselines for comparison with RL agents.

References:
    Absalom, A. R., & Struys, M. M. (2007). Overview of target
    controlled infusions and total intravenous anaesthesia.
    Academia Press.

    Moore, B. L., Pyeatt, L. D., Kulkarni, V., et al. (2011).
    Reinforcement learning for closed-loop propofol anesthesia.
    Artificial Intelligence in Medicine, 55(3), 159-168.
"""

from __future__ import annotations

import gymnasium as gym
import numpy as np


def _propofol_tci_action(obs: np.ndarray) -> np.ndarray:
    """Simplified proportional TCI controller for propofol.

    Targets BIS = 50 using proportional control on BIS error.
    """
    bis = obs[2]  # BIS value
    error = bis - 50.0  # positive = too light, negative = too deep

    # Proportional controller: increase rate if BIS too high (light)
    base_rate = 100.0  # mcg/kg/min maintenance
    Kp = 5.0
    rate = base_rate + Kp * error
    return np.array([np.clip(rate, 0.0, 200.0)], dtype=np.float32)


def _remifentanil_tci_action(obs: np.ndarray) -> np.ndarray:
    """Simplified proportional controller for remifentanil.

    Targets effect-site concentration matching the target.
    """
    Ce = obs[1]  # effect-site concentration
    target = obs[2]  # target concentration
    error = target - Ce  # positive = need more drug

    base_rate = 0.15  # mcg/kg/min
    Kp = 0.05
    rate = base_rate + Kp * error
    return np.array([np.clip(rate, 0.0, 0.5)], dtype=np.float32)


def _combined_tci_action(obs: np.ndarray) -> np.ndarray:
    """Combined TCI for propofol + remifentanil."""
    bis = obs[4]  # BIS
    error = bis - 50.0

    # Propofol: proportional on BIS error
    prop_base = 80.0
    prop_rate = prop_base + 4.0 * error
    prop_rate = np.clip(prop_rate, 0.0, 200.0)

    # Remifentanil: fixed moderate rate, adjust with BIS
    remi_base = 0.15
    remi_rate = remi_base + 0.01 * max(0, error)
    remi_rate = np.clip(remi_rate, 0.0, 0.5)

    return np.array([prop_rate, remi_rate], dtype=np.float32)


HEURISTIC_MAP = {
    "anesthesim/PropofolControl-v0": _propofol_tci_action,
    "anesthesim/RemifentanilDosing-v0": _remifentanil_tci_action,
    "anesthesim/CombinedAnesthesia-v0": _combined_tci_action,
}


def evaluate_heuristic(
    env_id: str,
    n_episodes: int = 50,
    seed: int = 42,
) -> dict:
    """Evaluate the heuristic agent on the given environment.

    Args:
        env_id: Gymnasium environment ID.
        n_episodes: Number of evaluation episodes.
        seed: Random seed.

    Returns:
        Dict with mean_reward, std_reward, episode_rewards.
    """
    import anesthesim  # noqa: F401

    env = gym.make(env_id)
    rng = np.random.default_rng(seed)
    action_fn = HEURISTIC_MAP.get(env_id)
    if action_fn is None:
        raise ValueError(f"No heuristic defined for {env_id}")

    rewards_list: list[float] = []

    for ep in range(n_episodes):
        obs, _ = env.reset(seed=int(rng.integers(0, 2**31)))
        total_reward = 0.0
        done = False
        while not done:
            action = action_fn(obs)
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += reward
            done = terminated or truncated
        rewards_list.append(total_reward)

    env.close()
    return {
        "mean_reward": float(np.mean(rewards_list)),
        "std_reward": float(np.std(rewards_list)),
        "episode_rewards": rewards_list,
    }
