"""Ultrafiltration control environment for intradialytic hypotension prevention.

The agent controls the ultrafiltration rate to remove excess fluid while
preventing dangerous drops in blood pressure (intradialytic hypotension).
"""

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

from nephrosim.models.patient import PatientProfile, create_patient
from nephrosim.models.cardiovascular import CardiovascularModel


class UltrafiltrationControlEnv(gym.Env):
    """Gymnasium environment for ultrafiltration rate control.

    Observation space (Box):
        - systolic_bp: Systolic blood pressure (mmHg), range [50, 200]
        - diastolic_bp: Diastolic blood pressure (mmHg), range [30, 130]
        - heart_rate: Heart rate (bpm), range [40, 160]
        - uf_rate: Current UF rate (mL/hr), range [0, 2000]
        - total_removed: Total fluid removed (L), range [0, 6]
        - target_volume: Target removal volume (L), range [1, 6]
        - elapsed_time: Normalized session progress [0, 1]
        - plasma_refill_rate: Estimated refill rate (mL/min), range [0, 50]

    Action space:
        Box(1) in [-1, 1] mapped to UF rate adjustment.
    """

    metadata = {"render_modes": []}

    def __init__(
        self,
        difficulty: str = "easy",
        session_duration: int = 240,
        dt: float = 1.0,
    ):
        super().__init__()
        self.difficulty = difficulty
        self.session_duration = session_duration
        self.dt = dt
        self.max_steps = int(session_duration / dt)

        self.observation_space = spaces.Box(
            low=np.array([50, 30, 40, 0, 0, 1, 0, 0], dtype=np.float32),
            high=np.array([200, 130, 160, 2000, 6, 6, 1, 50], dtype=np.float32),
        )
        self.action_space = spaces.Box(
            low=np.array([-1], dtype=np.float32),
            high=np.array([1], dtype=np.float32),
        )

        self.patient: PatientProfile | None = None
        self.cv_model: CardiovascularModel | None = None
        self.step_count = 0
        self.uf_rate = 0.0
        self.hypotension_count = 0

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        rng = np.random.default_rng(seed)

        self.patient = create_patient(self.difficulty, rng)
        self.cv_model = CardiovascularModel(
            baseline_sbp=self.patient.baseline_systolic_bp,
            baseline_dbp=self.patient.baseline_diastolic_bp,
            baseline_hr=self.patient.baseline_heart_rate,
            blood_volume=self.patient.body_weight * 70,  # ~70 mL/kg
            plasma_refill_coeff=self.patient.plasma_refill_coefficient,
            is_diabetic=self.patient.is_diabetic,
        )
        self.cv_model.reset()

        self.step_count = 0
        # Start with linear rate to achieve target in session
        target_ml = self.patient.target_fluid_removal * 1000
        self.uf_rate = float(target_ml / (self.session_duration / 60))
        self.hypotension_count = 0

        return self._get_obs(), {}

    def _get_obs(self) -> np.ndarray:
        assert self.cv_model is not None and self.patient is not None
        state = self.cv_model.get_state()
        noise = self.patient.noise_std

        rng = np.random.default_rng()
        sbp = state["systolic_bp"] + rng.normal(0, noise * 5)
        dbp = state["diastolic_bp"] + rng.normal(0, noise * 3)
        hr = state["heart_rate"] + rng.normal(0, noise * 3)

        return np.array([
            np.clip(sbp, 50, 200),
            np.clip(dbp, 30, 130),
            np.clip(hr, 40, 160),
            self.uf_rate,
            np.clip(state["total_uf_volume"] / 1000, 0, 6),
            self.patient.target_fluid_removal,
            self.step_count / self.max_steps,
            np.clip(self.cv_model.plasma_refill_rate(), 0, 50),
        ], dtype=np.float32)

    def step(self, action):
        assert self.cv_model is not None and self.patient is not None

        # Map action to UF rate adjustment
        uf_delta = float(action[0]) * 200.0  # +/- 200 mL/hr per step
        self.uf_rate = float(np.clip(self.uf_rate + uf_delta, 0, self.patient.max_uf_rate))

        # Advance cardiovascular model
        state = self.cv_model.step(self.uf_rate, self.dt)

        self.step_count += 1

        # Compute reward
        reward = 0.0

        # Progress toward target
        removed_L = state["total_uf_volume"] / 1000
        progress = removed_L / self.patient.target_fluid_removal
        reward += 0.5 * (self.uf_rate / 1000.0)  # small positive for removing fluid

        # Hypotension penalty
        if self.cv_model.is_hypotensive:
            reward -= 5.0
            self.hypotension_count += 1

        # Mild penalty for very high UF rate (risk factor)
        if self.uf_rate > self.patient.max_uf_rate * 0.9:
            reward -= 0.3

        # Terminal conditions
        terminated = self.step_count >= self.max_steps
        truncated = False

        if terminated:
            completion = min(1.0, removed_L / self.patient.target_fluid_removal)
            if completion >= 0.9 and self.hypotension_count == 0:
                reward += 10.0
            elif completion >= 0.9:
                reward += 3.0
            elif completion < 0.5:
                reward -= 5.0

        info = {
            "systolic_bp": state["systolic_bp"],
            "total_removed_L": removed_L,
            "target_L": self.patient.target_fluid_removal,
            "completion_fraction": progress,
            "hypotension_events": self.hypotension_count,
            "is_hypotensive": self.cv_model.is_hypotensive,
        }

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