sva-samples 0.8.1

// Concern: where on the grid each atom of a lane can be nonzero, and a lane summed over only those | Non-concern: one atom's value (point.rs), choosing the row | IO: (&Lane, span) -> intervals, a plane

use sva_formula::spectral_sum::atom::{Gauss, SpectralAtom};
use sva_formula::{Lane, exp_zero_at};

use super::point;
use crate::error::CollapseError;
use crate::grid::Grid;

pub(crate) type SampleInterval = (i64, i64);

pub(crate) const OPEN: SampleInterval = (i64::MIN, i64::MAX);

const REACH: i64 = 1 << 62;

const FINITE: f64 = 1e300;

/// Outside it `smooth_at` is exactly zero: its indicator, or a factor the engine's own `exp`
/// underflows while every other factor stays finite.
pub(crate) fn nonzero_interval(a: &SpectralAtom, grid: Grid) -> SampleInterval {
    if a.is_delta() {
        return OPEN;
    }
    let mut held = OPEN;
    if let Some(ind) = a.ind {
        held = meet(
            held,
            (grid.first_at(ind.l.value()), grid.first_at(ind.r.value())),
        );
    }
    if a.pole.is_some() || !reaches_finite(a, grid) {
        return held;
    }
    let zero = exp_zero_at();
    let exp = a.exp.filter(|e| e.sigma != 0.0);
    if let Some(e) = exp {
        let dead = |n: i64| e.sigma * (grid.instant(n) - e.mu) <= zero;
        held = meet(
            held,
            match e.sigma < 0.0 {
                true => (i64::MIN, end(dead)),
                false => (start(|n| !dead(n)), i64::MAX),
            },
        );
    }
    if let Some(g) = a.gauss {
        let grows = |n: i64, rising: bool| {
            exp.is_some_and(|e| {
                (e.sigma > 0.0) == rising || e.sigma * (grid.instant(n) - e.mu) > 700.0
            })
        };
        held = meet(held, gaussian(g, grid, zero, &grows));
    }
    held
}

fn gaussian(g: Gauss, grid: Grid, zero: f64, grows: &dyn Fn(i64, bool) -> bool) -> SampleInterval {
    let dead = |n: i64| {
        let x = grid.instant(n);
        -g.a * (x - g.mu) * (x - g.mu) <= zero
    };
    let centre = start(|n| grid.instant(n) >= g.mu).clamp(-REACH, REACH);
    let right = first_in(centre, REACH, dead).filter(|n| !grows(*n, true));
    let left =
        first_in(-REACH, centre, |n| !dead(n)).filter(|n| *n > -REACH && !grows(n - 1, false));
    (left.unwrap_or(i64::MIN), right.unwrap_or(i64::MAX))
}

fn reaches_finite(a: &SpectralAtom, grid: Grid) -> bool {
    let farthest = grid.instant(i64::MAX);
    (a.c.re.abs() + a.c.im.abs()) * (farthest + a.poly.at.abs()).powi(i32::from(a.poly.degree))
        < FINITE
}

/// The indices whose instant, increasing in the index, lies in `[l, r)`.
pub(crate) fn between(l: f64, r: f64, at: impl Fn(i64) -> f64) -> SampleInterval {
    (start(|n| l <= at(n)), end(|n| r <= at(n)))
}

pub(crate) fn meet(a: SampleInterval, b: SampleInterval) -> SampleInterval {
    let held = (a.0.max(b.0), a.1.min(b.1));
    match held.0 < held.1 {
        true => held,
        false => (0, 0),
    }
}

fn first_in(lo: i64, hi: i64, holds: impl Fn(i64) -> bool) -> Option<i64> {
    if lo > hi || !holds(hi) {
        return None;
    }
    let (mut no, mut yes) = (lo, hi);
    if holds(lo) {
        return Some(lo);
    }
    while i128::from(yes) - i128::from(no) > 1 {
        let mid = ((i128::from(no) + i128::from(yes)) / 2) as i64;
        match holds(mid) {
            true => yes = mid,
            false => no = mid,
        }
    }
    Some(yes)
}

fn start(holds: impl Fn(i64) -> bool) -> i64 {
    match first_in(-REACH, REACH, holds) {
        Some(n) if n == -REACH => i64::MIN,
        Some(n) => n,
        None => REACH,
    }
}

fn end(holds: impl Fn(i64) -> bool) -> i64 {
    match first_in(-REACH, REACH, holds) {
        Some(n) if n == -REACH => i64::MIN,
        Some(n) => n,
        None => i64::MAX,
    }
}

pub(crate) fn nonzero_intervals(lane: &Lane, grid: Grid) -> Vec<SampleInterval> {
    lane.atoms
        .iter()
        .map(|a| nonzero_interval(a, grid))
        .collect()
}

/// Samples `[from, to)` of one lane into `out`, each summed in lane order over the atoms
/// live there: one left out adds an exact zero to a sum that starts at +0.
pub(crate) fn sweep(
    lane: &Lane,
    intervals: &[SampleInterval],
    span: SampleInterval,
    grid: Grid,
    out: &mut [f64],
) -> Result<(), CollapseError> {
    sweep_by(intervals, span, out, |m, live| {
        Ok(point::eval_among(lane, live, (grid, m))?.re)
    })
}

/// Each sample `m` of `[from, to)` into `out` as `each(m, live)`, `live` the indices of the
/// intervals holding `m`, ascending: each index is visited only inside its own interval.
pub(crate) fn sweep_by(
    intervals: &[SampleInterval],
    (from, to): SampleInterval,
    out: &mut [f64],
    mut each: impl FnMut(i64, &[usize]) -> Result<f64, CollapseError>,
) -> Result<(), CollapseError> {
    let mut order: Vec<usize> = (0..intervals.len())
        .filter(|&i| {
            intervals[i].0 < intervals[i].1 && intervals[i].0 < to && intervals[i].1 > from
        })
        .collect();
    order.sort_by_key(|&i| intervals[i].0);
    let (mut live, mut next, mut n) = (Vec::<usize>::new(), 0, from);
    while n < to {
        while let Some(&i) = order.get(next).filter(|&&i| intervals[i].0 <= n) {
            let slot = live.partition_point(|&j| j < i);
            live.insert(slot, i);
            next += 1;
        }
        live.retain(|&i| intervals[i].1 > n);
        let stop = live
            .iter()
            .map(|&i| intervals[i].1)
            .chain(order.get(next).map(|&i| intervals[i].0))
            .fold(to, i64::min);
        for m in n..stop {
            out[(m - from) as usize] = each(m, &live)?;
        }
        n = stop;
    }
    Ok(())
}

/// Where every atom of a lane is windowed, the lane is zero outside their union.
pub(crate) fn spans(lane: &Lane, grid: Grid) -> Option<Vec<(i64, i64)>> {
    if !lane.is_finite_sum() || lane.atoms.iter().any(|a| a.ind.is_none()) {
        return None;
    }
    let mut spans: Vec<(i64, i64)> = lane
        .atoms
        .iter()
        .filter_map(|a| {
            let interval = a.ind?;
            let (from, to) = (
                grid.first_at(interval.l.value()),
                grid.first_at(interval.r.value()),
            );
            (from < to).then_some((from, to))
        })
        .collect();
    spans.sort_unstable();
    let mut merged: Vec<(i64, i64)> = Vec::with_capacity(spans.len());
    for (from, to) in spans {
        match merged.last_mut() {
            Some(held) if from <= held.1 => held.1 = held.1.max(to),
            _ => merged.push((from, to)),
        }
    }
    Some(merged)
}