//! Some structs and methods to make working with registration and interpolation methods in
//! SimpleITK more Rust friendly.

use crate::simple;
use anyhow::{Result, anyhow};
use autocxx::prelude::*;
use cxx::{CxxVector, UniquePtr, let_cxx_string};
use ndarray::{Array2, ArrayView2, AsArray, Ix2, array, s};
use serde::{Deserialize, Serialize};
use serde_yaml::{from_reader, to_writer};
use std::fs::File;
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut, Mul};
use std::path::PathBuf;
use num::Complex;
use tempfile::tempdir;

/// a trait marking number types that can be used in sitk:
/// (u/i)(8/16/32/64), (u/i)size, f(32/64)
pub trait PixelType: Clone {
    const PT: simple::PixelIDValueEnum;
}

macro_rules! pixel_type_impl {
    ($($T:ty: $sitk:expr $(,)?)*) => {
        $(
            impl PixelType for $T {
                const PT: simple::PixelIDValueEnum = $sitk;
            }
        )*
    };
}

pixel_type_impl! {
    u8: simple::PixelIDValueEnum::sitkUInt8,
    i8: simple::PixelIDValueEnum::sitkInt8,
    u16: simple::PixelIDValueEnum::sitkUInt16,
    i16: simple::PixelIDValueEnum::sitkInt16,
    u32: simple::PixelIDValueEnum::sitkUInt32,
    i32: simple::PixelIDValueEnum::sitkInt32,
    u64: simple::PixelIDValueEnum::sitkUInt64,
    i64: simple::PixelIDValueEnum::sitkInt64,
    f32: simple::PixelIDValueEnum::sitkFloat32,
    f64: simple::PixelIDValueEnum::sitkFloat64,
}

#[cfg(target_pointer_width = "64")]
pixel_type_impl!(usize: simple::PixelIDValueEnum::sitkUInt64);
#[cfg(target_pointer_width = "32")]
pixel_type_impl!(usize: simple::PixelIDValueEnum::sitkUInt32);
#[cfg(target_pointer_width = "64")]
pixel_type_impl!(isize: simple::PixelIDValueEnum::sitkInt64);
#[cfg(target_pointer_width = "32")]
pixel_type_impl!(isize: simple::PixelIDValueEnum::sitkInt32);

/// Struct holding a pointer to an image
pub struct Image<T: PixelType> {
    image: UniquePtr<simple::Image>,
    pixel_type: PhantomData<T>,
}

impl<T: PixelType> Deref for Image<T> {
    type Target = UniquePtr<simple::Image>;

    fn deref(&self) -> &Self::Target {
        &self.image
    }
}

impl<T: PixelType> Image<T> {
    /// encapsulate an itk::simple::Image
    pub fn new(image: UniquePtr<simple::Image>) -> Self {
        Self {
            image,
            pixel_type: PhantomData,
        }
    }

    /// take an ndarray Array2 and turn it into a SimpleITK image
    pub fn from_array<'a, A>(array: A) -> Self
    where
        T: 'a + PixelType,
        A: AsArray<'a, T, Ix2>,
    {
        let array = array.into();
        let shape = array.shape();
        let width = (shape[1] as u32).into();
        let height = (shape[0] as u32).into();
        let mut image = simple::Image::new3(width, height, T::PT).within_unique_ptr();
        image.pin_mut().MakeUnique();
        let buffer = image.pin_mut().GetBufferAsVoid();
        unsafe { std::ptr::copy(array.as_ptr(), buffer as *mut T, shape[0] * shape[1]) };
        Self {
            image,
            pixel_type: PhantomData,
        }
    }

    /// return as an ndarray Array2
    pub fn as_array(&self) -> Array2<T> {
        let width = u32::from(self.image.GetWidth()) as usize;
        let height = u32::from(self.image.GetHeight()) as usize;
        let mut array = Array2::<T>::uninit((height, width));
        let buffer = self.image.GetBufferAsVoid1();
        unsafe {
            std::ptr::copy(
                buffer as *const T,
                array.as_mut_ptr() as *mut T,
                width * height,
            );
            array.assume_init()
        }
    }
}

impl<'a, A, T> From<A> for Image<T>
where
    T: 'a + PixelType,
    A: AsArray<'a, T, Ix2>,
{
    fn from(value: A) -> Self {
        Self::from_array(value.into())
    }
}

impl<T> From<Image<T>> for Array2<T>
where
    T: PixelType,
{
    fn from(value: Image<T>) -> Self {
        value.as_array()
    }
}

/// a struct describing the transform
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct AffineTransform {
    /// flattened 2x2 rotation matrix + translation
    pub parameters: [f64; 6],
    /// error / significance on parameters
    pub dparameters: [f64; 6],
    /// the point about which rotations are performed
    pub origin: [f64; 2],
    /// the shape of images for which this transform is meant
    pub shape: [usize; 2],
}

impl Mul for AffineTransform {
    type Output = AffineTransform;

    #[allow(clippy::suspicious_arithmetic_impl)]
    fn mul(self, other: AffineTransform) -> AffineTransform {
        let m = self.matrix().dot(&other.matrix());
        let dm = self.dmatrix().dot(&other.matrix()) + self.matrix().dot(&other.dmatrix());
        AffineTransform {
            parameters: [
                m[[0, 0]],
                m[[0, 1]],
                m[[1, 0]],
                m[[1, 1]],
                m[[2, 0]],
                m[[2, 1]],
            ],
            dparameters: [
                dm[[0, 0]],
                dm[[0, 1]],
                dm[[1, 0]],
                dm[[1, 1]],
                dm[[2, 0]],
                dm[[2, 1]],
            ],
            origin: self.origin,
            shape: self.shape,
        }
    }
}

impl Eq for AffineTransform {}

impl AffineTransform {
    /// parameters: flat 2x2 part of matrix, translation; origin: center of rotation
    pub fn new(parameters: [f64; 6], origin: [f64; 2], shape: [usize; 2]) -> Self {
        Self {
            parameters,
            dparameters: [0f64; 6],
            origin,
            shape,
        }
    }

    /// find the affine transform which transforms moving into fixed
    pub fn register_affine<F, M, T>(fixed: F, moving: M) -> Result<AffineTransform>
    where
        F: Into<Image<T>>,
        M: Into<Image<T>>,
        T: PixelType,
    {
        Self::register(fixed, moving, true)
    }

    /// find the translation which transforms moving into fixed
    pub fn register_translation<F, M, T>(fixed: F, moving: M) -> Result<AffineTransform>
    where
        F: Into<Image<T>>,
        M: Into<Image<T>>,
        T: PixelType,
    {
        Self::register(fixed, moving, false)
    }

    /// find the transform which transforms moving into fixed
    pub fn register<F, M, T>(fixed: F, moving: M, affine: bool) -> Result<AffineTransform>
    where
        F: Into<Image<T>>,
        M: Into<Image<T>>,
        T: PixelType,
    {
        let tmp_folder = tempdir()?;
        let fixed = fixed.into();
        let moving = moving.into();
        let width = u32::from(fixed.GetWidth()) as usize;
        let height = u32::from(fixed.GetHeight()) as usize;
        let_cxx_string!(transform_name = if affine { "affine" } else { "translation" });
        let parameter_map = crate::ffi_extra::get_default_parameter_map(&transform_name);
        let mut tfilter = simple::ElastixImageFilter::new().within_box();
        tfilter.as_mut().LogToConsoleOff();
        tfilter.as_mut().LogToFileOff();
        tfilter.as_mut().SetLogToFile(false);
        tfilter.as_mut().SetFixedImage(&fixed);
        tfilter.as_mut().SetMovingImage(&moving);
        crate::ffi_extra::set_parameter_map(&mut tfilter, &parameter_map);
        tfilter.as_mut().SetParameter("WriteResultImage", "False");
        tfilter
            .as_mut()
            .SetOutputDirectory(tmp_folder.path().display().to_string());
        let _ = tfilter.as_mut().Execute().within_unique_ptr();
        let_cxx_string!(tp = "TransformParameters");
        let p = crate::ffi_extra::get_transform_parameter_map(tfilter.as_mut().deref_mut(), 0)
            .get(&tp)
            .iter()
            .map(|i| i.to_string_lossy().parse::<f64>())
            .collect::<Result<Vec<f64>, _>>()?;
        let parameters = if affine {
            [p[0], p[1], p[2], p[3], p[4], p[5]]
        } else {
            [1.0, 0.0, 0.0, 1.0, p[0], p[1]]
        };
        let origin = [((height - 1) as f64) / 2.0, ((width - 1) as f64) / 2.0];
        let shape = [height, width];
        Ok(AffineTransform::new(parameters, origin, shape))
    }

    /// create a transform from a xy translation
    pub fn from_translation(translation: [f64; 2]) -> Self {
        AffineTransform {
            parameters: [1f64, 0f64, 0f64, 1f64, translation[0], translation[1]],
            dparameters: [0f64; 6],
            origin: [0f64; 2],
            shape: [0usize; 2],
        }
    }

    /// read a transform from a file
    pub fn from_file(path: PathBuf) -> Result<Self> {
        let file = File::open(path)?;
        Ok(from_reader(file)?)
    }

    /// write a transform to a file
    pub fn to_file(&self, path: PathBuf) -> Result<()> {
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(path)?;
        to_writer(&mut file, self)?;
        Ok(())
    }

    /// true if transform does nothing
    pub fn is_unity(&self) -> bool {
        self.parameters == [1f64, 0f64, 0f64, 1f64, 0f64, 0f64]
    }

    /// transform an image using nearest neighbor interpolation
    pub fn transform_image_bspline<I, T>(&self, image: I) -> Result<Image<T>>
    where
        I: Into<Image<T>>,
        T: PixelType,
    {
        Self::interp(self, image, false)
    }

    /// transform an image using bspline interpolation
    pub fn transform_image_nearest_neighbor<I, T>(&self, image: I) -> Result<Image<T>>
    where
        I: Into<Image<T>>,
        T: PixelType,
    {
        Self::interp(self, image, true)
    }

    /// transform an image
    pub fn interp<I, T>(&self, image: I, nearest_neighbor: bool) -> Result<Image<T>>
    where
        I: Into<Image<T>>,
        T: PixelType,
    {
        let image = image.into();
        let width = u32::from(image.GetWidth()) as usize;
        let height = u32::from(image.GetHeight()) as usize;
        let origin = [((width - 1) as f64) / 2f64, ((height - 1) as f64) / 2f64];
        let p = self.parameters;
        let matrix = cxx_vector([p[0], p[1], p[2], p[3]]);
        let translation = cxx_vector([p[4], p[5]]);
        let fixed_center = cxx_vector(origin);
        let affine_transform =
            simple::AffineTransform::new3(&matrix, &translation, &fixed_center).within_unique_ptr();
        let transform = <_ as AsRef<simple::Transform>>::as_ref(affine_transform.as_ref().unwrap());
        let interpolator = if nearest_neighbor {
            simple::InterpolatorEnum::sitkBSpline
        } else {
            simple::InterpolatorEnum::sitkNearestNeighbor
        };
        Ok(Image::<T>::new(
            simple::Resample(
                &image,
                transform,
                interpolator,
                0.0,
                simple::PixelIDValueEnum::sitkUnknown,
                nearest_neighbor,
            )
            .within_unique_ptr(),
        ))
    }

    /// get coordinates resulting from transforming input coordinates, coordinates must have two
    /// columns: x & y
    pub fn transform_coordinates<'a, A, T>(&self, coordinates: A) -> Result<Array2<f64>>
    where
        T: 'a + Clone + Into<f64>,
        A: AsArray<'a, T, Ix2>,
    {
        let coordinates = coordinates.into();
        let s = coordinates.shape();
        if s[1] != 2 {
            return Err(anyhow!("coordinates must have two columns"));
        }
        let m = self.matrix();
        let mut res = Array2::zeros([s[0], s[1]]);
        for i in 0..s[0] {
            let a = array![
                coordinates[[i, 0]].clone().into(),
                coordinates[[i, 1]].clone().into(),
                1f64
            ]
            .to_owned();
            let b = m.dot(&a);
            res.slice_mut(s![i, ..]).assign(&b.slice(s![..2]));
        }
        Ok(res)
    }

    /// get the matrix defining the transform
    pub fn matrix(&self) -> Array2<f64> {
        Array2::from_shape_vec(
            (3, 3),
            vec![
                self.parameters[0],
                self.parameters[1],
                self.parameters[4],
                self.parameters[2],
                self.parameters[3],
                self.parameters[5],
                0f64,
                0f64,
                1f64,
            ],
        )
        .unwrap()
    }

    /// get the matrix describing the error of the transform
    pub fn dmatrix(&self) -> Array2<f64> {
        Array2::from_shape_vec(
            (3, 3),
            vec![
                self.dparameters[0],
                self.dparameters[1],
                self.dparameters[4],
                self.dparameters[2],
                self.dparameters[3],
                self.dparameters[5],
                0f64,
                0f64,
                1f64,
            ],
        )
        .unwrap()
    }

    /// get the inverse transform
    pub fn inverse(&self) -> Result<AffineTransform> {
        fn det(a: ArrayView2<f64>) -> f64 {
            (a[[0, 0]] * a[[1, 1]]) - (a[[0, 1]] * a[[1, 0]])
        }

        let m = self.matrix();
        let d = det(m.slice(s![..2, ..2]));
        if d == 0f64 {
            return Err(anyhow!("transform matrix is not invertible"));
        }
        let parameters = [
            det(m.slice(s![1.., 1..])) / d,
            -det(m.slice(s![..;2, 1..])) / d,
            -det(m.slice(s![1.., ..;2])) / d,
            det(m.slice(s![..;2, ..;2])) / d,
            det(m.slice(s![..2, 1..])) / d,
            -det(m.slice(s![..2, ..;2])) / d,
        ];

        Ok(AffineTransform {
            parameters,
            dparameters: [0f64; 6],
            origin: self.origin,
            shape: self.shape,
        })
    }

    /// adapt the transform to a new origin and shape
    pub fn adapt(&mut self, origin: [f64; 2], shape: [usize; 2]) {
        self.origin = [
            origin[0] + (((self.shape[0] - shape[0]) as f64) / 2f64),
            origin[1] + (((self.shape[1] - shape[1]) as f64) / 2f64),
        ];
        self.shape = shape;
    }
}

/// conveniently collect an iterator into a CxxVector
pub fn cxx_vector<T, I>(vec: I) -> UniquePtr<CxxVector<T>>
where
    I: IntoIterator<Item = T>,
    T: cxx::vector::VectorElement + cxx::ExternType<Kind = cxx::kind::Trivial>,
{
    let mut v = CxxVector::new();
    v.pin_mut().extend(vec);
    v
}

/// An example of generating julia fractals, for testing purposes.
pub fn julia_image(shift_x: f32, shift_y: f32) -> Result<Array2<u8>> {
    let imgx = 800;
    let imgy = 600;

    let scalex = 3.0 / imgx as f32;
    let scaley = 3.0 / imgy as f32;

    let mut im = Array2::<u8>::zeros((imgy, imgx));
    for x in 0..imgx {
        for y in 0..imgy {
            let cy = (y as f32 + shift_y) * scalex - 1.5;
            let cx = (x as f32 + shift_x) * scaley - 1.5;

            let c = Complex::new(-0.4, 0.6);
            let mut z = Complex::new(cy, cx);

            let mut i = 0;
            while i < 255 && z.norm() <= 2.0 {
                z = z * z + c;
                i += 1;
            }

            im[[y, x]] = i as u8;
        }
    }
    Ok(im)
}


#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Result;
    use ndarray::Array2;
    use tempfile::NamedTempFile;

    #[test]
    fn test_serialization() -> Result<()> {
        let file = NamedTempFile::new()?;
        let t = AffineTransform::new([1.2, 0.3, -0.4, 0.9, 10.2, -9.5], [59.5, 49.5], [120, 100]);
        t.to_file(file.path().to_path_buf())?;
        let s = AffineTransform::from_file(file.path().to_path_buf())?;
        assert_eq!(s, t);
        Ok(())
    }

    macro_rules! interp_tests_bspline {
        ($($name:ident: $t:ty $(,)?)*) => {
            $(
                #[test]
                fn $name() -> Result<()> {
                    let j = julia_image(-120f32, 10f32)?.mapv(|x| x as $t);
                    let k = julia_image(0f32, 0f32)?.mapv(|x| x as $t);
                    let shape = j.shape();
                    let origin = [
                        ((shape[1] - 1) as f64) / 2f64,
                        ((shape[0] - 1) as f64) / 2f64,
                    ];
                    let transform = AffineTransform::new([1., 0., 0., 1., 120., -10.], origin, [shape[0], shape[1]]);
                    let n: Array2<_> = transform.transform_image_bspline(j.view())?.into();
                    let d = (k.mapv(|x| x as f64) - n.mapv(|x| x as f64)).powi(2).sum();
                    assert!(d <= (shape[0] * shape[1]) as f64);
                    Ok(())
                }
            )*
        }
    }

    interp_tests_bspline! {
        interpbs_u8: u8,
        interpbs_i8: i8,
        interpbs_u16: u16,
        interpbs_i16: i16,
        interpbs_u32: u32,
        interpbs_i32: i32,
        interpbs_u64: u64,
        interpbs_i64: i64,
        interpbs_f32: f32,
        interpbs_f64: f64,
    }

    macro_rules! interp_tests_nearest_neighbor {
        ($($name:ident: $t:ty $(,)?)*) => {
            $(
                #[test]
                fn $name() -> Result<()> {
                    let j = julia_image(-120f32, 10f32)?.mapv(|x| x as $t);
                    let k = julia_image(0f32, 0f32)?.mapv(|x| x as $t);
                    let shape = j.shape();
                    let origin = [
                        ((shape[1] - 1) as f64) / 2f64,
                        ((shape[0] - 1) as f64) / 2f64,
                    ];
                    let j0 = j.clone();
                    let k0 = k.clone();
                    let transform = AffineTransform::new([1., 0., 0., 1., 120., -10.], origin, [shape[0], shape[1]]);
                    // make sure j & k weren't mutated
                    assert!(j.iter().zip(j0.iter()).map(|(a, b)| a == b).all(|x| x));
                    assert!(k.iter().zip(k0.iter()).map(|(a, b)| a == b).all(|x| x));
                    let n: Array2<_> = transform.transform_image_nearest_neighbor(j.view())?.into();
                    let d = (k.mapv(|x| x as f64) - n.mapv(|x| x as f64)).powi(2).sum();
                    assert!(d <= (shape[0] * shape[1]) as f64);
                    Ok(())
                }
            )*
        }
    }

    interp_tests_nearest_neighbor! {
        interpnn_u8: u8,
        interpnn_i8: i8,
        interpnn_u16: u16,
        interpnn_i16: i16,
        interpnn_u32: u32,
        interpnn_i32: i32,
        interpnn_u64: u64,
        interpnn_i64: i64,
        interpnn_f32: f32,
        interpnn_f64: f64,
    }

    macro_rules! registration_tests_translation {
        ($($name:ident: $t:ty $(,)?)*) => {
            $(
                #[test]
                fn $name() -> Result<()> {
                    let j = julia_image(0f32, 0f32)?.mapv(|x| x as $t);
                    let k = julia_image(10f32, 20f32)?.mapv(|x| x as $t);
                    let j0 = j.clone();
                    let k0 = k.clone();
                    let t = AffineTransform::register_translation(j.view(), k.view())?;
                    // make sure j & k weren't mutated
                    assert!(j.iter().zip(j0.iter()).map(|(a, b)| a == b).all(|x| x));
                    assert!(k.iter().zip(k0.iter()).map(|(a, b)| a == b).all(|x| x));
                    let mut m = Array2::eye(3);
                    m[[0, 2]] = -10f64;
                    m[[1, 2]] = -20f64;
                    let d = (t.matrix() - m).powi(2).sum();
                    assert!(d < 0.01, "d: {}, t: {:?}", d, t.parameters);
                    Ok(())
                }
            )*
        }
    }

    registration_tests_translation! {
        registration_translation_u8: u8,
        registration_translation_i8: i8,
        registration_translation_u16: u16,
        registration_translation_i16: i16,
        registration_translation_u32: u32,
        registration_translation_i32: i32,
        registration_translation_u64: u64,
        registration_translation_i64: i64,
        registration_translation_f32: f32,
        registration_translation_f64: f64,
    }

    macro_rules! registration_tests_affine {
        ($($name:ident: $t:ty $(,)?)*) => {
            $(
                #[test]
                fn $name() -> Result<()> {
                    let j = julia_image(0f32, 0f32)?.mapv(|x| x as $t);
                    let shape = j.shape();
                    let origin = [
                        ((shape[1] - 1) as f64) / 2f64,
                        ((shape[0] - 1) as f64) / 2f64,
                    ];
                    let s = AffineTransform::new([1.2, 0., 0., 1., 5., 7.], origin, [shape[0], shape[1]]);
                    let k: Array2<_> = s.transform_image_bspline(j.view())?.into();
                    let t = AffineTransform::register_affine(j.view(), k.view())?.inverse()?;
                    let d = (t.matrix() - s.matrix()).powi(2).sum();
                    assert!(d < 0.025, "d: {}, t: {:?}", d, t.parameters);
                    Ok(())
                }
            )*
        }
    }

    registration_tests_affine! {
        registration_tests_affine_u8: u8,
        registration_tests_affine_i8: i8,
        registration_tests_affine_u16: u16,
        registration_tests_affine_i16: i16,
        registration_tests_affine_u32: u32,
        registration_tests_affine_i32: i32,
        registration_tests_affine_u64: u64,
        registration_tests_affine_i64: i64,
        registration_tests_affine_f32: f32,
        registration_tests_affine_f64: f64,
    }
}