radiacode/train/vega_ml/synthetic_spectra/physics/spectrum_physics.py

"""
Spectrum Physics Module

Implements the physics of gamma spectrum generation including:
- Peak shape modeling (Gaussian with detector response)
- Background continuum generation
- Counting statistics (Poisson sampling)
- Detector efficiency modeling
"""

import numpy as np
from pathlib import Path
from scipy import special
from typing import Optional, Tuple, List
from dataclasses import dataclass

from ..config import DetectorConfig, get_default_config


@dataclass
class PeakParameters:
    """Parameters for a single gamma peak."""
    energy_kev: float
    intensity: float  # Emission probability (photons/decay)
    activity_bq: float  # Source activity in Becquerels
    live_time_s: float  # Acquisition time in seconds


def gaussian_peak(
    energy_bins: np.ndarray,
    peak_energy: float,
    sigma: float,
    amplitude: float
) -> np.ndarray:
    """
    Generate a Gaussian peak.
    
    Args:
        energy_bins: Array of energy bin centers (keV)
        peak_energy: Center energy of peak (keV)
        sigma: Standard deviation (keV)
        amplitude: Peak area (total counts)
    
    Returns:
        Array of counts in each bin
    """
    # Gaussian probability density
    prob = np.exp(-0.5 * ((energy_bins - peak_energy) / sigma) ** 2)
    prob /= (sigma * np.sqrt(2 * np.pi))
    
    # Scale by amplitude and bin width
    bin_width = energy_bins[1] - energy_bins[0] if len(energy_bins) > 1 else 1.0
    return amplitude * prob * bin_width


def calculate_fwhm(energy_kev: float, fwhm_at_662: float = 0.084) -> float:
    """
    Calculate FWHM at a given energy for scintillator detectors.
    
    FWHM scales as sqrt(E) for scintillators due to statistical fluctuations
    in light collection.
    
    FWHM(E) = FWHM_662 * sqrt(E/662) * 662 / E * E = FWHM_662 * sqrt(662/E) * E
    Actually: FWHM(E) / E = FWHM_662 / 662 * sqrt(662/E)
    So: FWHM(E) = E * FWHM_662 / 662 * sqrt(662/E) = FWHM_662 * sqrt(662 * E) / 662
                = FWHM_662 * sqrt(E / 662)
    
    Wait, let me recalculate:
    For scintillators, the relative resolution (FWHM/E) scales as 1/sqrt(E)
    FWHM(E)/E = (FWHM_662/662) * sqrt(662/E)
    FWHM(E) = FWHM_662 * sqrt(662 * E) / 662 = FWHM_662 * sqrt(E/662)
    
    At 662 keV: FWHM = FWHM_662 * sqrt(1) = FWHM_662 ✓
    At lower E: larger relative FWHM (worse resolution)
    At higher E: smaller relative FWHM (better resolution)
    
    Args:
        energy_kev: Energy in keV
        fwhm_at_662: FWHM at 662 keV as fraction (e.g., 0.084 for 8.4%)
    
    Returns:
        FWHM in keV at the given energy
    """
    # FWHM_662 is given as fraction, so at 662 keV, FWHM = 0.084 * 662 = ~55.6 keV
    fwhm_662_kev = fwhm_at_662 * 662.0
    # Scale by sqrt(E/662)
    fwhm_kev = fwhm_662_kev * np.sqrt(energy_kev / 662.0)
    return fwhm_kev


def fwhm_to_sigma(fwhm: float) -> float:
    """Convert FWHM to Gaussian sigma."""
    return fwhm / (2.0 * np.sqrt(2.0 * np.log(2.0)))  # ≈ FWHM / 2.355


def detector_efficiency(
    energy_kev: float,
    detector_config: Optional[DetectorConfig] = None
) -> float:
    """
    Calculate detector full-energy peak efficiency.
    
    For CsI and GAGG scintillators, efficiency varies with energy.
    This is a simplified model - real efficiency curves should be
    measured for each detector.
    
    Args:
        energy_kev: Gamma energy in keV
        detector_config: Detector configuration
    
    Returns:
        Efficiency as fraction (0-1)
    """
    if detector_config is None:
        detector_config = get_default_config()
    
    # Simplified efficiency model for ~1 cm³ scintillator
    # Low energy: efficiency increases (more stopping power)
    # High energy: efficiency decreases (photons pass through)
    # Peak around 100-300 keV for small scintillators
    
    # This is a phenomenological model
    # Real efficiency should be calibrated
    
    if energy_kev < 20:
        return 0.0
    
    # Simple model: efficiency peaks around 100-200 keV
    # Falls off at low energy (absorption in housing)
    # Falls off at high energy (less stopping power)
    
    # Low energy cutoff (absorption)
    low_eff = 1.0 - np.exp(-energy_kev / 50.0)
    
    # High energy falloff (escape)
    # For 1 cm³ CsI, efficiency drops significantly above ~500 keV
    high_eff = np.exp(-energy_kev / 2000.0)
    
    # Combine effects
    eff = 0.8 * low_eff * high_eff
    
    # Scale by detector volume
    volume_factor = (detector_config.detector_volume_cm3 / 1.0) ** (1/3)
    eff *= min(1.0, volume_factor)
    
    return max(0.0, min(1.0, eff))


def calculate_expected_counts(
    peak_params: PeakParameters,
    detector_config: Optional[DetectorConfig] = None
) -> float:
    """
    Calculate expected counts in a photopeak.
    
    λ = A * t * I * ε * T
    
    Where:
        A = activity (decays/s)
        t = live time (s)
        I = emission probability (photons/decay)
        ε = detector efficiency
        T = transmission factor (assumed 1 for now)
    
    Args:
        peak_params: Peak parameters
        detector_config: Detector configuration
    
    Returns:
        Expected number of counts in the photopeak
    """
    if detector_config is None:
        detector_config = get_default_config()
    
    efficiency = detector_efficiency(peak_params.energy_kev, detector_config)
    
    expected = (
        peak_params.activity_bq *
        peak_params.live_time_s *
        peak_params.intensity *
        efficiency
    )
    
    return expected


def generate_peak_spectrum(
    energy_bins: np.ndarray,
    peak_params: PeakParameters,
    detector_config: Optional[DetectorConfig] = None
) -> np.ndarray:
    """
    Generate a single gamma peak with detector response.
    
    Args:
        energy_bins: Array of energy bin centers (keV)
        peak_params: Peak parameters
        detector_config: Detector configuration
    
    Returns:
        Array of expected counts in each bin (not yet Poisson sampled)
    """
    if detector_config is None:
        detector_config = get_default_config()
    
    # Calculate expected counts
    amplitude = calculate_expected_counts(peak_params, detector_config)
    
    if amplitude <= 0:
        return np.zeros_like(energy_bins)
    
    # Calculate peak width
    fwhm_kev = calculate_fwhm(peak_params.energy_kev, detector_config.fwhm_at_662)
    sigma = fwhm_to_sigma(fwhm_kev)
    
    # Generate Gaussian peak
    peak = gaussian_peak(energy_bins, peak_params.energy_kev, sigma, amplitude)
    
    return peak


def generate_compton_continuum(
    energy_bins: np.ndarray,
    peak_energy: float,
    peak_counts: float,
    compton_to_peak_ratio: float = 0.5
) -> np.ndarray:
    """
    Generate simplified Compton continuum for a gamma line.
    
    The Compton continuum extends from 0 to the Compton edge.
    Compton edge energy = E * (1 - 1/(1 + 2*E/(511)))
    
    Args:
        energy_bins: Array of energy bin centers (keV)
        peak_energy: Energy of the gamma line (keV)
        peak_counts: Total counts in the photopeak
        compton_to_peak_ratio: Ratio of Compton counts to peak counts
    
    Returns:
        Array of Compton continuum counts
    """
    # Compton edge energy
    alpha = peak_energy / 511.0  # E / m_e c²
    compton_edge = peak_energy * (2 * alpha) / (1 + 2 * alpha)
    
    # Create continuum (simplified flat + edge shape)
    continuum = np.zeros_like(energy_bins)
    
    # Mask for energies below Compton edge
    mask = energy_bins < compton_edge
    
    if np.any(mask):
        # Simple model: roughly flat with enhancement near edge
        base_level = peak_counts * compton_to_peak_ratio / np.sum(mask)
        continuum[mask] = base_level
        
        # Add edge enhancement (Klein-Nishina-like shape)
        edge_region = (energy_bins > 0.8 * compton_edge) & (energy_bins < compton_edge)
        if np.any(edge_region):
            enhancement = 1.5 * np.exp(-((energy_bins[edge_region] - compton_edge) / (0.05 * compton_edge)) ** 2)
            continuum[edge_region] *= (1 + enhancement)
    
    return continuum


# =============================================================================
# BACKGROUND GENERATION
# =============================================================================

def generate_exponential_background(
    energy_bins: np.ndarray,
    amplitude: float = 100.0,
    decay_constant: float = 0.003
) -> np.ndarray:
    """
    Generate exponential background continuum.

    B(E) = A * exp(-b * E)

    Args:
        energy_bins: Array of energy bin centers (keV)
        amplitude: Background amplitude at E=0
        decay_constant: Exponential decay constant (1/keV)

    Returns:
        Array of background counts
    """
    return amplitude * np.exp(-decay_constant * energy_bins)


def generate_polynomial_background(
    energy_bins: np.ndarray,
    coefficients: List[float] = None
) -> np.ndarray:
    """
    Generate polynomial background.

    B(E) = Σ c_m * E^m

    Args:
        energy_bins: Array of energy bin centers (keV)
        coefficients: Polynomial coefficients [c0, c1, c2, ...]

    Returns:
        Array of background counts
    """
    if coefficients is None:
        coefficients = [10.0, -0.005, 1e-6]  # Default quadratic

    background = np.zeros_like(energy_bins)
    for m, c in enumerate(coefficients):
        background += c * (energy_bins ** m)

    return np.maximum(0, background)


def generate_realistic_continuum(
    energy_bins: np.ndarray,
    total_counts: float,
    detector_config: Optional[DetectorConfig] = None
) -> np.ndarray:
    """
    Generate realistic CsI(Tl) background continuum shape.

    Calibrated against real Radiacode 103 background measurements.
    Produces the characteristic asymmetric hump at ~110 keV with
    housing absorption at low energy, Compton plateau, and proper
    high-energy falloff.

    Shape components:
    - Asymmetric hump at ~110 keV (sigma_left=48, sigma_right=80 keV)
    - Housing absorption below ~40 keV: 1 - exp(-E/30)
    - Compton plateau around 200-260 keV from Pb-214/Bi-214 scatter
    - Compton tail: 0.38*exp(-E/170) + 0.06*exp(-E/500)
    - Noise floor at 0.3% of peak

    Args:
        energy_bins: Array of energy bin centers (keV)
        total_counts: Target total counts in the continuum
        detector_config: Detector configuration (unused, kept for API consistency)

    Returns:
        Array of background counts matching real CsI(Tl) continuum shape
    """
    E = energy_bins

    # Asymmetric hump at ~110 keV
    # Left side sharper (sigma=48), right side broader with Compton shoulder (sigma=80)
    hump_center = 110.0
    hump = np.where(
        E <= hump_center,
        np.exp(-0.5 * ((E - hump_center) / 48.0) ** 2),
        np.exp(-0.5 * ((E - hump_center) / 80.0) ** 2),
    )

    # Housing absorption at very low energy (< ~40 keV)
    absorption = 1.0 - np.exp(-E / 30.0)

    # Compton continuum tail
    tail = 0.38 * np.exp(-E / 170.0) + 0.06 * np.exp(-E / 500.0)

    # Compton plateau around 200-260 keV (Pb-214/Bi-214 scatter)
    compton_edge = np.maximum(0, 1.0 - ((E - 180.0) / 150.0) ** 2)
    compton_edge[E > 330] = 0
    compton_plateau = 0.12 * compton_edge

    # Noise floor
    noise_floor = 0.003

    # Combine continuum with absorption
    continuum = (hump + tail + compton_plateau) * absorption + noise_floor

    # Normalize to target total counts
    if continuum.sum() > 0 and total_counts > 0:
        continuum *= total_counts / continuum.sum()

    return continuum


def load_measured_background(
    path: str,
    energy_bins: np.ndarray,
    duration_seconds: float
) -> Optional[np.ndarray]:
    """
    Load a measured background spectrum from a .npy file and rescale it
    to match the target duration.

    The .npy file should contain a dict with keys 'counts' and 'duration'.

    Args:
        path: Path to the .npy background file
        energy_bins: Array of energy bin centers (keV) for alignment
        duration_seconds: Target duration to scale the spectrum to

    Returns:
        Background spectrum scaled to target duration, or None if file not found
    """
    bg_path = Path(path)
    if not bg_path.exists():
        return None

    try:
        bg_data = np.load(str(bg_path), allow_pickle=True).item()
        bg_counts = bg_data["counts"].astype(np.float64)
        bg_duration = float(bg_data["duration"])

        # Truncate or pad to match energy_bins length
        num_channels = len(energy_bins)
        if len(bg_counts) > num_channels:
            bg_counts = bg_counts[:num_channels]
        elif len(bg_counts) < num_channels:
            bg_counts = np.pad(bg_counts, (0, num_channels - len(bg_counts)))

        # Scale to target duration (cps * target_duration)
        if bg_duration > 0:
            scale = duration_seconds / bg_duration
            return bg_counts * scale
        return None
    except Exception:
        return None


def generate_environmental_background(
    energy_bins: np.ndarray,
    duration_seconds: float,
    background_cps: float = 5.0,
    include_k40: bool = True,
    include_radon: bool = True,
    include_thorium: bool = True,
    detector_config: Optional[DetectorConfig] = None,
    measured_background_path: Optional[str] = None
) -> Tuple[np.ndarray, List[str]]:
    """
    Generate realistic environmental background spectrum.

    Includes:
    - Realistic CsI(Tl) continuum shape (asymmetric hump + Compton tail)
    - Or measured background if path provided and file exists
    - K-40 peak (1460 keV) - ubiquitous in environment
    - Radon daughters (Pb-214, Bi-214) - indoor air
    - Thorium daughters (Pb-212, Tl-208) - building materials

    Args:
        energy_bins: Array of energy bin centers (keV)
        duration_seconds: Acquisition time
        background_cps: Average background count rate (cps)
        include_k40: Include potassium-40 peak
        include_radon: Include radon daughter peaks
        include_thorium: Include thorium daughter peaks
        detector_config: Detector configuration
        measured_background_path: Path to .npy file with measured background.
            If provided and file exists, used as the continuum base instead
            of the synthetic continuum. Isotope peaks are still added on top
            with stochastic variation for training diversity.

    Returns:
        Tuple of (background_spectrum, list_of_background_isotopes)
    """
    if detector_config is None:
        detector_config = get_default_config()

    background_isotopes = []

    # Use measured background if available, otherwise synthetic continuum
    total_continuum_counts = background_cps * duration_seconds * 0.7

    measured = None
    if measured_background_path:
        measured = load_measured_background(
            measured_background_path, energy_bins, duration_seconds
        )

    if measured is not None:
        # Scale measured background to match target CPS
        measured_total = measured.sum()
        if measured_total > 0 and total_continuum_counts > 0:
            # Blend: 70% measured shape, 30% synthetic for robustness
            synthetic = generate_realistic_continuum(
                energy_bins, total_counts=total_continuum_counts * 0.3,
                detector_config=detector_config
            )
            measured_scaled = measured * (total_continuum_counts * 0.7 / measured_total)
            background = measured_scaled + synthetic
        else:
            background = measured
    else:
        background = generate_realistic_continuum(
            energy_bins, total_counts=total_continuum_counts,
            detector_config=detector_config
        )

    # Add K-40 peak (very common)
    if include_k40:
        k40_activity = np.random.uniform(0.5, 5.0)  # Bq
        peak = generate_peak_spectrum(
            energy_bins,
            PeakParameters(
                energy_kev=1460.83,
                intensity=0.1066,
                activity_bq=k40_activity,
                live_time_s=duration_seconds
            ),
            detector_config
        )
        background += peak
        background_isotopes.append("K-40")

    # Add radon daughters
    if include_radon:
        radon_activity = np.random.uniform(0.1, 2.0)  # Bq

        # Pb-214 lines
        for energy, intensity in [(295.22, 0.1842), (351.93, 0.356)]:
            peak = generate_peak_spectrum(
                energy_bins,
                PeakParameters(
                    energy_kev=energy,
                    intensity=intensity,
                    activity_bq=radon_activity,
                    live_time_s=duration_seconds
                ),
                detector_config
            )
            background += peak

        # Bi-214 lines
        for energy, intensity in [(609.31, 0.4549), (1120.29, 0.1492), (1764.49, 0.1531)]:
            peak = generate_peak_spectrum(
                energy_bins,
                PeakParameters(
                    energy_kev=energy,
                    intensity=intensity,
                    activity_bq=radon_activity,
                    live_time_s=duration_seconds
                ),
                detector_config
            )
            background += peak

        background_isotopes.extend(["Pb-214", "Bi-214"])

    # Add thorium daughters
    if include_thorium:
        thorium_activity = np.random.uniform(0.05, 1.0)  # Bq

        # Ac-228 line
        peak = generate_peak_spectrum(
            energy_bins,
            PeakParameters(
                energy_kev=911.20,
                intensity=0.258,
                activity_bq=thorium_activity,
                live_time_s=duration_seconds
            ),
            detector_config
        )
        background += peak

        # Pb-212 line
        peak = generate_peak_spectrum(
            energy_bins,
            PeakParameters(
                energy_kev=238.63,
                intensity=0.436,
                activity_bq=thorium_activity,
                live_time_s=duration_seconds
            ),
            detector_config
        )
        background += peak

        # Tl-208 lines
        for energy, intensity in [(583.19, 0.845 * 0.36), (2614.51, 0.998 * 0.36)]:
            # Branching ratio of 36% for Tl-208 path
            peak = generate_peak_spectrum(
                energy_bins,
                PeakParameters(
                    energy_kev=energy,
                    intensity=intensity,
                    activity_bq=thorium_activity,
                    live_time_s=duration_seconds
                ),
                detector_config
            )
            background += peak

        background_isotopes.extend(["Ac-228", "Pb-212", "Tl-208"])

    return background, background_isotopes


def apply_poisson_noise(spectrum: np.ndarray) -> np.ndarray:
    """
    Apply Poisson counting statistics to a spectrum.
    
    Each bin is sampled from a Poisson distribution with
    lambda = expected counts in that bin.
    
    Args:
        spectrum: Array of expected counts (can be float)
    
    Returns:
        Array of actual counts (integers)
    """
    # Handle negative values (shouldn't happen but be safe)
    spectrum = np.maximum(0, spectrum)
    
    # Sample from Poisson distribution
    return np.random.poisson(spectrum).astype(np.float64)


def apply_electronic_noise(
    spectrum: np.ndarray,
    sigma: float = 0.5
) -> np.ndarray:
    """
    Apply small Gaussian electronic noise.
    
    Args:
        spectrum: Count spectrum
        sigma: Standard deviation of electronic noise (counts)
    
    Returns:
        Spectrum with added electronic noise
    """
    noise = np.random.normal(0, sigma, spectrum.shape)
    result = spectrum + noise
    return np.maximum(0, result)


# =============================================================================
# NORMALIZATION
# =============================================================================

def normalize_spectrum(
    spectrum: np.ndarray,
    method: str = "max"
) -> np.ndarray:
    """
    Normalize a spectrum for ML training.
    
    Args:
        spectrum: Raw count spectrum
        method: Normalization method
            - "max": Divide by maximum value (range 0-1)
            - "sum": Divide by total counts (probability distribution)
            - "log": Log transform then max normalize
            - "sqrt": Square root transform then max normalize
    
    Returns:
        Normalized spectrum
    """
    if method == "max":
        max_val = spectrum.max()
        if max_val > 0:
            return spectrum / max_val
        return spectrum
    
    elif method == "sum":
        total = spectrum.sum()
        if total > 0:
            return spectrum / total
        return spectrum
    
    elif method == "log":
        # Log transform (add 1 to handle zeros)
        log_spec = np.log1p(spectrum)
        max_val = log_spec.max()
        if max_val > 0:
            return log_spec / max_val
        return log_spec
    
    elif method == "sqrt":
        sqrt_spec = np.sqrt(spectrum)
        max_val = sqrt_spec.max()
        if max_val > 0:
            return sqrt_spec / max_val
        return sqrt_spec
    
    else:
        raise ValueError(f"Unknown normalization method: {method}")