leibniz 0.2.0

use ::burn::tensor::{Int, Tensor, TensorData};

use super::differentiate;
use crate::{
    base::{Config, geometry::Command},
    burn::{
        geometry::{Contour, contour},
        tests::Backend,
    },
};

#[test]
fn differentiates_contours() {
    let gradients = differentiate(
        &[square([0.0, 0.0], [1.0, 1.0])],
        image([[2.0]]),
        Config::default(),
    );

    assert_eq!(gradients.len(), 1);
    assert_nonzero(gradients[0].clone());
}

#[test]
fn gradient_scale_is_invariant_to_sample_count() {
    // The corrected normalization (weight_image division plus 1 / (height *
    // width)) makes the boundary gradient an unbiased estimate that does not
    // grow with the antialiasing sample count. The un-normalized version scaled
    // with x_sample_count * y_sample_count * height * width.
    let contours = [square([0.7, 0.7], [3.3, 3.3])];
    let signal = image([
        [1.0, 2.0, 3.0, 4.0],
        [2.0, 3.0, 4.0, 5.0],
        [3.0, 4.0, 5.0, 6.0],
        [4.0, 5.0, 6.0, 7.0],
    ]);
    let coarse = magnitude(differentiate(
        &contours,
        signal.clone(),
        Config::new(2, 2, 0.5),
    ));
    let fine = magnitude(differentiate(&contours, signal, Config::new(4, 4, 0.5)));

    assert!(coarse > 0.0 && fine > 0.0);
    assert!(
        (coarse / fine - 1.0).abs() < 0.25,
        "coarse {coarse}, fine {fine}"
    );
}

#[test]
#[should_panic]
fn rejects_empty_contours() {
    let _ = differentiate::<Backend>(&[], image([[2.0]]), Config::default());
}

#[test]
#[should_panic]
fn rejects_invalid_config() {
    let _ = differentiate(
        &[square([0.0, 0.0], [1.0, 1.0])],
        image([[2.0]]),
        Config::new(0, 2, 0.5),
    );
}

#[test]
#[should_panic]
fn rejects_invalid_contour_shape() {
    let _ = differentiate(
        &[Contour::new(
            vec![Command::Quadratic],
            Tensor::<Backend, 3>::from_data(
                TensorData::new(vec![0.0; 6], [1, 3, 2]),
                &Default::default(),
            ),
        )],
        image([[2.0]]),
        Config::default(),
    );
}

#[test]
#[should_panic]
fn rejects_invalid_signal() {
    let _ = differentiate(
        &[square([0.0, 0.0], [1.0, 1.0])],
        Tensor::<Backend, 2>::from_data(
            TensorData::new(Vec::<f32>::new(), [0, 1]),
            &Default::default(),
        ),
        Config::default(),
    );
}

pub fn assert_close(actual: f32, expected: f32) {
    assert!((actual - expected).abs() < 1e-5);
}

pub fn assert_floats<const N: usize>(tensor: Tensor<Backend, 1>, expected: [f32; N]) {
    let actual = tensor.into_data().to_vec::<f32>().unwrap();

    assert_eq!(actual.len(), expected.len());
    for (actual, expected) in actual.into_iter().zip(expected) {
        assert_close(actual, expected);
    }
}

pub fn assert_ints<const N: usize>(tensor: Tensor<Backend, 1, Int>, expected: [i64; N]) {
    let actual = tensor.into_data().to_vec::<i64>().unwrap();

    assert_eq!(actual, expected);
}

pub fn assert_matrix<const ROWS: usize, const COLUMNS: usize>(
    tensor: Tensor<Backend, 2>,
    expected: [[f32; COLUMNS]; ROWS],
) {
    let actual = tensor.into_data().to_vec::<f32>().unwrap();
    let expected = expected.into_iter().flatten().collect::<Vec<_>>();

    assert_eq!(actual.len(), expected.len());
    for (actual, expected) in actual.into_iter().zip(expected) {
        assert_close(actual, expected);
    }
}

pub fn assert_nonzero<const D: usize>(tensor: Tensor<Backend, D>) {
    let actual = tensor.into_data().to_vec::<f32>().unwrap();

    assert!(actual.iter().any(|value| value.abs() > 1e-6));
}

pub fn image<const HEIGHT: usize, const WIDTH: usize>(
    values: [[f32; WIDTH]; HEIGHT],
) -> Tensor<Backend, 2> {
    Tensor::<Backend, 2>::from_data(TensorData::from(values), &Default::default())
}

pub fn one_segment() -> Contour<Backend> {
    segments([[[0.0, 0.0], [1.0, 0.0]]])
}

pub fn samples<const N: usize>(values: [f32; N]) -> Tensor<Backend, 1> {
    Tensor::<Backend, 1>::from_data(TensorData::from(values), &Default::default())
}

pub fn segments<const N: usize>(values: [[[f32; 2]; 2]; N]) -> Contour<Backend> {
    Contour::new(
        vec![Command::Quadratic; N],
        Tensor::<Backend, 3>::from_data(TensorData::from(values), &Default::default()),
    )
}

pub fn small_triangle() -> Contour<Backend> {
    segments([
        [[0.0, 0.0], [0.75, 0.0]],
        [[1.5, 0.0], [1.5, 1.0]],
        [[1.5, 2.0], [0.75, 1.0]],
    ])
}

pub fn square(min: [f32; 2], max: [f32; 2]) -> Contour<Backend> {
    Contour::new(
        vec![Command::Linear; 4],
        Tensor::<Backend, 3>::from_data(
            TensorData::from(square_segments(
                [min[0], min[1]],
                [max[0], min[1]],
                [max[0], max[1]],
                [min[0], max[1]],
            )),
            &Default::default(),
        ),
    )
}

pub fn triangle() -> Contour<Backend> {
    segments([
        [[0.0, 0.0], [1.5, 0.0]],
        [[3.0, 0.0], [3.0, 2.0]],
        [[3.0, 4.0], [1.5, 2.0]],
    ])
}

fn interpolate(start: [f32; 2], end: [f32; 2], t: f32) -> [f32; 2] {
    [
        start[0] + (end[0] - start[0]) * t,
        start[1] + (end[1] - start[1]) * t,
    ]
}

fn magnitude(gradients: Vec<contour::Arguments<Backend>>) -> f32 {
    gradients
        .iter()
        .flat_map(|gradient| gradient.clone().into_data().to_vec::<f32>().unwrap())
        .map(|value| value.abs())
        .sum()
}

fn square_segments(a: [f32; 2], b: [f32; 2], c: [f32; 2], d: [f32; 2]) -> [[[f32; 2]; 2]; 4] {
    [
        [a, interpolate(a, b, 0.5)],
        [b, interpolate(b, c, 0.5)],
        [c, interpolate(c, d, 0.5)],
        [d, interpolate(d, a, 0.5)],
    ]
}