use crate::*;

pub trait Kernel {
    fn convolve<P>(&self, img: &image::Image<P>) -> image::Image<P>
    where
        P: image::Pixel;
}

fn sum<const N: usize>(kernel: &[[f32; N]; N]) -> f32 {
    let mut x = 0.0;
    for i in 0..N {
        for j in 0..N {
            x += kernel[j][i];
        }
    }
    x
}

impl<const N: usize> Kernel for [[f32; N]; N] {
    fn convolve<P>(&self, img: &image::Image<P>) -> image::Image<P>
    where
        P: image::Pixel,
    {
        let l = N / 2;
        let mag = sum(self);

        let extent = image::Extent {
            x: img.extent.x - 2 * l,
            y: img.extent.y - 2 * l,
        };

        let mut pixels = Vec::with_capacity(extent.size() as usize);
        for y in 0..extent.y {
            for x in 0..extent.x {
                let mut pixel = P::default();
                for i in 0..N {
                    for j in 0..N {
                        let img_x = x as usize + i;
                        let img_y = y as usize + j;
                        let idx = img_y * img.extent.x as usize + img_x;
                        pixel += img.pixels[idx] * self[j][i];
                    }
                }
                pixels.push(pixel / mag);
            }
        }

        image::Image { extent, pixels }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn size_of_result_img() {
        let input = image::Image {
            extent: image::Extent { x: 6, y: 10 },
            pixels: vec![0.0; 60],
        };
        let kernel = [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]];
        let output = kernel.convolve(&input);
        assert_eq!(output.extent.x, 4);
        assert_eq!(output.extent.y, 8);
    }

    #[test]
    fn single_pixel_result() {
        let input = image::Image {
            extent: image::Extent { x: 3, y: 3 },
            pixels: vec![
                image::SRGBA::new(0.1, 0.0, 0.0, 1.0),
                image::SRGBA::new(0.0, 0.2, 0.0, 1.0),
                image::SRGBA::new(0.0, 0.0, 0.3, 1.0),
                image::SRGBA::new(0.2, 0.0, 0.0, 1.0),
                image::SRGBA::new(0.0, 0.3, 0.0, 1.0),
                image::SRGBA::new(0.0, 0.0, 0.1, 1.0),
                image::SRGBA::new(0.3, 0.0, 0.0, 1.0),
                image::SRGBA::new(0.0, 0.1, 0.0, 1.0),
                image::SRGBA::new(0.0, 0.0, 0.2, 1.0),
            ],
        };
        let kernel = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]];
        let output = kernel.convolve(&input);
        assert_eq!(output.extent.x, 1);
        assert_eq!(output.extent.y, 1);
        assert_eq!(
            output.pixels[0].r,
            (0.1 * 0.1 + 0.2 * 0.4 + 0.3 * 0.7) / 4.5,
        );
        assert_eq!(
            output.pixels[0].g,
            (0.2 * 0.2 + 0.3 * 0.5 + 0.1 * 0.8) / 4.5,
        );
        assert_eq!(
            output.pixels[0].b,
            (0.3 * 0.3 + 0.1 * 0.6 + 0.2 * 0.9) / 4.5,
        );
        assert_eq!(output.pixels[0].a, 1.0);
    }
}