"""SepsisManagement-v0: Combined multi-intervention sepsis management.

The agent simultaneously manages fluid administration, vasopressor dosing,
and antibiotic timing for a septic patient. This is the most challenging
environment, requiring balancing multiple competing interventions.

Observation space (10D continuous):
    - MAP (mmHg)
    - Lactate (mmol/L)
    - Urine output (mL/kg/h)
    - Tissue damage (0-1)
    - Intravascular volume excess (L)
    - Current vasopressor dose (mcg/kg/min)
    - Bacteria load (normalized)
    - SOFA score (0-12)
    - Hours since admission
    - Antibiotic active (0 or 1)

Action space (3D continuous):
    - Fluid bolus: 0-1000 mL
    - Vasopressor dose change: -0.1 to +0.1 mcg/kg/min
    - Administer antibiotic: 0 or 1 (thresholded at 0.5)

Reward:
    - Composite of MAP maintenance, lactate clearance, fluid balance, organ function
    - Large bonus for successful stabilization (MAP > 65, lactate < 2, SOFA decreasing)
    - Large penalty for death
"""

from __future__ import annotations

import gymnasium as gym
import numpy as np
from gymnasium import spaces
from numpy.typing import NDArray

from sepsisim.models.cardiovascular import CardiovascularModel
from sepsisim.models.inflammation import InflammationModel
from sepsisim.models.lactate import LactateModel
from sepsisim.models.sofa import compute_sofa_score


class SepsisManagementEnv(gym.Env):
    """Gymnasium environment for comprehensive sepsis management."""

    metadata = {"render_modes": []}

    def __init__(
        self,
        severity: str = "medium",
        dt: float = 1.0,
    ) -> None:
        super().__init__()

        self.dt = dt
        self.severity = severity

        severity_map = {"easy": 0.25, "medium": 0.45, "hard": 0.65}
        self._init_bacteria = severity_map.get(severity, 0.45)

        self.observation_space = spaces.Box(
            low=np.array(
                [20.0, 0.1, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                dtype=np.float32,
            ),
            high=np.array(
                [150.0, 30.0, 2.0, 1.0, 10.0, 2.0, 1.0, 12.0, 72.0, 1.0],
                dtype=np.float32,
            ),
        )

        self.action_space = spaces.Box(
            low=np.array([0.0, -0.1, 0.0], dtype=np.float32),
            high=np.array([1000.0, 0.1, 1.0], dtype=np.float32),
        )

        self._inflammation: InflammationModel | None = None
        self._cardio: CardiovascularModel | None = None
        self._lactate_model: LactateModel | None = None
        self._rng: np.random.Generator | None = None

        self._inf_state: NDArray[np.float64] | None = None
        self._lactate: float = 0.0
        self._volume: float = 0.0
        self._vaso_dose: float = 0.0
        self._map: float = 0.0
        self._uo: float = 0.0
        self._abx_active: bool = False
        self._abx_given_step: int = -1
        self._step_count: int = 0
        self._total_fluid: float = 0.0
        self._prev_sofa: int = 12
        self._prev_lactate: float = 10.0

    def reset(
        self, *, seed: int | None = None, options: dict | None = None
    ) -> tuple[NDArray[np.float32], dict]:
        super().reset(seed=seed)
        self._rng = np.random.default_rng(seed)

        self._inflammation = InflammationModel(antibiotic_efficacy=0.0)
        self._cardio = CardiovascularModel()
        self._lactate_model = LactateModel()

        self._inf_state = self._inflammation.reset(
            bacteria_load=self._init_bacteria, rng=self._rng
        )
        self._volume = 0.0
        self._vaso_dose = 0.0
        self._abx_active = False
        self._abx_given_step = -1
        self._total_fluid = 0.0
        self._step_count = 0

        tissue_damage = float(self._inf_state[3])
        self._map = self._cardio.compute_map(
            tissue_damage, self._volume, self._vaso_dose, self._rng
        )
        self._lactate = self._lactate_model.reset(
            severity=self._init_bacteria, rng=self._rng
        )
        self._uo = self._cardio.compute_urine_output(self._map, self._rng)

        sofa = compute_sofa_score(self._map, self._vaso_dose, self._uo, self._lactate)
        self._prev_sofa = sofa
        self._prev_lactate = self._lactate

        return self._get_obs(), self._get_info()

    def step(
        self, action: NDArray[np.float32]
    ) -> tuple[NDArray[np.float32], float, bool, bool, dict]:
        assert self._inflammation is not None
        assert self._cardio is not None
        assert self._lactate_model is not None
        assert self._inf_state is not None

        fluid_ml = float(np.clip(action[0], 0.0, 1000.0))
        dose_change = float(np.clip(action[1], -0.1, 0.1))
        give_abx = float(action[2]) > 0.5

        self._vaso_dose = float(np.clip(self._vaso_dose + dose_change, 0.0, 2.0))
        self._total_fluid += fluid_ml

        if give_abx and not self._abx_active:
            self._abx_active = True
            self._abx_given_step = self._step_count
            self._inflammation.antibiotic_efficacy = 0.5

        self._inf_state = self._inflammation.step(self._inf_state, self.dt)
        tissue_damage = float(self._inf_state[3])

        self._volume = self._cardio.update_volume(
            self._volume, fluid_ml, self.dt
        )
        self._map = self._cardio.compute_map(
            tissue_damage, self._volume, self._vaso_dose, self._rng
        )
        self._lactate = self._lactate_model.step(
            self._lactate, self._map, tissue_damage, self.dt
        )
        self._uo = self._cardio.compute_urine_output(self._map, self._rng)

        sofa = compute_sofa_score(
            self._map, self._vaso_dose, self._uo, self._lactate
        )

        reward = self._compute_reward(tissue_damage, sofa)

        self._prev_sofa = sofa
        self._prev_lactate = self._lactate
        self._step_count += 1

        terminated = tissue_damage >= 0.9
        truncated = False

        if terminated:
            reward -= 50.0

        stabilized = (
            self._map >= 65.0
            and self._lactate < 2.0
            and sofa <= 3
            and self._step_count >= 6
        )
        if stabilized:
            reward += 10.0

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

    def _compute_reward(self, tissue_damage: float, sofa: int) -> float:
        reward = 0.0

        if 65.0 <= self._map <= 90.0:
            reward += 1.5
        elif self._map < 55.0:
            reward -= 3.0
        elif self._map < 65.0:
            reward -= 1.0 * (65.0 - self._map) / 10.0

        if self._lactate < 2.0:
            reward += 1.0
        elif self._lactate > 4.0:
            reward -= 0.5 * (self._lactate - 4.0)

        lactate_delta = self._prev_lactate - self._lactate
        reward += 0.5 * np.clip(lactate_delta, -2.0, 2.0)

        sofa_delta = self._prev_sofa - sofa
        reward += 0.3 * sofa_delta

        if self._volume > 4.0:
            reward -= 0.5 * (self._volume - 4.0)
        if self._vaso_dose > 0.5:
            reward -= 0.3 * (self._vaso_dose - 0.5)

        if self._uo >= 0.5:
            reward += 0.5

        if not self._abx_active and self._step_count >= 3:
            reward -= 1.0

        reward -= 0.1 * tissue_damage

        return float(reward)

    def _get_obs(self) -> NDArray[np.float32]:
        assert self._inf_state is not None
        sofa = compute_sofa_score(
            self._map, self._vaso_dose, self._uo, self._lactate
        )
        return np.array([
            self._map,
            self._lactate,
            self._uo,
            float(self._inf_state[3]),
            self._volume,
            self._vaso_dose,
            float(self._inf_state[0]),
            float(sofa),
            float(self._step_count),
            1.0 if self._abx_active else 0.0,
        ], dtype=np.float32)

    def _get_info(self) -> dict:
        assert self._inf_state is not None
        return {
            "map_mmhg": self._map,
            "lactate": self._lactate,
            "urine_output": self._uo,
            "tissue_damage": float(self._inf_state[3]),
            "bacteria": float(self._inf_state[0]),
            "vasopressor_dose": self._vaso_dose,
            "total_fluid_ml": self._total_fluid,
            "intravascular_volume": self._volume,
            "antibiotic_active": self._abx_active,
            "step": self._step_count,
        }
