"""2D beam geometry utilities for radiation therapy simulation."""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray


def generate_beam_profile(
    angle_deg: float,
    grid_shape: tuple[int, int] = (64, 64),
    beam_width_pixels: float = 10.0,
) -> NDArray[np.float64]:
    """Generate a 2D beam profile mask at a given angle.

    Args:
        angle_deg: Beam angle in degrees.
        grid_shape: Shape of the output grid.
        beam_width_pixels: Width of the beam in pixels (FWHM).

    Returns:
        2D array with beam intensity profile (0 to 1).
    """
    rows, cols = grid_shape
    center = np.array([rows / 2, cols / 2])
    angle_rad = np.deg2rad(angle_deg)

    direction = np.array([-np.sin(angle_rad), -np.cos(angle_rad)])
    perp = np.array([direction[1], -direction[0]])

    y_coords, x_coords = np.mgrid[0:rows, 0:cols]
    coords = np.stack([y_coords, x_coords], axis=-1).astype(np.float64)

    to_point = coords - center
    lateral = np.abs(np.sum(to_point * perp, axis=-1))

    sigma = beam_width_pixels / (2.0 * np.sqrt(2.0 * np.log(2.0)))
    profile = np.exp(-0.5 * (lateral / sigma) ** 2)

    return profile


def compute_beam_intersection(
    beam_profile: NDArray[np.float64],
    structure_mask: NDArray[np.float64],
) -> float:
    """Compute the dose contribution of a beam to a structure.

    Args:
        beam_profile: 2D beam intensity profile.
        structure_mask: Binary mask of the structure (tumor or OAR).

    Returns:
        Sum of beam intensity within the structure.
    """
    return float(np.sum(beam_profile * structure_mask))


def create_tumor_mask(
    center: tuple[float, float],
    radius: float,
    grid_shape: tuple[int, int] = (64, 64),
) -> NDArray[np.float64]:
    """Create a circular binary mask for a tumor.

    Args:
        center: (row, col) center of the tumor.
        radius: Radius in pixels.
        grid_shape: Shape of the output grid.

    Returns:
        Binary mask (0 or 1) as float64 array.
    """
    rows, cols = grid_shape
    y_coords, x_coords = np.mgrid[0:rows, 0:cols]
    dist = np.sqrt((y_coords - center[0]) ** 2 + (x_coords - center[1]) ** 2)
    return (dist <= radius).astype(np.float64)


def create_oar_mask(
    center: tuple[float, float],
    radius: float,
    grid_shape: tuple[int, int] = (64, 64),
) -> NDArray[np.float64]:
    """Create a circular binary mask for an organ at risk.

    Args:
        center: (row, col) center of the OAR.
        radius: Radius in pixels.
        grid_shape: Shape of the output grid.

    Returns:
        Binary mask (0 or 1) as float64 array.
    """
    return create_tumor_mask(center, radius, grid_shape)


def compute_dvh(
    dose_grid: NDArray[np.float64],
    structure_mask: NDArray[np.float64],
    num_bins: int = 100,
) -> NDArray[np.float64]:
    """Compute cumulative dose-volume histogram for a structure.

    Args:
        dose_grid: 2D dose distribution.
        structure_mask: Binary mask of the structure.
        num_bins: Number of DVH bins.

    Returns:
        1D array of length num_bins, representing fraction of volume receiving
        at least each dose level (cumulative DVH).
    """
    voxel_doses = dose_grid[structure_mask > 0.5]
    if len(voxel_doses) == 0:
        return np.zeros(num_bins, dtype=np.float64)

    max_dose = np.max(voxel_doses) + 1e-10
    dose_levels = np.linspace(0, max_dose, num_bins)

    dvh = np.array([np.mean(voxel_doses >= d) for d in dose_levels])
    return dvh
