"""
Detector-agnostic continuum extraction for gamma-ray spectra.

Extracts the detector's intrinsic response curve (continuum) from measured
background data. Isotope photopeaks are subtracted, then the residual is
smoothed to produce a clean continuum shape that reflects only the detector's
physics — no isotope signatures.

Works with any scintillator or semiconductor detector.
"""

import numpy as np
from scipy.ndimage import gaussian_filter1d


# Common environmental isotope lines (keV) — subtracted regardless of detector.
_ENV_PEAKS = [
    (241.0, 0.04),
    (295.22, 0.1842),
    (351.93, 0.3560),
    (609.31, 0.4549),
    (911.20, 0.2580),
    (1120.29, 0.1492),
    (1460.83, 0.1066),
    (1764.49, 0.1531),
    (2614.51, 0.3586),
]

_E_OFFSET = 0.33
_E_SLOPE = 2.97


def _sigma_ch(E_keV):
    """Peak sigma in channels at energy E_keV (sqrt(E) resolution scaling)."""
    fwhm_keV = 0.08 * E_keV * (E_keV / 662.0) ** 0.5
    sigma_keV = fwhm_keV / 2.355
    return max(sigma_keV / _E_SLOPE, 2.0)


def _subtract_peaks(counts, energy_axis):
    """Remove known isotope photopeaks from spectrum."""
    continuum = counts.copy()
    channels = np.arange(len(counts), dtype=np.float64)

    for line_energy, _ in _ENV_PEAKS:
        idx = int(np.argmin(np.abs(energy_axis - line_energy)))
        if idx < 0 or idx >= len(counts):
            continue

        sig = _sigma_ch(line_energy)
        far = int(5 * sig) + 3

        lo_start = max(0, idx - far - int(3 * sig))
        lo_end = max(0, idx - far)
        hi_start = min(len(counts), idx + far)
        hi_end = min(len(counts), idx + far + int(3 * sig))

        baseline_regions = []
        if lo_end > lo_start:
            baseline_regions.append(continuum[lo_start:lo_end])
        if hi_end > hi_start:
            baseline_regions.append(continuum[hi_start:hi_end])

        if not baseline_regions:
            continue

        local_bg = float(np.median(np.concatenate(baseline_regions)))
        peak_height = continuum[idx] - local_bg

        if peak_height > 0:
            gaussian = peak_height * np.exp(-0.5 * ((channels - idx) / sig) ** 2)
            continuum -= gaussian

    return np.maximum(continuum, 0)


def extract_continuum(counts, energy_axis=None):
    """Extract the detector's intrinsic response continuum.

    Removes isotope photopeaks, then smooths with a wide Gaussian filter
    to produce a clean curve showing only the detector's continuum shape.

    Parameters
    ----------
    counts : array
        Raw accumulated counts per channel.
    energy_axis : array, optional
        Energy axis in keV.

    Returns
    -------
    array — smooth continuum (peak-subtracted, Gaussian-smoothed)
    """
    counts = np.asarray(counts, dtype=np.float64)
    n_channels = len(counts)

    if energy_axis is None:
        energy_axis = _E_OFFSET + _E_SLOPE * np.arange(n_channels, dtype=np.float64)

    continuum = _subtract_peaks(counts, energy_axis)

    # Wide Gaussian smooth (sigma ~1.5% of channels ≈ 45 keV)
    sigma = max(15, n_channels // 60)
    continuum_smooth = gaussian_filter1d(continuum, sigma=sigma)
    continuum_smooth = np.maximum(continuum_smooth, 0)

    return continuum_smooth