"""Urea clearing environment for hemodialysis Kt/V optimization.

The agent controls blood flow rate and dialysate flow rate to maximize
urea clearance (measured by Kt/V) during a 4-hour dialysis session.
"""

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

from nephrosim.models.patient import PatientProfile, create_patient
from nephrosim.models.urea_kinetics import UreaKineticsModel


class UreaClearingEnv(gym.Env):
    """Gymnasium environment for urea clearance optimization.

    Observation space (Dict -> flattened to Box):
        - urea_ecf: ECF urea concentration (mg/dL), range [0, 150]
        - urea_icf: ICF urea concentration (mg/dL), range [0, 150]
        - elapsed_time: Normalized session progress [0, 1]
        - blood_flow_rate: Current Qb (mL/min), range [200, 500]
        - dialysate_flow_rate: Current Qd (mL/min), range [300, 800]
        - volume_removed: Cumulative urea removed (mg), range [0, 5000]
        - kt_v_current: Current Kt/V, range [0, 3]

    Action space:
        Box(2) in [-1, 1] mapped to:
        - dialysate_flow_rate adjustment
        - blood_flow_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([0, 0, 0, 200, 300, 0, 0], dtype=np.float32),
            high=np.array([150, 150, 1, 500, 800, 5000, 3], dtype=np.float32),
        )
        self.action_space = spaces.Box(
            low=np.array([-1, -1], dtype=np.float32),
            high=np.array([1, 1], dtype=np.float32),
        )

        self.patient: PatientProfile | None = None
        self.model: UreaKineticsModel | None = None
        self.step_count = 0
        self.qb = 350.0
        self.qd = 500.0
        self.total_urea_removed = 0.0
        self.cumulative_clearance = 0.0
        self.clearance_steps = 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.model = UreaKineticsModel(
            ecf_volume=self.patient.ecf_volume * 1000,
            icf_volume=self.patient.icf_volume * 1000,
            koa=self.patient.koa_urea,
            intercompartment_kc=self.patient.intercompartment_transfer_coeff,
            urea_generation_rate=self.patient.urea_generation_rate,
        )
        self.model.reset(self.patient.pre_dialysis_bun, self.patient.pre_dialysis_bun)

        self.step_count = 0
        self.qb = 350.0
        self.qd = 500.0
        self.total_urea_removed = 0.0
        self.cumulative_clearance = 0.0
        self.clearance_steps = 0

        return self._get_obs(), {}

    def _get_obs(self) -> np.ndarray:
        assert self.model is not None
        noise = self.patient.noise_std if self.patient else 0.0
        rng = np.random.default_rng()

        c_ecf = self.model.state[0] + rng.normal(0, noise * 5)
        c_icf = self.model.state[1] + rng.normal(0, noise * 5)
        elapsed = self.step_count / self.max_steps

        avg_clearance = (
            self.cumulative_clearance / self.clearance_steps
            if self.clearance_steps > 0
            else 0.0
        )
        kt_v = self.model.compute_kt_v(avg_clearance, self.step_count * self.dt)

        return np.array([
            np.clip(c_ecf, 0, 150),
            np.clip(c_icf, 0, 150),
            elapsed,
            self.qb,
            self.qd,
            np.clip(self.total_urea_removed, 0, 5000),
            np.clip(kt_v, 0, 3),
        ], dtype=np.float32)

    def step(self, action):
        assert self.model is not None

        # Map actions to flow rate adjustments
        qd_delta = float(action[0]) * 50.0  # +/- 50 mL/min per step
        qb_delta = float(action[1]) * 25.0  # +/- 25 mL/min per step

        self.qd = float(np.clip(self.qd + qd_delta, 300, 800))
        self.qb = float(np.clip(self.qb + qb_delta, 200, 500))

        # Advance model
        result = self.model.step(self.qb, self.qd, self.dt)
        self.total_urea_removed += result["urea_removed_mg"]
        self.cumulative_clearance += result["clearance"]
        self.clearance_steps += 1

        self.step_count += 1

        # Compute reward
        avg_clearance = self.cumulative_clearance / self.clearance_steps
        kt_v = self.model.compute_kt_v(avg_clearance, self.step_count * self.dt)

        # Reward for Kt/V progress
        reward = result["clearance"] / 1000.0  # small positive for clearance

        # Penalty for extreme settings
        if self.qb > 450 or self.qb < 220:
            reward -= 0.5
        if self.qd > 750 or self.qd < 320:
            reward -= 0.5

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

        if terminated:
            if kt_v >= 1.2:
                reward += 10.0  # KDOQI target achieved
            if kt_v >= 1.4:
                reward += 5.0  # above minimum target
            if kt_v < 1.0:
                reward -= 5.0  # inadequate dialysis

        info = {
            "kt_v": kt_v,
            "c_ecf": result["c_ecf"],
            "c_icf": result["c_icf"],
            "clearance": result["clearance"],
            "total_urea_removed": self.total_urea_removed,
        }

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