"""DoseFractionation-v0: Sequential daily dose fractionation decisions."""

from __future__ import annotations

from typing import Any

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

from oncosim.physics.tissue_models import surviving_fraction, tcp, ntcp_lkb


DIFFICULTY_PRESETS = {
    "radiosensitive": {
        "alpha": 0.4,
        "beta": 0.04,
        "td50": 55.0,
        "tumor_doubling_days": 7,
    },
    "standard": {
        "alpha": 0.3,
        "beta": 0.03,
        "td50": 50.0,
        "tumor_doubling_days": 5,
    },
    "radioresistant": {
        "alpha": 0.15,
        "beta": 0.015,
        "td50": 45.0,
        "tumor_doubling_days": 3,
    },
}


class DoseFractionationEnv(gym.Env):
    """Environment for radiation dose fractionation decisions.

    The agent decides the dose for each treatment fraction. The goal is to
    maximize tumor control probability (TCP) while minimizing normal tissue
    complication probability (NTCP).
    """

    metadata = {"render_modes": [None]}

    def __init__(
        self,
        difficulty: str = "standard",
        max_fractions: int = 30,
        tcp_weight: float = 1.0,
        ntcp_weight: float = 2.0,
        render_mode: str | None = None,
    ):
        super().__init__()
        self.max_fractions = max_fractions
        self.tcp_weight = tcp_weight
        self.ntcp_weight = ntcp_weight
        self.render_mode = render_mode

        preset = DIFFICULTY_PRESETS[difficulty]
        self.alpha = preset["alpha"]
        self.beta = preset["beta"]
        self.td50 = preset["td50"]
        self.tumor_doubling_days = preset["tumor_doubling_days"]

        self.observation_space = spaces.Dict({
            "fraction_number": spaces.Box(0, max_fractions, shape=(1,), dtype=np.float64),
            "tumor_volume": spaces.Box(0, 2, shape=(1,), dtype=np.float64),
            "cumulative_dose": spaces.Box(0, 200, shape=(1,), dtype=np.float64),
            "healthy_tissue_state": spaces.Box(0, 1, shape=(1,), dtype=np.float64),
            "tcp": spaces.Box(0, 1, shape=(1,), dtype=np.float64),
            "ntcp": spaces.Box(0, 1, shape=(1,), dtype=np.float64),
        })

        # Action: dose in Gy for this fraction (continuous, 0.5 to 4.0 Gy)
        self.action_space = spaces.Box(
            low=np.float32(0.5), high=np.float32(4.0), shape=(1,), dtype=np.float32
        )

        self._fraction: int = 0
        self._tumor_volume: float = 1.0
        self._cumulative_dose: float = 0.0
        self._tissue_damage: float = 0.0
        self._doses_given: list[float] = []

    def _compute_tcp(self) -> float:
        if not self._doses_given:
            return 0.0
        avg_dose = np.mean(self._doses_given)
        return tcp(avg_dose, len(self._doses_given), self.alpha, self.beta)

    def _compute_ntcp(self) -> float:
        return ntcp_lkb(self._cumulative_dose, self.td50)

    def _get_obs(self) -> dict[str, np.ndarray]:
        return {
            "fraction_number": np.array([self._fraction], dtype=np.float64),
            "tumor_volume": np.array([self._tumor_volume], dtype=np.float64),
            "cumulative_dose": np.array([self._cumulative_dose], dtype=np.float64),
            "healthy_tissue_state": np.array([self._tissue_damage], dtype=np.float64),
            "tcp": np.array([self._compute_tcp()], dtype=np.float64),
            "ntcp": np.array([self._compute_ntcp()], 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._fraction = 0
        self._tumor_volume = 1.0
        self._cumulative_dose = 0.0
        self._tissue_damage = 0.0
        self._doses_given = []
        return self._get_obs(), {}

    def step(
        self, action: np.ndarray
    ) -> tuple[dict[str, np.ndarray], float, bool, bool, dict[str, Any]]:
        dose = float(np.clip(action[0], 0.5, 4.0))

        # Record previous state
        prev_tcp = self._compute_tcp()
        prev_ntcp = self._compute_ntcp()

        # Apply dose
        self._fraction += 1
        self._cumulative_dose += dose
        self._doses_given.append(dose)

        # Tumor response: LQ cell kill
        sf = surviving_fraction(dose, self.alpha, self.beta)
        self._tumor_volume *= sf

        # Tumor regrowth between fractions
        growth_factor = np.log(2) / self.tumor_doubling_days
        self._tumor_volume *= np.exp(growth_factor)
        self._tumor_volume = min(self._tumor_volume, 2.0)

        # Tissue damage accumulation (simplified)
        self._tissue_damage = min(1.0, self._cumulative_dose / (2.0 * self.td50))

        # Compute new metrics
        new_tcp = self._compute_tcp()
        new_ntcp = self._compute_ntcp()

        # Reward: TCP improvement minus NTCP increase
        delta_tcp = new_tcp - prev_tcp
        delta_ntcp = new_ntcp - prev_ntcp
        reward = self.tcp_weight * delta_tcp - self.ntcp_weight * delta_ntcp

        # Termination conditions
        terminated = False
        if self._fraction >= self.max_fractions:
            terminated = True
        elif new_tcp > 0.95:
            terminated = True
            reward += 1.0  # Bonus for achieving high TCP
        elif new_ntcp > 0.5:
            terminated = True
            reward -= 1.0  # Penalty for excessive toxicity

        info = {
            "fraction": self._fraction,
            "dose_given": dose,
            "tumor_volume": self._tumor_volume,
            "tcp": new_tcp,
            "ntcp": new_ntcp,
            "cumulative_dose": self._cumulative_dose,
        }

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