"""Clinical heuristic baseline agents for OncoSim environments."""

from __future__ import annotations

import gymnasium as gym
import numpy as np


class BeamSelectionHeuristic:
    """Selects equispaced beam angles (standard clinical approach)."""

    def __init__(self, env: gym.Env, num_beams: int = 7):
        self.num_beams = num_beams
        self._step = 0
        # Equispaced angles: e.g., for 7 beams -> 0, 51, 103, 154, 206, 257, 309 deg
        # Map to nearest discrete angle index (0-35, each = 10 deg)
        spacing = 360.0 / num_beams
        self._angles = [int(round((i * spacing) / 10.0)) % 36 for i in range(num_beams)]

    def act(self, obs: dict | np.ndarray) -> int:
        if self._step < len(self._angles):
            action = self._angles[self._step]
        else:
            action = 0
        self._step += 1
        return action

    def reset(self) -> None:
        self._step = 0


class DoseFractionationHeuristic:
    """Uses standard 2 Gy per fraction (conventional fractionation)."""

    def act(self, obs: dict | np.ndarray) -> np.ndarray:
        return np.array([2.0], dtype=np.float32)

    def reset(self) -> None:
        pass


class AdaptiveRTHeuristic:
    """Conservative approach: never replans, keeps dose at 1.0x."""

    def act(self, obs: dict | np.ndarray) -> np.ndarray:
        return np.array([0, 2], dtype=np.int64)  # No replan, dose_factor=1.0

    def reset(self) -> None:
        pass


def evaluate_heuristic(
    env: gym.Env,
    agent: BeamSelectionHeuristic | DoseFractionationHeuristic | AdaptiveRTHeuristic,
    n_episodes: int = 100,
    seed: int = 42,
) -> dict[str, float]:
    """Evaluate a heuristic agent.

    Returns:
        Dict with mean_reward, std_reward, min_reward, max_reward.
    """
    rewards = []
    for ep in range(n_episodes):
        obs, _ = env.reset(seed=seed + ep)
        agent.reset()
        total_reward = 0.0
        done = False
        while not done:
            action = agent.act(obs)
            obs, reward, terminated, truncated, _ = env.step(action)
            total_reward += reward
            done = terminated or truncated
        rewards.append(total_reward)

    return {
        "mean_reward": float(np.mean(rewards)),
        "std_reward": float(np.std(rewards)),
        "min_reward": float(np.min(rewards)),
        "max_reward": float(np.max(rewards)),
    }
