"""
Theoretical natural background spectrum for CsI(Tl) detectors (Radiacode 103).

Shape calibrated against real Radiacode 103 background measurements.
The CsI(Tl) crystal (1 cm³, 8.4% FWHM) produces a spectrum with:
- A dominant low-energy hump peaking around 100-120 keV
- Exponential decay at higher energies
- Subtle photopeaks from natural isotopes
"""

import numpy as np
from app.config import ENERGY_OFFSET, ENERGY_SLOPE, NUM_CHANNELS


# Photopeak lines: (energy_keV, relative_weight)
# Weights tuned so peaks are visible above local continuum at typical CPS
NATURAL_BG_LINES = [
    (295.22, 0.10),   # Pb-214
    (351.93, 0.18),   # Pb-214
    (609.31, 0.15),   # Bi-214
    (911.20, 0.08),   # Ac-228
    (968.97, 0.05),   # Ac-228
    (1120.29, 0.06),  # Bi-214
    (1460.83, 0.12),  # K-40
    (1764.49, 0.08),  # Bi-214
    (2614.51, 0.18),  # Tl-208
]


def _gaussian(x, center, sigma, amplitude):
    return amplitude * np.exp(-0.5 * ((x - center) / sigma) ** 2)


def generate_theoretical_bg(cps: float = 6.0, live_time_s: float = 3600.0):
    channels = np.arange(NUM_CHANNELS, dtype=np.float64)
    energy_axis = ENERGY_OFFSET + ENERGY_SLOPE * channels
    total_counts = cps * live_time_s

    # ── 1. Main hump: asymmetric peak at ~105 keV ──
    # Real data: rises from ~60 at 10keV to ~280 at 100-120keV, then falls
    hump_center = 110.0
    hump = np.zeros(NUM_CHANNELS, dtype=np.float64)
    low_mask = energy_axis <= hump_center
    hump[low_mask] = _gaussian(energy_axis[low_mask], hump_center, 55.0, 1.0)
    hump[~low_mask] = _gaussian(energy_axis[~low_mask], hump_center, 50.0, 1.0)

    # ── 2. Compton continuum tail ──
    # Real data: ~136@200, ~80@250, ~44@295, ~14@400, ~5@600
    tail = 0.45 * np.exp(-energy_axis / 240) + 0.04 * np.exp(-energy_axis / 700)

    # ── 3. Low-energy noise floor ──
    noise_floor = 0.008

    # ── 4. Combine continuum ──
    continuum = hump + tail + noise_floor

    # ── 5. Photopeaks ──
    # CsI(Tl) 8.4% FWHM at 662 keV, scaling as sqrt(E)
    # sigma(E) = FWHM(E) / 2.355 = 0.084 * sqrt(E * 662) / 662 / 2.355
    # Simplified: sigma = 23.6 * sqrt(E/662) keV
    def sigma_keV(E):
        return max(12.0, 23.6 * np.sqrt(max(E, 1.0) / 662.0))

    peak_frac = 0.08  # 8% of total counts in resolved photopeaks
    total_weight = sum(w for _, w in NATURAL_BG_LINES)

    peaks = np.zeros(NUM_CHANNELS, dtype=np.float64)
    for line_energy, weight in NATURAL_BG_LINES:
        sig = sigma_keV(line_energy)
        peak_counts = total_counts * peak_frac * (weight / total_weight)
        amplitude = peak_counts / (sig * np.sqrt(2 * np.pi))
        peaks += _gaussian(energy_axis, line_energy, sig, amplitude)

    # ── 6. Combine and normalize ──
    raw = continuum + peaks / total_counts  # peaks normalized later
    raw *= total_counts / raw.sum()

    # ── 7. Poisson-like noise ──
    rng = np.random.default_rng(42)
    noise = rng.normal(0, 1, NUM_CHANNELS) * np.sqrt(np.maximum(raw, 1.0)) * 0.25
    raw += noise

    # Floor at 0.9 for log scale
    spectrum = np.clip(raw, 0.9, None)

    key_lines = [
        (295.22, "Pb-214"), (351.93, "Pb-214"),
        (609.31, "Bi-214"), (911.20, "Ac-228"),
        (1120.29, "Bi-214"), (1460.83, "K-40"),
        (1764.49, "Bi-214"), (2614.51, "Tl-208"),
    ]

    return {
        "energy_kev": [round(float(E), 2) for E in energy_axis],
        "counts": [round(float(c), 1) for c in spectrum],
        "cps": round(cps, 2),
        "live_time_s": round(live_time_s, 1),
        "lines": [
            {"energy_keV": E, "name": name} for E, name in key_lines
        ],
    }


def generate_continuum_only(cps: float = 6.0, live_time_s: float = 3600.0):
    """Generate only the CsI(Tl) continuum shape (no photopeaks, no noise).

    This matches the model used in training (generate_realistic_continuum in
    spectrum_physics.py) for direct comparison with measured backgrounds.
    """
    channels = np.arange(NUM_CHANNELS, dtype=np.float64)
    energy_axis = ENERGY_OFFSET + ENERGY_SLOPE * channels
    total_counts = cps * live_time_s

    # Asymmetric hump at ~110 keV
    hump_center = 110.0
    hump = np.where(
        energy_axis <= hump_center,
        np.exp(-0.5 * ((energy_axis - hump_center) / 55.0) ** 2),
        np.exp(-0.5 * ((energy_axis - hump_center) / 50.0) ** 2),
    )

    # Compton continuum tail
    tail = 0.45 * np.exp(-energy_axis / 240.0) + 0.04 * np.exp(-energy_axis / 700.0)

    # Noise floor
    noise_floor = 0.008

    continuum = hump + tail + noise_floor

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

    return {
        "energy_kev": [round(float(E), 2) for E in energy_axis],
        "counts": [round(float(c), 1) for c in continuum],
        "cps": round(cps, 2),
        "live_time_s": round(live_time_s, 1),
    }