"""CombinedAnesthesia-v0: Multi-drug propofol + remifentanil anesthesia control.

The agent simultaneously controls propofol and remifentanil infusion rates to
maintain BIS in the target range, exploiting the synergistic drug interaction
modeled by the Greco response surface.
"""

from __future__ import annotations

import gymnasium as gym
import numpy as np
from gymnasium import spaces

from anesthesim.models.marsh_pk import MarshPKModel
from anesthesim.models.minto_pk import MintoPKModel
from anesthesim.models.pd_models import GrecoInteractionModel


DIFFICULTY_TIERS = {
    "easy": {
        "pk_variability": 0.0, "surgical_stimulation": False,
        "bis_noise": 0.0,
    },
    "medium": {
        "pk_variability": 0.15, "surgical_stimulation": False,
        "bis_noise": 2.0,
    },
    "hard": {
        "pk_variability": 0.3, "surgical_stimulation": True,
        "bis_noise": 3.0,
    },
}


class CombinedAnesthesiaEnv(gym.Env):
    """Gymnasium environment for combined propofol-remifentanil anesthesia.

    Observation space (Box, 8 dimensions):
        - propofol_plasma: Propofol plasma concentration (mcg/mL)
        - propofol_effect: Propofol effect-site concentration (mcg/mL)
        - remifentanil_plasma: Remifentanil plasma concentration (ng/mL)
        - remifentanil_effect: Remifentanil effect-site concentration (ng/mL)
        - BIS: Current Bispectral Index
        - BIS_error: BIS - 50 (signed error from target midpoint)
        - propofol_rate_norm: Propofol rate / max rate
        - remifentanil_rate_norm: Remifentanil rate / max rate

    Action space (Box, 2 dimensions):
        - propofol_rate: Propofol infusion mcg/kg/min [0, 200]
        - remifentanil_rate: Remifentanil infusion mcg/kg/min [0, 0.5]

    Reward:
        - +1.0 for BIS in [40, 60]
        - -0.5 for BIS in [30, 40) or (60, 70]
        - -2.0 for BIS < 30 or BIS > 70
        - +0.2 synergy bonus for using lower total drug while in target
        - -0.01 * rate_change for stability
    """

    metadata = {"render_modes": []}

    def __init__(
        self,
        difficulty: str = "medium",
        max_steps: int = 200,
        dt: float = 0.5,
        weight_kg: float = 70.0,
    ) -> None:
        super().__init__()

        tier = DIFFICULTY_TIERS.get(difficulty, DIFFICULTY_TIERS["medium"])
        self.difficulty = difficulty
        self.max_steps = max_steps
        self.dt = dt
        self.weight_kg = weight_kg
        self._pk_variability = tier["pk_variability"]
        self._surgical_stimulation = tier["surgical_stimulation"]
        self._bis_noise = tier["bis_noise"]

        self.propofol_pk = MarshPKModel(weight_kg=weight_kg, dt=dt)
        self.remifentanil_pk = MintoPKModel(weight_kg=weight_kg, dt=dt)
        self.pd_model = GrecoInteractionModel()

        self.observation_space = spaces.Box(
            low=np.array([0.0, 0.0, 0.0, 0.0, 0.0, -100.0, 0.0, 0.0], dtype=np.float32),
            high=np.array([20.0, 20.0, 15.0, 15.0, 100.0, 100.0, 1.0, 1.0], dtype=np.float32),
        )

        self.action_space = spaces.Box(
            low=np.array([0.0, 0.0], dtype=np.float32),
            high=np.array([200.0, 0.5], dtype=np.float32),
        )

        self.steps = 0
        self.in_target_steps = 0
        self.prev_prop_rate = 0.0
        self.prev_remi_rate = 0.0
        self._rng: np.random.Generator | None = None

    def reset(self, *, seed: int | None = None, options: dict | None = None):
        super().reset(seed=seed)
        self._rng = np.random.default_rng(seed)

        self.propofol_pk = MarshPKModel(weight_kg=self.weight_kg, dt=self.dt)
        self.remifentanil_pk = MintoPKModel(weight_kg=self.weight_kg, dt=self.dt)

        if self._pk_variability > 0:
            self.propofol_pk.apply_variability(self._rng, self._pk_variability)
            self.remifentanil_pk.apply_variability(self._rng, self._pk_variability)

        self.steps = 0
        self.in_target_steps = 0
        self.prev_prop_rate = 0.0
        self.prev_remi_rate = 0.0

        return self._get_obs(), {}

    def step(self, action):
        prop_rate = float(np.clip(action[0], 0.0, 200.0))
        remi_rate = float(np.clip(action[1], 0.0, 0.5))

        # Advance PK models
        prop_out = self.propofol_pk.step(prop_rate)
        remi_out = self.remifentanil_pk.step(remi_rate)

        Ce_prop = prop_out["effect_site_concentration"]
        Ce_remi = remi_out["effect_site_concentration"]

        # Compute BIS via interaction model
        bis_clean = self.pd_model.predict_BIS(Ce_prop, Ce_remi)

        # Surgical stimulation
        if self._surgical_stimulation and self._rng is not None:
            if self._rng.random() < 0.05:
                bis_clean = min(100.0, bis_clean + self._rng.uniform(5.0, 15.0))

        # Sensor noise
        if self._bis_noise > 0 and self._rng is not None:
            bis = self.pd_model.add_noise(bis_clean, self._rng, self._bis_noise)
        else:
            bis = bis_clean

        # Compute reward
        reward = 0.0
        in_target = 40.0 <= bis <= 60.0

        if in_target:
            reward += 1.0
            self.in_target_steps += 1
            # Synergy bonus: reward efficient drug use
            total_drug_frac = (prop_rate / 200.0 + remi_rate / 0.5) / 2.0
            if total_drug_frac < 0.5:
                reward += 0.2
        elif 30.0 <= bis < 40.0 or 60.0 < bis <= 70.0:
            reward -= 0.5
        else:
            reward -= 2.0

        # Rate change penalty
        prop_change = abs(prop_rate - self.prev_prop_rate) / 200.0
        remi_change = abs(remi_rate - self.prev_remi_rate) / 0.5
        reward -= 0.01 * (prop_change + remi_change)
        self.prev_prop_rate = prop_rate
        self.prev_remi_rate = remi_rate

        self.steps += 1
        terminated = bis < 10.0
        truncated = self.steps >= self.max_steps

        info = {
            "propofol_Ce": Ce_prop,
            "remifentanil_Ce": Ce_remi,
            "BIS": bis,
            "BIS_clean": bis_clean,
            "in_target": in_target,
            "target_fraction": self.in_target_steps / max(1, self.steps),
        }

        return self._get_obs(), reward, terminated, truncated, info

    def _get_obs(self) -> np.ndarray:
        Ce_prop = min(self.propofol_pk.Ce, 20.0)
        Ce_remi = min(self.remifentanil_pk.Ce, 15.0)
        bis = self.pd_model.predict_BIS(Ce_prop, Ce_remi)
        if self._bis_noise > 0 and self._rng is not None:
            bis = self.pd_model.add_noise(bis, self._rng, self._bis_noise)
        return np.array([
            min(self.propofol_pk.plasma_concentration, 20.0),
            Ce_prop,
            min(self.remifentanil_pk.plasma_concentration, 15.0),
            Ce_remi,
            bis,
            bis - 50.0,
            self.prev_prop_rate / 200.0,
            self.prev_remi_rate / 0.5,
        ], dtype=np.float32)
