"""Phosphate management environment for weekly binder dose optimization.

The agent optimizes phosphate binder dosing across a 7-day cycle with
3 dialysis sessions per week, keeping serum phosphate within target range.
"""

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

from nephrosim.models.patient import PatientProfile, create_patient
from nephrosim.models.phosphate_kinetics import PhosphateKineticsModel


class PhosphateMgmtEnv(gym.Env):
    """Gymnasium environment for weekly phosphate management.

    Observation space (Box):
        - serum_phosphate: Current serum phosphate (mg/dL), range [1, 12]
        - calcium_phosphate_product: Ca*P product, range [10, 120]
        - dietary_phosphate: Daily dietary intake (mg), range [500, 2000]
        - binder_dose: Current binder dose (mg), range [0, 2000]
        - session_count: Dialysis sessions completed this week [0, 3]
        - days_since_session: Days since last dialysis [0, 3]

    Action space:
        Box(1) in [0, 1] mapped to binder dose (0-2000 mg).

    Episode: 21 steps (3 per day for 7 days).
    Dialysis sessions occur on steps 0, 6, 12 (Mon/Wed/Fri at start of day).
    """

    metadata = {"render_modes": []}
    DIALYSIS_STEPS = {0, 6, 12}  # MWF schedule
    STEPS_PER_DAY = 3  # morning, afternoon, evening

    def __init__(self, difficulty: str = "easy"):
        super().__init__()
        self.difficulty = difficulty
        self.max_steps = 21  # 7 days * 3 periods

        self.observation_space = spaces.Box(
            low=np.array([1, 10, 500, 0, 0, 0], dtype=np.float32),
            high=np.array([12, 120, 2000, 2000, 3, 3], dtype=np.float32),
        )
        self.action_space = spaces.Box(
            low=np.array([0], dtype=np.float32),
            high=np.array([1], dtype=np.float32),
        )

        self.patient: PatientProfile | None = None
        self.phos_model: PhosphateKineticsModel | None = None
        self.step_count = 0
        self.session_count = 0
        self.steps_since_session = 0
        self.binder_dose = 0.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.phos_model = PhosphateKineticsModel(
            ecf_volume=self.patient.ecf_volume * 1000,
            icf_volume=self.patient.icf_volume * 1000,
            dietary_phosphate_mg_per_day=self.patient.dietary_phosphate,
        )
        self.phos_model.reset(self.patient.pre_dialysis_phosphate, 5.0)

        self.step_count = 0
        self.session_count = 0
        self.steps_since_session = 3  # start as if last session was yesterday
        self.binder_dose = 0.0

        return self._get_obs(), {}

    def _get_obs(self) -> np.ndarray:
        assert self.phos_model is not None and self.patient is not None
        return np.array([
            np.clip(self.phos_model.state[0], 1, 12),
            np.clip(self.phos_model.state[0] * 9.5, 10, 120),
            self.patient.dietary_phosphate,
            self.binder_dose,
            min(self.session_count, 3),
            min(self.steps_since_session // self.STEPS_PER_DAY, 3),
        ], dtype=np.float32)

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

        # Map action to binder dose
        self.binder_dose = float(action[0]) * 2000.0

        # Check if this step is a dialysis session
        is_dialysis = self.step_count in self.DIALYSIS_STEPS

        if is_dialysis:
            # Run 4-hour dialysis session (compressed to one step)
            result = self.phos_model.step_dialysis(
                blood_flow_rate=350.0,
                dialysate_flow_rate=500.0,
                binder_dose_mg=self.binder_dose,
                dt=240.0,  # 4 hours
            )
            self.session_count += 1
            self.steps_since_session = 0
        else:
            # Inter-dialytic period (8 hours per step: 3 steps per day)
            result = self.phos_model.step_interdialytic(
                binder_dose_mg=self.binder_dose,
                dt_hours=8.0,
            )
            self.steps_since_session += 1

        self.step_count += 1

        # Reward: keep phosphate in target range (3.5-5.5 mg/dL)
        p = result["p_ecf"]
        reward = 0.0

        if 3.5 <= p <= 5.5:
            reward += 1.0  # in range
            # Bonus for being close to center (4.5)
            reward += 0.5 * (1.0 - abs(p - 4.5) / 1.0)
        else:
            reward -= abs(p - 4.5)  # penalty proportional to deviation
            if p > 7.0:
                reward -= 3.0  # severe hyperphosphatemia

        # Penalty for calcium-phosphate product > 55
        if result["calcium_phosphate_product"] > 55:
            reward -= 2.0

        # Small cost for binder use (encourage efficiency)
        reward -= 0.1 * (self.binder_dose / 2000.0)

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

        info = {
            "serum_phosphate": p,
            "calcium_phosphate_product": result["calcium_phosphate_product"],
            "binder_dose": self.binder_dose,
            "is_dialysis_step": is_dialysis,
            "session_count": self.session_count,
        }

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