"""Radiobiological tissue response models for radiation therapy.

Implements the Linear-Quadratic model, Tumor Control Probability (TCP),
Normal Tissue Complication Probability (NTCP), and Biologically Effective Dose (BED).
"""

from __future__ import annotations

import numpy as np
from scipy.special import erf


# Default radiobiological parameters
TUMOR_PARAMS = {
    "alpha": 0.3,       # Gy^-1, typical tumor
    "beta": 0.03,       # Gy^-2 (alpha/beta = 10 Gy)
    "N0": 1e9,          # Initial clonogen count
    "doubling_time": 5, # Days for tumor cell doubling
}

NORMAL_TISSUE_PARAMS = {
    "late_responding": {
        "alpha": 0.1,   # Gy^-1
        "beta": 0.033,  # Gy^-2 (alpha/beta ~ 3 Gy)
        "TD50": 50.0,   # Gy, tolerance dose for 50% complication
        "m": 0.18,      # Slope parameter for NTCP
        "n": 1.0,       # Volume dependence parameter
    },
    "early_responding": {
        "alpha": 0.35,
        "beta": 0.035,  # alpha/beta = 10 Gy
        "TD50": 55.0,
        "m": 0.15,
        "n": 0.7,
    },
}


def surviving_fraction(dose: float, alpha: float, beta: float) -> float:
    """Compute surviving fraction using the Linear-Quadratic model.

    SF = exp(-alpha * d - beta * d^2)

    Args:
        dose: Radiation dose in Gy.
        alpha: Linear damage coefficient (Gy^-1).
        beta: Quadratic damage coefficient (Gy^-2).

    Returns:
        Surviving fraction in [0, 1].
    """
    return float(np.exp(-alpha * dose - beta * dose**2))


def tcp(
    dose_per_fraction: float,
    num_fractions: int,
    alpha: float = 0.3,
    beta: float = 0.03,
    n0: float = 1e9,
) -> float:
    """Compute Tumor Control Probability using the Poisson model.

    TCP = exp(-N0 * SF^n)

    Args:
        dose_per_fraction: Dose per fraction in Gy.
        num_fractions: Number of fractions delivered.
        alpha: Linear LQ parameter.
        beta: Quadratic LQ parameter.
        n0: Initial number of clonogens.

    Returns:
        TCP value in [0, 1].
    """
    sf = surviving_fraction(dose_per_fraction, alpha, beta)
    surviving_cells = n0 * (sf ** num_fractions)
    return float(np.exp(-surviving_cells))


def ntcp_lkb(
    mean_dose: float,
    td50: float = 50.0,
    m: float = 0.18,
    n: float = 1.0,
) -> float:
    """Compute Normal Tissue Complication Probability using Lyman-Kutcher-Burman model.

    NTCP = 0.5 * (1 + erf(t / sqrt(2))), where t = (mean_dose - TD50) / (m * TD50)

    Args:
        mean_dose: Mean dose to the organ at risk in Gy.
        td50: Tolerance dose for 50% complication probability.
        m: Slope parameter controlling steepness.
        n: Volume dependence parameter.

    Returns:
        NTCP value in [0, 1].
    """
    if td50 <= 0 or m <= 0:
        return 0.0
    t = (mean_dose - td50) / (m * td50)
    return float(0.5 * (1.0 + erf(t / np.sqrt(2.0))))


def bed(dose_per_fraction: float, num_fractions: int, alpha_beta: float = 10.0) -> float:
    """Compute Biologically Effective Dose.

    BED = n * d * (1 + d / (alpha/beta))

    Args:
        dose_per_fraction: Dose per fraction in Gy.
        num_fractions: Number of fractions.
        alpha_beta: Alpha/beta ratio in Gy.

    Returns:
        BED in Gy.
    """
    return num_fractions * dose_per_fraction * (1.0 + dose_per_fraction / alpha_beta)
