numerical_analysis 0.1.2

use misc_iterators::MultiIndex;
use misc_iterators::raster_line::DDA;
use nalgebra as na;

type Vec2 = na::Vector2<f32>;
type IVec2 = na::Vector2<i32>;

pub fn fast_sweeping_parametric_curve<const N: usize>(
    corner: &Vec2,
    extent: &Vec2,
    dx: f32,
    f: &dyn Fn(&Vec2) -> f32,
    curve: &dyn Fn(f32) -> Vec2,
    range: (f32, f32),
    resolution: usize,
) -> Vec<f32>
{
    let mut dimensions = [0; N];
    for (i, dim) in dimensions.iter_mut().enumerate()
    {
        *dim = (extent[i] / dx) as i32;
    }

    let grid_size = dimensions.iter().fold(1, |prod, factor| prod * factor);
    let mut grid = vec![f32::MAX; grid_size as usize];

    rasterize_parametric_curve(&mut grid, corner, extent, dx, curve, range, resolution);

    fast_sweeping(
        &[corner.x, corner.y],
        &[extent.x, extent.y],
        dx,
        &|p: &[f32; 2]| f(&Vec2::new(p[0], p[1])),
        &mut grid,
    );

    grid
}

pub fn rasterize_parametric_curve(
    grid: &mut [f32],
    corner: &Vec2,
    extent: &Vec2,
    dx: f32,
    curve: &dyn Fn(f32) -> Vec2,
    range: (f32, f32),
    resolution: usize,
)
{
    let stretch = range.1 - range.0;
    let mut points = Vec::new();
    for i in 0..resolution
    {
        let t = i as f32 / (resolution - 1) as f32;
        let t = t * stretch + range.0;
        points.push(curve(t));
    }

    for i in 0..points.len() - 1
    {
        let point = points[i];

        let relative_pos = point - corner;
        let box_corner = relative_pos.map(|x| x - (x % dx)) + corner;
        let index = relative_pos.map(|x| x.div_euclid(dx));

        let resolution: IVec2 = na::try_convert(extent / dx).unwrap();

        let dir = points[i + 1] - points[i];
        let distance = dir.norm();
        let dir = dir / distance;

        let mut line_iter = DDA::new(
            [point.x, point.y],
            [box_corner.x, box_corner.y],
            [dir.x, dir.y],
            [index.x as isize, index.y as isize],
            dx,
        );

        loop
        {
            let (t, [j, k]) = line_iter.next().unwrap();
            if t >= distance
            {
                break;
            }

            grid[(k * resolution.x as isize + j) as usize] =
                distance_to_segment(&(point + dir * t), &point, &points[i + 1]);
        }
    }
}

fn project_point_to_segment(point: &Vec2, origin: &Vec2, end: &Vec2) -> Vec2
{
    let len = (end - origin).norm();
    let dir = (end - origin) / len;
    let v = point - origin;

    let proj_distance = v.dot(&dir).clamp(0.0, len);

    origin + dir * proj_distance
}

fn distance_to_segment(point: &Vec2, origin: &Vec2, end: &Vec2) -> f32
{
    let proj = project_point_to_segment(point, origin, end);
    (point - proj).norm()
}

// Implementation of:
// https://graphics.stanford.edu/courses/cs468-03-fall/Papers/zhao_fastsweep1.pdf
// Also read:
// https://math.berkeley.edu/~sethian/2006/Papers/sethian.vonkarman_1.pdf
/// Implementation of Hongkai Zhao's fast-sweeping algorithm. Computes a
/// distance field for an initialised grid of starting values.
///
/// The grid is assumed to be "infinite" almost everywhere, except on the
/// boundary of the wavefront to be expanded, where it is assumed to be
/// initialised to the corresponding distance of the cell centre to the
/// boundary.
///
/// `dx` is the sidelength of a grid cell.
/// `f` is a function that returns the speed of the wavefront at a given point;
/// almost  universally, this will be `|_| 1.0`.
pub fn fast_sweeping<const N: usize>(
    corner: &[f32; N],
    extent: &[f32; N],
    dx: f32,
    f: &dyn Fn(&[f32; N]) -> f32,
    grid: &mut [f32],
)
{
    let mut dimensions = [0; N];
    for (i, dim) in dimensions.iter_mut().enumerate()
    {
        *dim = (extent[i] / dx) as usize;
    }

    let coords_to_position = |coords: &[usize; N]| {
        let mut position = [0.; N];
        for i in 0..N
        {
            position[i] = coords[i] as f32 * dx + corner[i];
        }

        position
    };

    for _ in 0..2
    {
        for step_dirs in MultiIndex::new([-1; N], [3; N], [2; N])
        {
            sweep(
                grid,
                &dimensions,
                dx,
                f,
                &coords_to_position,
                &step_dirs.map(|i| i as i32),
            );
        }
    }
}

fn check_neighbours<const N: usize>(
    grid: &mut [f32],
    grid_dimensions: &[usize],
    dx: f32,
    coords: &[usize; N],
    coords_to_position: &dyn Fn(&[usize; N]) -> [f32; N],
    f: &dyn Fn(&[f32; N]) -> f32,
)
{
    // j * grid_dimensions[1] as usize + i
    let grid_index = |index: &[usize; N]| {
        let mut grid_index = 0;
        for (dim, i) in index.iter().rev().enumerate()
        {
            grid_index *= grid_dimensions[N - 1 - dim];
            grid_index += *i;
        }
        grid_index
    };

    // TODO(low): Change to a slice once generic const expressions become a thing.
    let mut us = vec![f32::MAX; N + 1];
    for i in 0..N
    {
        let mut offset_coord = coords.clone();
        offset_coord[i] = (coords[i] + 1).min(grid_dimensions[i] - 1);
        let u_next = grid[grid_index(&offset_coord)];

        let mut offset_coord = coords.clone();
        offset_coord[i] = (coords[i] as isize - 1).max(0) as usize;
        let u_prev = grid[grid_index(&offset_coord)];

        us[i] = u_next.min(u_prev).max(0.);
    }

    let pos = coords_to_position(&coords);

    grid[grid_index(coords)] =
        grid[grid_index(coords)].min(fast_sweep_quadratic_solve(&mut us, f(&pos) * dx));
}

fn sweep<const N: usize>(
    grid: &mut [f32],
    grid_dimensions: &[usize],
    dx: f32,
    f: &dyn Fn(&[f32; N]) -> f32,
    coords_to_position: &dyn Fn(&[usize; N]) -> [f32; N],
    step_directions: &[i32; N],
)
{
    // Initialize the start and end indices of the current sweep for each dimension.
    let mut iteration_start_bounds = [0; N];
    let mut iteration_end_bounds = [0; N];
    for (dim, &step_dir) in step_directions.iter().enumerate()
    {
        if step_dir > 0
        {
            iteration_start_bounds[dim] = 0;
            iteration_end_bounds[dim] = grid_dimensions[dim] as isize - 1;
        }
        else
        {
            iteration_start_bounds[dim] = grid_dimensions[dim] as isize - 1;
            iteration_end_bounds[dim] = 0;
        };
    }

    for multi_index in MultiIndex::new(
        iteration_start_bounds,
        iteration_end_bounds,
        step_directions.clone(),
    )
    {
        check_neighbours(
            grid,
            grid_dimensions,
            dx,
            &multi_index.map(|i| i as usize),
            coords_to_position,
            f,
        );
    }
}

// equation 2.6 in Hongkai Zhao's paper:
// https://graphics.stanford.edu/courses/cs468-03-fall/Papers/zhao_fastsweep1.pdf
fn fast_sweep_quadratic_solve(a_coeffs: &mut [f32], fh: f32) -> f32
{
    // Fast sort.
    a_coeffs.sort_unstable_by(|a, b| a.total_cmp(b));
    // Special case if the solution is determined by the first coefficient.
    // Sw Zaho's paper.
    if a_coeffs[0] + fh < a_coeffs[1]
    {
        return a_coeffs[0] + fh;
    }
    // We are solving a sequence of quadratic equations of the form:
    // $ sum (x - a_i)^2 = (fh)^2 $ so our terms are:
    // $ a = p x^2 $
    // $ b = -2(sum a_i) $
    // $ c = (sum a_i^2) - fh $
    // Refer to Zhao's paper for details.
    for p in 2..a_coeffs.len()
    {
        let sum = a_coeffs[..p].iter().sum::<f32>();
        let sum_squares = a_coeffs[..p].iter().fold(0.0, |sum, x| sum + x * x);

        // Coefficients of the quadratic with index p.
        let a = p as f32;
        let b = -2.0 * sum;
        let c = sum_squares - fh * fh;

        let det = b * b - 4.0 * a * c;

        let s1 = -b - det.sqrt();
        let s2 = -b + det.sqrt();
        // Pick largest root.
        let solution = s1.max(s2) / (2.0 * a);

        if p == a_coeffs.len() - 1 || solution <= a_coeffs[p + 1]
        {
            return solution;
        }
    }

    panic!("Something on the SDF quadratic solver is broken.");
}