"""JAX-jitted line profile integration kernels.

All functions compute the fraction of a normalized profile that falls within
each wavelength bin ``[low, high]``.  They are pure JAX with no numpyro
dependency and are designed to be called from within :func:`jax.jit`-compiled
model code.
"""

from __future__ import annotations

from typing import Final

import jax.numpy as jnp
from jax import Array, config, jit
from jax.scipy.special import erf, erfc
from jax.typing import ArrayLike

# Conversion: FWHM to sigma for the half-variance parametrization of erf.
# sigma = FWHM / (2 sqrt(2 ln 2)); erf uses sqrt(2)*sigma, so the factor is:
_HALFVAR_SIGMA_TO_FWHM: Final[Array] = 2 * jnp.sqrt(jnp.log(2))

# Pseudo-Voigt polynomial coefficients from Thompson et al. (1987)
# DOI: 10.1107/S0021889887087090
_VOIGT_FWHM_CS: Final[Array] = jnp.array([1, 2.69268, 2.42843, 4.47163, 0.07842, 1])
_VOIGT_ETA_CS: Final[Array] = jnp.array([1.33603, -0.47719, 0.11116])

# Threshold for when exp(x*x) overflows
_OVERFLOW_THRESHOLD: Final[float] = 26.0 if config.jax_enable_x64 else 9.0

# Conversion factor from exponential (Laplace) scale to FWHM
# pdf = (1/(2*b)) * exp(-|x - μ|/b)
# max(pdf) = 1/(2*b), half max = 1/(4*b)
# 1/(4*b) = 1/(2*b) * exp(-|x - μ|/b) => exp(-|x - μ|/b) = 1/2
# => |x - μ|/b = ln(2) => FWHM = 2*b*ln(2)
_EXP_SCALE_TO_FWHM: Final[Array] = 2 * jnp.log(2)

# Precompute constants
_INV_SQRT2PI: Final[Array] = 1.0 / jnp.sqrt(2 * jnp.pi)
_SQRT6: Final[Array] = jnp.sqrt(6.0)
_SQRT24: Final[Array] = jnp.sqrt(24.0)


@jit
def integrate_gaussian(
    low: ArrayLike,
    high: ArrayLike,
    center: ArrayLike,
    lsf_fwhm: ArrayLike,
    fwhm: ArrayLike,
) -> Array:
    """Integrate a Gaussian profile over wavelength bins.

    Parameters
    ----------
    low : ArrayLike
        Lower bin edges.
    high : ArrayLike
        Upper bin edges.
    center : ArrayLike
        Line center wavelength.
    lsf_fwhm : ArrayLike
        Instrumental line spread function FWHM at the line center.
    fwhm : ArrayLike
        Intrinsic Gaussian FWHM.

    Returns
    -------
    jnp.ndarray
        Integrated fraction per bin.
    """
    # Add LSF in quadrature
    total_fwhm = jnp.sqrt(lsf_fwhm**2 + fwhm**2)
    inv_sigma = _HALFVAR_SIGMA_TO_FWHM / total_fwhm
    t_low = (low - center) * inv_sigma
    t_high = (high - center) * inv_sigma
    return (erf(t_high) - erf(t_low)) / 2


def _integrate_cauchy(
    low: ArrayLike,
    high: ArrayLike,
    center: ArrayLike,
    lsf_fwhm: ArrayLike,
    fwhm: ArrayLike,
) -> Array:
    """Private: Integrate a pure Cauchy (Lorentzian) profile over wavelength bins.

    This function is kept for reference but is no longer used directly.
    Cauchy profiles are now implemented as PseudoVoigt with LSF=0.

    Parameters
    ----------
    low : ArrayLike
        Lower bin edges.
    high : ArrayLike
        Upper bin edges.
    center : ArrayLike
        Line center wavelength.
    lsf_fwhm : ArrayLike
        Instrumental line spread function FWHM at the line center.
    fwhm : ArrayLike
        Intrinsic Lorentzian FWHM.

    Returns
    -------
    jnp.ndarray
        Integrated fraction per bin.
    """
    # For Cauchy, we don't convolve with LSF - it's a pure Lorentzian
    # This function is kept for reference but no longer used
    inv_hwhm = 2 / fwhm
    t_low = (low - center) * inv_hwhm
    t_high = (high - center) * inv_hwhm
    return (jnp.arctan(t_high) - jnp.arctan(t_low)) / jnp.pi


@jit
def integrate_voigt(
    low: ArrayLike,
    high: ArrayLike,
    center: ArrayLike,
    lsf_fwhm: ArrayLike,
    fwhm_g: ArrayLike,
    fwhm_l: ArrayLike,
) -> Array:
    """Integrate a pseudo-Voigt profile over wavelength bins.

    Uses the Thompson et al. (1987) approximation: a weighted sum of
    Gaussian and Lorentzian components with an effective FWHM computed
    from the individual Gaussian and Lorentzian fwhms.

    Parameters
    ----------
    low : ArrayLike
        Lower bin edges.
    high : ArrayLike
        Upper bin edges.
    center : ArrayLike
        Line center wavelength.
    lsf_fwhm : ArrayLike
        Instrumental line spread function FWHM at the line center.
    fwhm_g : ArrayLike
        Intrinsic Gaussian component FWHM.
    fwhm_l : ArrayLike
        Lorentzian component FWHM.

    Returns
    -------
    jnp.ndarray
        Integrated fraction per bin.

    References
    ----------
    Thompson, Cox & Hastings (1987), J. Appl. Cryst. 20, 79-83.
    DOI: 10.1107/S0021889887087090
    """
    # Add LSF in quadrature to Gaussian component only
    fwhm_g_tot = jnp.sqrt(lsf_fwhm**2 + fwhm_g**2)

    # Effective FWHM from the 5th-order polynomial
    pows = jnp.arange(_VOIGT_FWHM_CS.size)
    fwhm = jnp.sum(_VOIGT_FWHM_CS * (fwhm_g_tot**pows) * (fwhm_l ** pows[::-1])) ** (
        1 / 5
    )

    # Lorentzian mixing fraction η
    fwhm_ratio = fwhm_l / fwhm
    eta = jnp.sum(_VOIGT_ETA_CS * (fwhm_ratio ** jnp.arange(1, len(_VOIGT_ETA_CS) + 1)))

    lorentzian = eta * _integrate_cauchy(low, high, center, lsf_fwhm, fwhm)
    gaussian = (1 - eta) * integrate_gaussian(low, high, center, lsf_fwhm, fwhm)
    return lorentzian + gaussian


# Don't need this since we always convolve with LSF
# But keeping anyways just in case
@jit
def _integrate_laplace(
    low: ArrayLike,
    high: ArrayLike,
    center: ArrayLike,
    lsf_fwhm: ArrayLike,
    fwhm: ArrayLike,
) -> Array:
    """Integrate a Laplace (double-exponential) profile over wavelength bins.

    The LSF is **not** convolved --- this profile is a pure Laplace
    distribution.

    Parameters
    ----------
    low : ArrayLike
        Lower bin edges.
    high : ArrayLike
        Upper bin edges.
    center : ArrayLike
        Line center wavelength.
    lsf_fwhm : ArrayLike
        Instrumental line spread function FWHM at the line center.
    fwhm : ArrayLike
        Laplace FWHM (related to scale parameter *b* by
        ``FWHM = 2 b ln 2``).

    Returns
    -------
    jnp.ndarray
        Integrated fraction per bin.
    """
    # For Laplace, we don't convolve with LSF - it's a pure Laplace distribution
    b = fwhm / _EXP_SCALE_TO_FWHM

    def _cdf(x):
        t = (x - center) / b
        return 0.5 + 0.5 * jnp.sign(t) * (1.0 - jnp.exp(-jnp.abs(t)))

    return _cdf(high) - _cdf(low)


def _integrandGL(t: ArrayLike, a: ArrayLike) -> Array:
    """
    Compute exponential correction integrand for Gaussian-Laplace convolution.

    Evaluates exp(-a²) * [exp(2ta) * erfcx(a+t) - exp(-2ta) * erfcx(a-t)].
    Exploits odd symmetry to cover overflow at large a + t.

    Parameters
    ----------
    t : jnp.ndarray
        Normalized distance from center
    a : jnp.ndarray
        Convolution parameter (σλ/2)

    Returns
    -------
    jnp.ndarray
        Integrand value with numerical stability
    """
    # Exploit odd symmetry: I(-t,a) = -I(t,a)
    # Then we only need one overflow protection
    t_abs = jnp.abs(t)

    # Shorthand terms
    ta = t_abs + a
    twota = 2 * t_abs * a

    # Overflow protection
    posterm = jnp.where(ta > _OVERFLOW_THRESHOLD, 0, jnp.exp(twota) * erfc(ta))

    return jnp.sign(t) * jnp.exp(a * a) * (posterm - jnp.exp(-twota) * erfc(-t_abs + a))


@jit
def integrate_gaussianLaplace(
    low: ArrayLike,
    high: ArrayLike,
    center: ArrayLike,
    lsf_fwhm: ArrayLike,
    fwhm_g: ArrayLike,
    fwhm_l: ArrayLike,
) -> Array:
    """
    Integrate exponentially modified Gaussian (EMG) emission lines over wavelength bins.

    Computes the integral of Gaussian convolved with symmetric double exponential
    (Laplace distribution). This profile combines instrumental broadening (Gaussian)
    with natural broadening (Laplacian). Uses asymptotic approximations for
    numerical stability.

    Parameters
    ----------
    low : jnp.ndarray
        Low wavelength edges of bins
    high : jnp.ndarray
        High wavelength edges of bins
    center : jnp.ndarray
        Line centers
    lsf_fwhm : jnp.ndarray
        Instrumental LSF FWHM
    fwhm_g : jnp.ndarray
        Intrinsic Gaussian component FWHM
    fwhm_l : jnp.ndarray
        Laplacian component FWHM (natural broadening)

    Returns
    -------
    jnp.ndarray
        Integrated flux in each bin for each line
    """
    # This function breaks for pure Laplace, does that matter? lsf is never zero?

    # Add LSF in quadrature to Gaussian component
    fwhm_g_total = jnp.sqrt(lsf_fwhm**2 + fwhm_g**2)

    # Convert FWHM to distribution parameters
    σ = fwhm_g_total / _HALFVAR_SIGMA_TO_FWHM
    λ = _EXP_SCALE_TO_FWHM / fwhm_l

    # Normalized distance from center
    scale = 1 / σ
    t_low = (low - center) * scale
    t_high = (high - center) * scale

    # Convolution parameter
    a = σ * λ / 2

    # Gaussian component (via error function CDF)
    gaussian_cdf = (erf(t_high) - erf(t_low)) / 2

    # Exponential correction
    exp_correction = _integrandGL(t_high, a) - _integrandGL(t_low, a)

    # Guard against large a, the Gaussian limit
    return gaussian_cdf + jnp.where(a > _OVERFLOW_THRESHOLD, 0, exp_correction / 4)


def _integrandGH(t_halfvar: ArrayLike, h3_eff: ArrayLike, h4_eff: ArrayLike) -> Array:
    """
    Antiderivative of the Gauss-Hermite correction at halfvar coordinate t_halfvar.

    Converts to the standard coordinate ``y = t_halfvar * sqrt(2)``, then evaluates
    ``-g(y) * [h3_eff * He_2(y) + h4_eff * He_3(y)]``, the analytic antiderivative
    of ``g(y) * [h3_eff * He_3(y) + h4_eff * He_4(y)]`` w.r.t. x, where
    g(y) = exp(-y²/2) is the standard normal kernel and He_n are the probabilists'
    Hermite polynomials: He_2(y) = y²-1, He_3(y) = y(y²-3).

    Parameters
    ----------
    t_halfvar : ArrayLike
        Halfvar-normalized coordinate ``(x - center) / sigma_tot_halfvar``.
    h3_eff : ArrayLike
        Effective h3 coefficient: ``h3 * (sigma_g / sigma_tot)**3 / sqrt(6)``.
    h4_eff : ArrayLike
        Effective h4 coefficient: ``h4 * (sigma_g / sigma_tot)**4 / sqrt(24)``.

    Returns
    -------
    jnp.ndarray
        Antiderivative value.
    """
    t = t_halfvar * jnp.sqrt(2.0)
    g = jnp.exp(-t * t / 2)
    # He_2(y) = y²-1,  He_3(y) = y(y²-3)
    return g * (h3_eff * (t * t - 1) + h4_eff * t * (t * t - 3))


@jit
def integrate_split_normal(
    low: ArrayLike,
    high: ArrayLike,
    center: ArrayLike,
    lsf_fwhm: ArrayLike,
    fwhm_blue: ArrayLike,
    fwhm_red: ArrayLike,
) -> Array:
    """
    Integrate a split-normal (two-sided Gaussian) profile over wavelength bins.

    The split-normal distribution has different standard deviations on each side
    of the mean. The left side (blue, shorter wavelengths) uses fwhm_blue, and
    the right side (red, longer wavelengths) uses fwhm_red.

    Parameters
    ----------
    low : ArrayLike
        Lower bin edges.
    high : ArrayLike
        Upper bin edges.
    center : ArrayLike
        Line center wavelength.
    lsf_fwhm : ArrayLike
        Instrumental line spread function FWHM at the line center.
    fwhm_blue : ArrayLike
        Blue side (left) Gaussian FWHM.
    fwhm_red : ArrayLike
        Red side (right) Gaussian FWHM.

    Returns
    -------
    jnp.ndarray
        Integrated fraction per bin.
    """
    # Add LSF in quadrature to both components
    total_fwhm_blue = jnp.sqrt(lsf_fwhm**2 + fwhm_blue**2)
    total_fwhm_red = jnp.sqrt(lsf_fwhm**2 + fwhm_red**2)

    # Convert FWHM to sigma for the half-variance parametrization of erf
    inv_sigma_blue = _HALFVAR_SIGMA_TO_FWHM / total_fwhm_blue
    inv_sigma_red = _HALFVAR_SIGMA_TO_FWHM / total_fwhm_red

    # Calculate normalization factor to ensure the total integral is 1
    # The split-normal PDF is proportional to:
    #   exp(-(x-center)^2/(2*sigma_blue^2)) for x <= center
    #   exp(-(x-center)^2/(2*sigma_red^2)) for x > center
    # The normalization constant is: 1 / [sqrt(2*pi) * (sigma_blue + sigma_red) / 2]
    # But since we're working with CDFs, we need a different approach

    # Normalization factors to ensure proper weighting
    # The weights should be proportional to the inverse of the sigmas
    total_weight = inv_sigma_blue + inv_sigma_red
    norm_blue = inv_sigma_blue / total_weight
    norm_red = inv_sigma_red / total_weight

    def _split_normal_cdf(
        x, center, inv_sigma_blue, inv_sigma_red, norm_blue, norm_red
    ):
        # For x <= center, use blue side
        # For x > center, use red side
        t_blue = (x - center) * inv_sigma_blue
        t_red = (x - center) * inv_sigma_red

        # CDF for each side
        cdf_blue = (1 + erf(t_blue)) / 2
        cdf_red = (1 + erf(t_red)) / 2

        # Combine them with proper weighting to ensure normalization to unity
        # The blue side contributes norm_blue * cdf_blue
        # The red side contributes norm_blue + norm_red * cdf_red (to ensure continuity)
        return jnp.where(
            x <= center, norm_blue * cdf_blue, norm_blue + norm_red * cdf_red
        )

    # Calculate CDF at high and low edges
    cdf_high = _split_normal_cdf(
        high, center, inv_sigma_blue, inv_sigma_red, norm_blue, norm_red
    )
    cdf_low = _split_normal_cdf(
        low, center, inv_sigma_blue, inv_sigma_red, norm_blue, norm_red
    )

    return cdf_high - cdf_low


@jit
def integrate_gaussHermite(
    low: ArrayLike,
    high: ArrayLike,
    center: ArrayLike,
    fwhm_lsf: ArrayLike,
    fwhm_g: ArrayLike,
    h3: ArrayLike,
    h4: ArrayLike,
) -> Array:
    """
    Integrate Gauss-Hermite emission lines over wavelength bins.

    Parameters
    ----------
    low : jnp.ndarray
        Low wavelength edges of bins
    high : jnp.ndarray
        High wavelength edges of bins
    center : jnp.ndarray
        Line centers
    fwhm_lsf : jnp.ndarray
        Instrumental LSF FWHM
    fwhm_g : jnp.ndarray
        Gaussian component FWHM
    h3 : jnp.ndarray
        Gauss-Hermite h3 coefficient
    h4 : jnp.ndarray
        Gauss-Hermite h4 coefficient

    Returns
    -------
    jnp.ndarray
        Integrated flux in each bin for each line
    """
    # Effective FWHM from convolution of LSF and Gaussian component
    fwhm_tot = jnp.sqrt(fwhm_lsf**2 + fwhm_g**2)

    # Convert FWHM to sigma
    inv_sigma_tot = _HALFVAR_SIGMA_TO_FWHM / fwhm_tot

    # Sigma ratio: GH moments scale as r^n under convolution with a Gaussian LSF
    r = fwhm_g / fwhm_tot
    r3 = r * r * r

    # Normalized bin edges
    t_low = (low - center) * inv_sigma_tot
    t_high = (high - center) * inv_sigma_tot

    # Gaussian CDF contribution
    gaussian_cdf = (erf(t_high) - erf(t_low)) / 2

    # GH coefficients scaled by sigma ratio (moment theorem for Gaussian convolution)
    c3 = h3 * r3 / _SQRT6
    c4 = h4 * r3 * r / _SQRT24

    # GH correction via analytic antiderivative
    gh_correction = _integrandGH(t_high, c3, c4) - _integrandGH(t_low, c3, c4)

    return gaussian_cdf - _INV_SQRT2PI * gh_correction


# _integrate_single_line and _integrate_lines live in unite.line.profiles,
# where each Profile subclass owns its JAX branch via jax_branch().
