"""BeamSelection-v0: Select optimal beam angles for 2D IMRT planning."""

from __future__ import annotations

from typing import Any

import gymnasium as gym
import numpy as np
from gymnasium import spaces
from numpy.typing import NDArray

from oncosim.physics.dose_calc import calculate_beam_dose
from oncosim.physics.beam_geometry import create_tumor_mask, create_oar_mask


# Difficulty presets
DIFFICULTY_PRESETS = {
    "easy": {
        "tumor_center": (32, 32),
        "tumor_radius": 10.0,
        "oar_centers": [(10, 10), (54, 54), (10, 54)],
        "oar_radii": [5.0, 5.0, 5.0],
        "num_beams": 5,
    },
    "medium": {
        "tumor_center": (30, 35),
        "tumor_radius": 7.0,
        "oar_centers": [(25, 25), (40, 45), (20, 50)],
        "oar_radii": [6.0, 5.0, 5.0],
        "num_beams": 7,
    },
    "hard": {
        "tumor_center": (32, 32),
        "tumor_radius": 5.0,
        "oar_centers": [(28, 38), (38, 28), (26, 26)],
        "oar_radii": [7.0, 6.0, 6.0],
        "num_beams": 7,
    },
}

NUM_ANGLES = 36  # 0, 10, 20, ..., 350 degrees


class BeamSelectionEnv(gym.Env):
    """Environment for sequential beam angle selection in IMRT planning.

    The agent selects beam angles one at a time. After each selection,
    the dose distribution updates. The goal is to maximize tumor coverage
    while minimizing dose to organs at risk (OARs).
    """

    metadata = {"render_modes": [None]}

    def __init__(
        self,
        difficulty: str = "medium",
        grid_size: int = 64,
        oar_weight: float = 0.5,
        render_mode: str | None = None,
    ):
        super().__init__()
        self.grid_size = grid_size
        self.oar_weight = oar_weight
        self.render_mode = render_mode

        preset = DIFFICULTY_PRESETS[difficulty]
        self.tumor_center = preset["tumor_center"]
        self.tumor_radius = preset["tumor_radius"]
        self.oar_centers = preset["oar_centers"]
        self.oar_radii = preset["oar_radii"]
        self.num_beams = preset["num_beams"]
        self.num_oars = len(self.oar_centers)

        # Create structure masks
        gs = (grid_size, grid_size)
        self.tumor_mask = create_tumor_mask(self.tumor_center, self.tumor_radius, gs)
        self.oar_masks = np.stack(
            [create_oar_mask(c, r, gs) for c, r in zip(self.oar_centers, self.oar_radii)]
        )

        # Precompute beam doses for all candidate angles
        self._beam_doses = np.stack(
            [calculate_beam_dose(a * 10.0, gs) for a in range(NUM_ANGLES)]
        )

        self.observation_space = spaces.Dict({
            "dose_grid": spaces.Box(0, np.inf, shape=(grid_size, grid_size), dtype=np.float64),
            "tumor_mask": spaces.Box(0, 1, shape=(grid_size, grid_size), dtype=np.float64),
            "oar_masks": spaces.Box(0, 1, shape=(self.num_oars, grid_size, grid_size), dtype=np.float64),
            "selected_beams": spaces.Box(0, 1, shape=(NUM_ANGLES,), dtype=np.float64),
            "num_selected": spaces.Box(0, self.num_beams, shape=(1,), dtype=np.float64),
        })

        self.action_space = spaces.Discrete(NUM_ANGLES)

        self._dose_grid: NDArray[np.float64] = np.zeros(gs, dtype=np.float64)
        self._selected: NDArray[np.float64] = np.zeros(NUM_ANGLES, dtype=np.float64)
        self._n_selected: int = 0

    def _get_obs(self) -> dict[str, NDArray]:
        return {
            "dose_grid": self._dose_grid.copy(),
            "tumor_mask": self.tumor_mask.copy(),
            "oar_masks": self.oar_masks.copy(),
            "selected_beams": self._selected.copy(),
            "num_selected": np.array([self._n_selected], dtype=np.float64),
        }

    def _compute_plan_quality(self) -> tuple[float, float]:
        """Compute tumor coverage and OAR dose metrics."""
        tumor_dose = self._dose_grid * self.tumor_mask
        tumor_mean = np.sum(tumor_dose) / (np.sum(self.tumor_mask) + 1e-10)

        oar_total = 0.0
        for i in range(self.num_oars):
            oar_dose = self._dose_grid * self.oar_masks[i]
            oar_total += np.sum(oar_dose) / (np.sum(self.oar_masks[i]) + 1e-10)
        oar_mean = oar_total / max(self.num_oars, 1)

        return float(tumor_mean), float(oar_mean)

    def reset(
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, NDArray], dict[str, Any]]:
        super().reset(seed=seed)
        gs = (self.grid_size, self.grid_size)
        self._dose_grid = np.zeros(gs, dtype=np.float64)
        self._selected = np.zeros(NUM_ANGLES, dtype=np.float64)
        self._n_selected = 0
        return self._get_obs(), {}

    def step(self, action: int) -> tuple[dict[str, NDArray], float, bool, bool, dict[str, Any]]:
        # Penalty for selecting an already-chosen angle
        duplicate_penalty = 0.0
        if self._selected[action] > 0.5:
            duplicate_penalty = -0.5

        # Add beam dose
        self._dose_grid += self._beam_doses[action]
        self._selected[action] = 1.0
        self._n_selected += 1

        prev_coverage, prev_oar = self._compute_plan_quality()
        coverage, oar_dose = self._compute_plan_quality()

        reward = coverage - self.oar_weight * oar_dose + duplicate_penalty

        terminated = self._n_selected >= self.num_beams
        truncated = False

        info = {
            "tumor_coverage": coverage,
            "oar_dose": oar_dose,
            "num_selected": self._n_selected,
            "plan_quality": coverage - self.oar_weight * oar_dose,
        }

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