"""PropofolControl-v0: BIS-guided propofol infusion rate control.

The agent controls propofol infusion rate to maintain the Bispectral Index
(BIS) within the target anesthetic window of 40-60, modeling a closed-loop
total intravenous anesthesia (TIVA) scenario.
"""

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.pd_models import HillPDModel


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


class PropofolControlEnv(gym.Env):
    """Gymnasium environment for propofol infusion control.

    Observation space (Box, 6 dimensions):
        - plasma_concentration: Propofol plasma level (mcg/mL)
        - effect_site_concentration: Propofol effect-site level (mcg/mL)
        - BIS: Current Bispectral Index (0-100)
        - BIS_error: Difference from target midpoint (BIS - 50)
        - infusion_rate_normalized: Current rate / max rate
        - time_fraction: Current step / max steps

    Action space (Box, 1 dimension):
        - infusion_rate: Propofol infusion in mcg/kg/min [0, 200]

    Reward:
        - +1.0 for BIS in [40, 60] (target window)
        - -0.5 for BIS in [30, 40) or (60, 70] (marginal)
        - -2.0 for BIS < 30 (dangerously deep) or BIS > 70 (too light)
        - -0.01 * |rate_change| for rate change penalty
    """

    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.pk_model = MarshPKModel(weight_kg=weight_kg, dt=dt)
        self.pd_model = HillPDModel()

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

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

        self.steps = 0
        self.in_target_steps = 0
        self.prev_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.pk_model = MarshPKModel(weight_kg=self.weight_kg, dt=self.dt)

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

        self.steps = 0
        self.in_target_steps = 0
        self.prev_rate = 0.0

        return self._get_obs(), {}

    def step(self, action):
        rate = float(np.clip(action[0], 0.0, 200.0))

        # Advance PK model
        pk_out = self.pk_model.step(rate)
        Ce = pk_out["effect_site_concentration"]

        # Compute BIS
        bis_clean = self.pd_model.predict_BIS(Ce)

        # Surgical stimulation: randomly increase BIS (lightening)
        if self._surgical_stimulation and self._rng is not None:
            if self._rng.random() < 0.05:  # 5% chance per step
                bis_clean = min(100.0, bis_clean + self._rng.uniform(5.0, 15.0))

        # Add 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
        if 40.0 <= bis <= 60.0:
            reward += 1.0
            self.in_target_steps += 1
        elif 30.0 <= bis < 40.0 or 60.0 < bis <= 70.0:
            reward -= 0.5
        else:
            reward -= 2.0

        # Rate change penalty
        rate_change = abs(rate - self.prev_rate) / 200.0
        reward -= 0.01 * rate_change
        self.prev_rate = rate

        self.steps += 1
        terminated = bis < 10.0  # Extremely deep anesthesia
        truncated = self.steps >= self.max_steps

        info = {
            "plasma_concentration": pk_out["plasma_concentration"],
            "effect_site_concentration": Ce,
            "BIS": bis,
            "BIS_clean": bis_clean,
            "in_target": 40.0 <= bis <= 60.0,
            "target_fraction": self.in_target_steps / max(1, self.steps),
        }

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

    def _get_obs(self) -> np.ndarray:
        Cp = min(self.pk_model.plasma_concentration, 20.0)
        Ce = min(self.pk_model.Ce, 20.0)
        bis = self.pd_model.predict_BIS(Ce)
        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([
            Cp,
            Ce,
            bis,
            bis - 50.0,
            self.prev_rate / 200.0,
            self.steps / max(1, self.max_steps),
        ], dtype=np.float32)
