"""Analytical dose calculation using simplified pencil beam model."""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray


def calculate_beam_dose(
    angle_deg: float,
    grid_shape: tuple[int, int] = (64, 64),
    pixel_size_mm: float = 2.0,
    beam_width_mm: float = 20.0,
    source_distance_mm: float = 1000.0,
    mu: float = 0.004,
) -> NDArray[np.float64]:
    """Calculate 2D dose distribution for a single beam at a given angle.

    Uses a simplified pencil beam model: dose = D0 * exp(-mu * depth) * gaussian_spread.

    Args:
        angle_deg: Beam angle in degrees (0 = from right, 90 = from top).
        grid_shape: Shape of the dose grid (rows, cols).
        pixel_size_mm: Physical size of each pixel in mm.
        beam_width_mm: FWHM of the beam in mm.
        source_distance_mm: Source-to-isocenter distance in mm.
        mu: Linear attenuation coefficient (1/mm).

    Returns:
        2D numpy array of dose values (normalized, arbitrary units).
    """
    rows, cols = grid_shape
    center = np.array([rows / 2, cols / 2])
    angle_rad = np.deg2rad(angle_deg)

    # Direction vector (beam travels inward from edge)
    direction = np.array([-np.sin(angle_rad), -np.cos(angle_rad)])
    # Source position (far from center, along beam direction)
    source_pos = center - direction * (source_distance_mm / pixel_size_mm)

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

    # Vector from source to each point
    to_point = coords - source_pos
    # Depth along beam direction
    depth = np.sum(to_point * direction, axis=-1) * pixel_size_mm
    # Lateral distance from beam axis
    perp_direction = np.array([direction[1], -direction[0]])
    lateral = np.abs(np.sum(to_point * perp_direction, axis=-1)) * pixel_size_mm

    # Beam sigma from FWHM
    sigma = beam_width_mm / (2.0 * np.sqrt(2.0 * np.log(2.0)))

    # Dose: exponential attenuation * gaussian lateral falloff
    dose = np.exp(-mu * np.maximum(depth, 0)) * np.exp(-0.5 * (lateral / sigma) ** 2)

    # Zero out regions behind the source
    dose[depth < 0] = 0.0

    return dose


def calculate_plan_dose(
    angles_deg: list[float],
    weights: list[float] | None,
    grid_shape: tuple[int, int] = (64, 64),
    pixel_size_mm: float = 2.0,
    beam_width_mm: float = 20.0,
    mu: float = 0.004,
) -> NDArray[np.float64]:
    """Calculate composite dose from multiple beams.

    Args:
        angles_deg: List of beam angles in degrees.
        weights: Beam weights. If None, uniform weights are used.
        grid_shape: Shape of the dose grid.
        pixel_size_mm: Pixel size in mm.
        beam_width_mm: Beam width FWHM in mm.
        mu: Linear attenuation coefficient.

    Returns:
        2D numpy array of composite dose.
    """
    if weights is None:
        weights = [1.0 / len(angles_deg)] * len(angles_deg)

    total_dose = np.zeros(grid_shape, dtype=np.float64)
    for angle, weight in zip(angles_deg, weights):
        total_dose += weight * calculate_beam_dose(
            angle, grid_shape, pixel_size_mm, beam_width_mm, mu=mu
        )

    return total_dose
