"""RemifentanilDosing-v0: Effect-site concentration targeting for remifentanil.

The agent controls remifentanil infusion rate to track a target effect-site
concentration, modeling intraoperative analgesia management.
"""

from __future__ import annotations

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

from anesthesim.models.minto_pk import MintoPKModel


DIFFICULTY_TIERS = {
    "easy": {
        "pk_variability": 0.0, "target_changes": False,
        "noise_std": 0.0, "fixed_target": 4.0,
    },
    "medium": {
        "pk_variability": 0.15, "target_changes": True,
        "noise_std": 0.3, "fixed_target": 4.0,
    },
    "hard": {
        "pk_variability": 0.3, "target_changes": True,
        "noise_std": 0.5, "fixed_target": 4.0,
    },
}


class RemifentanilDosingEnv(gym.Env):
    """Gymnasium environment for remifentanil infusion control.

    Observation space (Box, 6 dimensions):
        - plasma_concentration: Remifentanil plasma level (ng/mL)
        - effect_site_concentration: Remifentanil effect-site level (ng/mL)
        - target_concentration: Current target Ce (ng/mL)
        - concentration_error: Ce - target (ng/mL)
        - infusion_rate_normalized: Current rate / max rate
        - time_fraction: Current step / max steps

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

    Reward:
        - +1.0 for |Ce - target| < 0.5 ng/mL
        - Scaled penalty proportional to error otherwise
        - -3.0 for Ce > 8 ng/mL (respiratory depression risk)
    """

    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._target_changes = tier["target_changes"]
        self._noise_std = tier["noise_std"]
        self._base_target = tier["fixed_target"]

        self.pk_model = MintoPKModel(weight_kg=weight_kg, dt=dt)

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

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

        self.steps = 0
        self.on_target_steps = 0
        self.target_Ce = self._base_target
        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 = MintoPKModel(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.on_target_steps = 0
        self.target_Ce = self._base_target
        self.prev_rate = 0.0

        return self._get_obs(), {}

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

        # Target changes at medium/hard difficulty
        if self._target_changes and self._rng is not None:
            if self.steps > 0 and self.steps % 50 == 0:
                self.target_Ce = self._rng.uniform(2.0, 6.0)

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

        # Add measurement noise
        if self._noise_std > 0 and self._rng is not None:
            Ce_observed = max(0.0, Ce + self._rng.normal(0.0, self._noise_std))
        else:
            Ce_observed = Ce

        # Compute reward
        error = abs(Ce - self.target_Ce)
        if error < 0.5:
            reward = 1.0
            self.on_target_steps += 1
        elif error < 1.5:
            reward = 0.5 - 0.5 * (error - 0.5)
        else:
            reward = -0.5 * min(error, 5.0)

        # Respiratory depression penalty
        if Ce > 8.0:
            reward -= 3.0

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

        self.steps += 1
        terminated = Ce > 12.0  # Severe respiratory depression
        truncated = self.steps >= self.max_steps

        info = {
            "plasma_concentration": pk_out["plasma_concentration"],
            "effect_site_concentration": Ce,
            "target_concentration": self.target_Ce,
            "error": error,
            "on_target": error < 0.5,
            "target_fraction": self.on_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, 15.0)
        Ce = min(self.pk_model.Ce, 15.0)
        return np.array([
            Cp,
            Ce,
            self.target_Ce,
            Ce - self.target_Ce,
            self.prev_rate / 0.5,
            self.steps / max(1, self.max_steps),
        ], dtype=np.float32)
