"""Full dialysis session environment combining urea clearance and ultrafiltration.

The agent jointly optimizes blood flow rate, dialysate flow rate,
ultrafiltration rate, and dialysate temperature for a complete session.
Multi-objective: maximize Kt/V while preventing 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.urea_kinetics import UreaKineticsModel
from nephrosim.models.cardiovascular import CardiovascularModel


class FullDialysisSessionEnv(gym.Env):
    """Gymnasium environment for complete dialysis session management.

    Observation space (Box, 11 dims):
        - urea_ecf, urea_icf, kt_v_current (from urea model)
        - systolic_bp, diastolic_bp, heart_rate (from CV model)
        - uf_rate, total_removed_L, target_volume_L
        - blood_flow_rate, elapsed_time

    Action space:
        Box(3) in [-1, 1]:
        - blood_flow_rate adjustment
        - dialysate_flow_rate adjustment
        - ultrafiltration_rate adjustment

    Difficulty: always 'hard' (diabetic elderly patient with variable response).
    """

    metadata = {"render_modes": []}

    def __init__(self, session_duration: int = 240, dt: float = 1.0):
        super().__init__()
        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, 50, 30, 40, 0, 0, 1, 200, 0], dtype=np.float32),
            high=np.array([150, 150, 3, 200, 130, 160, 2000, 6, 6, 500, 1], dtype=np.float32),
        )
        self.action_space = spaces.Box(
            low=np.array([-1, -1, -1], dtype=np.float32),
            high=np.array([1, 1, 1], dtype=np.float32),
        )

        self.patient: PatientProfile | None = None
        self.urea_model: UreaKineticsModel | None = None
        self.cv_model: CardiovascularModel | None = None
        self.step_count = 0
        self.qb = 350.0
        self.qd = 500.0
        self.uf_rate = 0.0
        self.cumulative_clearance = 0.0
        self.clearance_steps = 0
        self.hypotension_count = 0

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

        # Always hard difficulty
        self.patient = create_patient("hard", rng)

        self.urea_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.urea_model.reset(self.patient.pre_dialysis_bun, self.patient.pre_dialysis_bun)

        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,
            plasma_refill_coeff=self.patient.plasma_refill_coefficient,
            is_diabetic=self.patient.is_diabetic,
        )
        self.cv_model.reset()

        self.step_count = 0
        self.qb = 350.0
        self.qd = 500.0
        target_ml = self.patient.target_fluid_removal * 1000
        self.uf_rate = float(target_ml / (self.session_duration / 60))
        self.cumulative_clearance = 0.0
        self.clearance_steps = 0
        self.hypotension_count = 0

        return self._get_obs(), {}

    def _get_obs(self) -> np.ndarray:
        assert self.urea_model is not None
        assert self.cv_model is not None
        assert self.patient is not None

        cv_state = self.cv_model.get_state()
        avg_k = (
            self.cumulative_clearance / self.clearance_steps
            if self.clearance_steps > 0 else 0.0
        )
        kt_v = self.urea_model.compute_kt_v(avg_k, self.step_count * self.dt)

        return np.array([
            np.clip(self.urea_model.state[0], 0, 150),
            np.clip(self.urea_model.state[1], 0, 150),
            np.clip(kt_v, 0, 3),
            np.clip(cv_state["systolic_bp"], 50, 200),
            np.clip(cv_state["diastolic_bp"], 30, 130),
            np.clip(cv_state["heart_rate"], 40, 160),
            self.uf_rate,
            np.clip(cv_state["total_uf_volume"] / 1000, 0, 6),
            self.patient.target_fluid_removal,
            self.qb,
            self.step_count / self.max_steps,
        ], dtype=np.float32)

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

        # Map actions
        qb_delta = float(action[0]) * 25.0
        qd_delta = float(action[1]) * 50.0
        uf_delta = float(action[2]) * 200.0

        self.qb = float(np.clip(self.qb + qb_delta, 200, 500))
        self.qd = float(np.clip(self.qd + qd_delta, 300, 800))
        self.uf_rate = float(np.clip(self.uf_rate + uf_delta, 0, self.patient.max_uf_rate))

        # Advance both models
        urea_result = self.urea_model.step(self.qb, self.qd, self.dt)
        cv_state = self.cv_model.step(self.uf_rate, self.dt)

        self.cumulative_clearance += urea_result["clearance"]
        self.clearance_steps += 1
        self.step_count += 1

        # Compute multi-objective reward
        avg_k = self.cumulative_clearance / self.clearance_steps
        kt_v = self.urea_model.compute_kt_v(avg_k, self.step_count * self.dt)

        reward = 0.0

        # Kt/V progress
        reward += urea_result["clearance"] / 1000.0

        # Fluid removal progress
        removed_L = cv_state["total_uf_volume"] / 1000
        reward += 0.3 * (self.uf_rate / 1000.0)

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

        # Parameter range penalties
        if self.qb > 450 or self.qb < 220:
            reward -= 0.3
        if self.uf_rate > self.patient.max_uf_rate * 0.9:
            reward -= 0.3

        terminated = self.step_count >= self.max_steps
        truncated = False

        if terminated:
            completion = min(1.0, removed_L / self.patient.target_fluid_removal)
            # Kt/V reward
            if kt_v >= 1.2:
                reward += 8.0
            if kt_v >= 1.4:
                reward += 4.0
            if kt_v < 1.0:
                reward -= 5.0
            # Fluid removal reward
            if completion >= 0.9 and self.hypotension_count == 0:
                reward += 8.0
            elif completion >= 0.9:
                reward += 2.0
            if completion < 0.5:
                reward -= 5.0

        info = {
            "kt_v": kt_v,
            "systolic_bp": cv_state["systolic_bp"],
            "total_removed_L": removed_L,
            "completion_fraction": removed_L / self.patient.target_fluid_removal,
            "hypotension_events": self.hypotension_count,
            "c_ecf": urea_result["c_ecf"],
        }

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