numerical_analysis 0.1.2

use core::f64;
use core::f64::consts::PI;

use misc_iterators::IterDimensions;

use crate::EvalResult;
use crate::RealField;
#[cfg(not(target_arch = "spirv"))]
use crate::binomial;
use crate::to_index;

// TODO(medium): try chebyshev polynomials.

#[derive(Default, Clone, Debug)]
pub struct NormalizedMultiDimensionalFourierBasis<
    S: RealField,
    const ARRAY_SIZE: usize = 20,
> {
    pub aabb_min: [S; 3],
    pub aabb_max: [S; 3],
}

impl<S: RealField, const ARRAY_SIZE: usize>
    NormalizedMultiDimensionalFourierBasis<S, ARRAY_SIZE>
{
    pub fn new(aabb_min: [S; 3], aabb_max: [S; 3]) -> Self { Self { aabb_min, aabb_max } }

    pub fn aabb(&self) -> [[S; 3]; 2] { [self.aabb_min, self.aabb_max] }

    // TODO: this logic and the one on `legendre.rs` can be merged, all that changes
    // is the function evaluation and normalization.
    pub fn eval(
        &self,
        x: [S; 3],
        max_basis_index: usize,
        #[cfg(not(feature = "std"))] out: &mut [S; ARRAY_SIZE],
    ) -> EvalResult<S>
    {
        let construct = |i| {
            let pi = S::from(f64::consts::PI).unwrap();
            // Normalize input.
            let a = self.aabb_min[i];
            let b = self.aabb_max[i];
            // map to [-PI, PI]
            let x_p = pi * x[i] * (S::from(2).unwrap() / (b - a)) - ((a + b) / (b - a));

            #[cfg(feature = "std")]
            assert!(x_p >= S::from(-1.0001).unwrap() && x_p <= S::from(1.0001).unwrap());

            #[cfg(not(feature = "std"))]
            let mut res = [S::default(); 6];
            #[cfg(feature = "std")]
            let mut res = vec![S::default(); max_basis_index];

            for j in 0..max_basis_index
            {
                res[j] = fourier_basis(x_p, j) / fourier_norm_on_standard_interval(j);
            }

            res
        };

        let evals = [construct(0), construct(1), construct(2)];

        #[cfg(not(feature = "std"))]
        let mut res = out;
        #[cfg(feature = "std")]
        let mut res = {
            let total_basis_count = binomial(max_basis_index + 3, 3);
            vec![S::from(0.).unwrap(); total_basis_count]
        };

        // Iterate over all integer tuples (i, j, k) such that
        // i + j + k <= max_basis_index.
        for i in 0..max_basis_index + 3
        {
            for j in 1..i
            {
                for k in 0..j
                {
                    let w = i - j - 1;
                    let v = j - k - 1;
                    let u = k;

                    let indices = [u, v, w];
                    let mut prod = S::from(1).unwrap();
                    for l in 0..3
                    {
                        prod *= evals[l][indices[l]];
                    }

                    res[to_index(&indices)] = prod;
                }
            }
        }

        #[cfg(feature = "std")]
        return res;

        #[cfg(not(feature = "std"))]
        []
    }

    #[cfg(feature = "std")]
    pub fn project<F>(&self, max_index: usize, mut function: F) -> Vec<S>
    where
        F: FnMut([S; 3]) -> S,
    {
        let mut res = Vec::new();

        let aabb = self.aabb();
        let aabb = [
            [aabb[0][0], aabb[1][0]],
            [aabb[0][1], aabb[1][1]],
            [aabb[0][2], aabb[1][2]],
        ];
        // Iterate over every basis funciton.
        for i in 2..max_index + 3
        {
            for j in 1..i
            {
                for k in 0..j
                {
                    let w = i - j - 1;
                    let v = j - k - 1;
                    let u = k;

                    let indices = [u, v, w];
                    let coeff = fourier_projection(&mut function, aabb, indices, 100);

                    res.push(coeff);
                }
            }
        }

        res
    }

    pub fn evaluate_function_vector(
        &self,
        p: [S; 3],
        max_index: usize,
        #[cfg(feature = "std")] coefficients: &[S],
        #[cfg(not(feature = "std"))] coefficients: &[S; ARRAY_SIZE],
    ) -> S
    {
        let mut res = S::default();
        let mut index = 0;
        // Iterate over every basis funciton.
        for i in 2..max_index + 3
        {
            for j in 1..i
            {
                for k in 0..j
                {
                    let w = i - j - 1;
                    let v = j - k - 1;
                    let u = k;

                    let indices = [u, v, w];
                    let mut val = coefficients[index];
                    for d in 0..3
                    {
                        let range = self.aabb_max[d] - self.aabb_min[d];
                        let t = (p[d] - self.aabb_min[d]) / range;
                        let xp = t * S::from(2).unwrap() - S::from(1).unwrap();

                        val *= fourier_basis(xp, indices[d])
                            / fourier_norm_on_standard_interval(indices[d]);
                    }

                    res += val;
                    index += 1;
                }
            }
        }

        res
    }
}

/// Defined on [-1, 1].
pub fn fourier_basis<S: RealField>(t: S, basis_index: usize) -> S
{
    let pi = S::from(PI).unwrap();

    let n = basis_index / 2;
    if basis_index % 2 == 0
    {
        ((t * pi) * S::from(n).unwrap()).cos()
    }
    else
    {
        ((t * pi) * S::from(n + 1).unwrap()).sin()
    }
}

pub fn fourier_norm_on_standard_interval<S: RealField>(basis_index: usize) -> S
{
    // Only for the very first function (constant) is the norm different.
    if basis_index == 0
    {
        S::from(core::f64::consts::SQRT_2).unwrap()
    }
    else
    {
        S::from(1).unwrap()
    }
}

pub fn fourier_roots<S: RealField>(period: S, basis_index: usize) -> [S; 2]
{
    let n = basis_index / 2;
    if basis_index % 2 == 0
    {
        let v = period / S::from(n).unwrap();
        [-v, v]
    }
    else
    {
        [S::default(), period / S::from(2 * n).unwrap()]
    }
}

fn fourier_projection<F, S: RealField, const DIM: usize>(
    mut function: F,
    interval: [[S; 2]; DIM],
    basis_index: [usize; DIM],
    resolution: usize,
) -> S
where
    F: FnMut([S; DIM]) -> S,
{
    let trapezoidal_weight = |i: usize| {
        if i == 0 || i == resolution - 1
        {
            S::from(0.5).unwrap()
        }
        else
        {
            S::from(1).unwrap()
        }
    };

    let mut res = S::default();
    for xi in IterDimensions::new([resolution; DIM])
    {
        let mut point = [S::default(); DIM];
        for d in 0..DIM
        {
            let range = interval[d][1] - interval[d][0];
            let t = S::from(xi[d]).unwrap() / S::from(resolution - 1).unwrap();
            let x = t * range + interval[d][0];

            point[d] = x;
        }

        let mut val = function(point);
        for d in 0..DIM
        {
            let range = interval[d][1] - interval[d][0];
            let dx = range / S::from(resolution - 1).unwrap();
            let t = S::from(xi[d]).unwrap() / S::from(resolution - 1).unwrap();
            let xp = t * S::from(2).unwrap() - S::from(1).unwrap();

            val *= fourier_basis(xp, basis_index[d]) * trapezoidal_weight(xi[d]) * dx
                / fourier_norm_on_standard_interval(basis_index[d]);
        }

        res += val
    }

    res
}

#[cfg(test)]
mod tests
{
    use core::f64::consts::PI;

    use super::*;

    #[test]
    fn test_fourier_quadrature()
    {
        let test = fourier_projection(|_: [f64; 1]| 1., [[-1., 1.]], [0], 100);
        assert!((test - 2_f64.sqrt()).abs() < 0.001, "{}", test);
        let test = fourier_projection(|[x]| (PI * x).cos(), [[-1., 1.]], [2], 100);
        assert!((test - 1.).abs() < 0.001, "{}", test);
        let test = fourier_projection(|[x, y]| 1., [[-1., 1.], [-1., 1.]], [0, 0], 100);
        assert!((test - 2_f64).abs() < 0.001, "{}", test);
        let test = fourier_projection(
            |[x, y, z]| (PI * x).cos(),
            [[-1., 1.]; 3],
            [2, 0, 0],
            100,
        );
        assert!((test - 2_f64).abs() < 0.001, "{}", test);
        let test = fourier_projection(
            |[x, y, z]| f64::sqrt(x * x + y * y + z * z) - 0.5,
            [[-1., 1.]; 3],
            [2, 2, 0],
            100,
        );
        assert!((test - -0.0750718).abs() < 0.001, "{}", test);

        // let basis =
        //     NormalizedMultiDimensionalFourierBasis::<_, 20>::new([-1.; 3],
        // [1.; 3]);

        // let sdf = |v: [f64; 3]| (v[0] * v[0] + v[1] * v[1] + v[2] *
        // v[2]).sqrt() - 0.5;

        // // let sdf = |v: [f64; 3]| 5.;

        // let coeffs = basis.project(4, sdf);

        // println!("{:.2?}", coeffs);

        // let x = basis.evaluate_function_vector([0., 0.5, 0.], 4, &coeffs);
        // println!("{:?}", x);
        // let x = basis.evaluate_function_vector([0., 1., 0.], 4, &coeffs);
        // println!("{:?}", x);
        // let x = basis.evaluate_function_vector([0., 0., 0.], 4, &coeffs);
        // println!("{:?}", x);
    }
}