"""AdaptiveRT-v0: Adaptive replanning decisions during radiation therapy."""

from __future__ import annotations

from typing import Any

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


DIFFICULTY_PRESETS = {
    "stable": {
        "response_volatility": 0.02,
        "drift_rate": 0.005,
        "max_replans": 5,
    },
    "moderate": {
        "response_volatility": 0.08,
        "drift_rate": 0.015,
        "max_replans": 4,
    },
    "variable": {
        "response_volatility": 0.15,
        "drift_rate": 0.03,
        "max_replans": 3,
    },
}


class AdaptiveRTEnv(gym.Env):
    """Environment for adaptive radiotherapy replanning decisions.

    The agent observes the treatment progress and tumor response, then
    decides whether to replan and how to adjust the dose. Replanning
    improves plan quality but incurs a cost (clinical time and resources).
    """

    metadata = {"render_modes": [None]}

    def __init__(
        self,
        difficulty: str = "moderate",
        treatment_sessions: int = 30,
        replan_cost: float = 0.3,
        render_mode: str | None = None,
    ):
        super().__init__()
        self.treatment_sessions = treatment_sessions
        self.replan_cost = replan_cost
        self.render_mode = render_mode

        preset = DIFFICULTY_PRESETS[difficulty]
        self.response_volatility = preset["response_volatility"]
        self.drift_rate = preset["drift_rate"]
        self.max_replans = preset["max_replans"]

        self.observation_space = spaces.Dict({
            "treatment_progress": spaces.Box(0, 1, shape=(1,), dtype=np.float64),
            "plan_quality": spaces.Box(0, 1, shape=(1,), dtype=np.float64),
            "tumor_response": spaces.Box(-1, 1, shape=(1,), dtype=np.float64),
            "dose_deviation": spaces.Box(0, 2, shape=(1,), dtype=np.float64),
            "replans_remaining": spaces.Box(0, self.max_replans, shape=(1,), dtype=np.float64),
        })

        # Action: [replan (0=no, 1=yes), dose_adjustment (0-4 maps to 0.8-1.2x)]
        self.action_space = spaces.MultiDiscrete([2, 5])

        self._session: int = 0
        self._plan_quality: float = 0.0
        self._tumor_response: float = 0.0
        self._dose_deviation: float = 0.0
        self._replans_remaining: int = 0
        self._base_plan_quality: float = 0.0

    def _get_obs(self) -> dict[str, np.ndarray]:
        progress = self._session / max(self.treatment_sessions, 1)
        return {
            "treatment_progress": np.array([progress], dtype=np.float64),
            "plan_quality": np.array([self._plan_quality], dtype=np.float64),
            "tumor_response": np.array([self._tumor_response], dtype=np.float64),
            "dose_deviation": np.array([self._dose_deviation], dtype=np.float64),
            "replans_remaining": np.array([self._replans_remaining], dtype=np.float64),
        }

    def reset(
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, np.ndarray], dict[str, Any]]:
        super().reset(seed=seed)
        self._session = 0
        self._plan_quality = 0.85  # Initial plan quality (good but not perfect)
        self._base_plan_quality = 0.85
        self._tumor_response = 0.0
        self._dose_deviation = 0.0
        self._replans_remaining = self.max_replans
        return self._get_obs(), {}

    def step(
        self, action: np.ndarray
    ) -> tuple[dict[str, np.ndarray], float, bool, bool, dict[str, Any]]:
        replan = int(action[0])
        dose_adj_idx = int(action[1])

        # Dose adjustment factor: 0->0.8, 1->0.9, 2->1.0, 3->1.1, 4->1.2
        dose_factor = 0.8 + dose_adj_idx * 0.1

        self._session += 1
        prev_quality = self._plan_quality

        # Simulate tumor response (random perturbations)
        response_change = self.np_random.normal(0, self.response_volatility)
        self._tumor_response = np.clip(
            self._tumor_response + response_change, -1.0, 1.0
        )

        # Plan quality degrades due to anatomical changes
        quality_drift = -self.drift_rate * (1.0 + abs(self._tumor_response))
        self._plan_quality = np.clip(self._plan_quality + quality_drift, 0.0, 1.0)

        # Dose deviation increases over time
        self._dose_deviation = min(
            2.0, self._dose_deviation + abs(self._tumor_response) * 0.1
        )

        # Apply dose adjustment (modifies deviation)
        if dose_factor != 1.0:
            self._dose_deviation *= abs(1.0 - (dose_factor - 1.0))

        reward = 0.0

        # Replanning action
        if replan == 1 and self._replans_remaining > 0:
            self._replans_remaining -= 1
            quality_improvement = 0.1 + 0.05 * self._dose_deviation
            self._plan_quality = min(1.0, self._plan_quality + quality_improvement)
            self._dose_deviation *= 0.3  # Replan reduces deviation significantly
            reward -= self.replan_cost
        elif replan == 1 and self._replans_remaining <= 0:
            reward -= 0.1  # Penalty for trying to replan with none remaining

        # Reward for maintaining/improving plan quality
        quality_change = self._plan_quality - prev_quality
        reward += quality_change + self._plan_quality * 0.1

        terminated = self._session >= self.treatment_sessions
        truncated = False

        # Terminal bonus/penalty
        if terminated:
            if self._plan_quality > 0.7:
                reward += 2.0 * self._plan_quality
            else:
                reward -= 1.0

        info = {
            "session": self._session,
            "plan_quality": float(self._plan_quality),
            "tumor_response": float(self._tumor_response),
            "dose_deviation": float(self._dose_deviation),
            "replans_remaining": self._replans_remaining,
            "replanned": bool(replan == 1 and self._replans_remaining >= 0),
        }

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