use crate::bspline::{BSpline, BSplineTrait};
use crate::error::Error;
use crate::metric::FixedMu;
use crate::register::{Registration, RegistrationResult, RegistrationStep};
use itertools::Itertools;
use ndarray::{Array, Array2, ArrayD, AsArray, Dimension, Ix2, s};
use ndarray_linalg::Inverse;
use num::cast::AsPrimitive;
use serde::{Deserialize, Serialize};
use serde_yaml::{from_reader, to_writer};
use std::fs::File;
use std::marker::PhantomData;
use std::ops::Mul;
use std::path::PathBuf;

/// get coordinates resulting from transforming input coordinates, coordinate must have N
/// columns: x, y, z, ...
#[inline]
pub fn transform_point<'a, T>(point: &'a [T], center: &'a [f64], parameters: &'a [f64]) -> Vec<f64>
where
    T: AsPrimitive<f64>,
{
    // [a, b, c] x = a*x + b*y + c
    // [d, e, f] y = d*x + e*y + f
    // [0, 0, 1] 1 = 1
    debug_assert_eq!(center.len(), point.len());
    debug_assert_eq!((center.len() + 1) * center.len(), parameters.len());

    let m = center.len();
    let n = m * m;
    let point = point
        .iter()
        .zip_eq(center.iter())
        .map(|(p, c)| p.as_() - c)
        .collect::<Vec<_>>();
    parameters
        .iter()
        .take(n)
        .chunks(m)
        .into_iter()
        .zip_eq(parameters.iter().skip(n))
        .zip_eq(center.iter())
        .map(|((r, t), c)| r.zip_eq(point.iter()).map(|(i, p)| i * p).sum::<f64>() + t + c)
        .collect()
}

/// a struct describing the transform
/// generic parameter N = # image dimensions
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct Transform<D: Dimension = Ix2> {
    /// flattened NxN rotation matrix + N translation parameters
    pub parameters: Vec<f64>,
    /// error / significance on parameters
    pub dparameters: Vec<f64>,
    /// the point about which rotations are performed
    pub center: Vec<f64>,
    /// the shape of images for which this transform is meant
    pub shape: Vec<usize>,
    ndim: usize,
    dimension: PhantomData<D>,
}

impl<D: Dimension> Mul for Transform<D> {
    type Output = Transform<D>;

    fn mul(self, rhs: Self) -> Self::Output {
        &self * &rhs
    }
}

impl<D: Dimension> Mul<&Transform<D>> for Transform<D> {
    type Output = Transform<D>;

    fn mul(self, rhs: &Transform<D>) -> Self::Output {
        &self * rhs
    }
}

impl<D: Dimension> Mul<Transform<D>> for &Transform<D> {
    type Output = Transform<D>;

    fn mul(self, rhs: Transform<D>) -> Self::Output {
        self * &rhs
    }
}

impl<D: Dimension> Mul<&Transform<D>> for &Transform<D> {
    type Output = Transform<D>;

    #[allow(clippy::suspicious_arithmetic_impl)]
    fn mul(self, rhs: &Transform<D>) -> Self::Output {
        let m = self.matrix().dot(&rhs.matrix());
        let dm = self.dmatrix().dot(&rhs.matrix()) + self.matrix().dot(&rhs.dmatrix());
        Self::Output {
            parameters: m
                .slice(s![..self.ndim, ..self.ndim])
                .flatten()
                .iter()
                .chain(m.slice(s![..self.ndim, self.ndim]).iter())
                .cloned()
                .collect(),
            dparameters: dm
                .slice(s![..self.ndim, ..])
                .flatten()
                .iter()
                .chain(dm.slice(s![..self.ndim, ..]).iter())
                .cloned()
                .collect(),
            center: self.center.clone(),
            shape: self.shape.clone(),
            ndim: self.ndim,
            dimension: self.dimension,
        }
    }
}

impl<D: Dimension> Eq for Transform<D> {}

impl<D: Dimension> Default for Transform<D> {
    /// the unit transform
    fn default() -> Self {
        let ndim = D::NDIM.expect("number of dimensions must be known to initialise transform");
        let mut parameters = vec![0.0; ndim * ndim + ndim];
        for i in 0..ndim {
            parameters[(ndim + 1) * i] = 1.0;
        }

        Self {
            parameters,
            dparameters: vec![0f64; ndim * ndim + ndim],
            center: vec![0f64; ndim],
            shape: vec![0usize; ndim],
            ndim,
            dimension: PhantomData,
        }
    }
}

impl<D: Dimension> Transform<D> {
    /// parameters: flat NxN part of matrix + translation; center: center of rotation
    pub fn new_with_center(parameters: Vec<f64>, center: Vec<f64>, shape: Vec<usize>) -> Self {
        let ndim = if let Some(ndim) = D::NDIM {
            ndim
        } else {
            center.len()
        };
        debug_assert_eq!(parameters.len(), ndim * ndim + ndim);
        debug_assert_eq!(center.len(), ndim);
        debug_assert_eq!(shape.len(), ndim);
        Self {
            parameters,
            dparameters: vec![0.0; ndim * ndim + ndim],
            center,
            shape,
            ndim,
            dimension: PhantomData,
        }
    }

    pub fn new(parameters: Vec<f64>, shape: Vec<usize>) -> Self {
        let ndim = if let Some(ndim) = D::NDIM {
            ndim
        } else {
            shape.len()
        };
        debug_assert_eq!(parameters.len(), ndim * ndim + ndim);
        debug_assert_eq!(shape.len(), ndim);
        let center = shape.iter().map(|s| ((*s as f64) - 1.0) / 2.0).collect();
        Self {
            parameters,
            dparameters: vec![0.0; ndim * ndim + ndim],
            center,
            shape,
            ndim,
            dimension: PhantomData,
        }
    }

    /// find the affine transform which transforms moving into fixed
    pub fn register_affine<'a, F, M, T>(fixed: F, moving: M) -> Result<Self, Error>
    where
        F: AsArray<'a, T, D>,
        M: AsArray<'a, T, D>,
        T: 'a + Clone + AsPrimitive<f64>,
    {
        Registration::new_affine()?.register(fixed, moving)
    }

    /// find the translation which transforms moving into fixed
    pub fn register_translation<'a, F, M, T>(fixed: F, moving: M) -> Result<Self, Error>
    where
        F: AsArray<'a, T, D>,
        M: AsArray<'a, T, D>,
        T: 'a + Clone + AsPrimitive<f64>,
    {
        Registration::new_translation()?.register(fixed, moving)
    }

    /// find the transform which transforms moving into fixed, using fixed_mu to specify which
    /// parameters to keep fixed
    pub fn register<'a, F, M, T, G>(
        fixed: F,
        moving: M,
        fixed_mu: G,
        steps: Option<Vec<RegistrationStep>>,
        initial_guess: Option<Vec<f64>>,
    ) -> Result<Self, Error>
    where
        F: AsArray<'a, T, D>,
        M: AsArray<'a, T, D>,
        T: 'a + Clone + AsPrimitive<f64>,
        G: Into<FixedMu>,
    {
        let mut registration = Registration::new(fixed_mu);
        if let Some(steps) = steps {
            registration.set_steps(steps);
        }
        if let Some(initial_guess) = initial_guess {
            registration.set_initial_guess(initial_guess);
        }
        registration.register(fixed, moving)
    }

    /// find the transform which transforms moving into fixed and return the results of each
    /// optimization step
    pub fn register_debug<'a, F, M, T, G>(
        fixed: F,
        moving: M,
        fixed_mu: G,
        steps: Option<Vec<RegistrationStep>>,
        initial_guess: Option<Vec<f64>>,
    ) -> Result<(Self, Vec<RegistrationResult>), Error>
    where
        F: AsArray<'a, T, D>,
        M: AsArray<'a, T, D>,
        T: 'a + Clone + AsPrimitive<f64>,
        G: Into<FixedMu>,
    {
        let mut registration = Registration::new(fixed_mu);
        if let Some(steps) = steps {
            registration.set_steps(steps);
        }
        if let Some(initial_guess) = initial_guess {
            registration.set_initial_guess(initial_guess);
        }
        registration.register_debug(fixed, moving)
    }

    /// create a transform from a translation
    pub fn from_translation(translation: &[f64]) -> Self {
        let ndim = if let Some(ndim) = D::NDIM {
            ndim
        } else {
            translation.len()
        };
        debug_assert_eq!(translation.len(), ndim);
        let mut parameters = vec![0.0; ndim * ndim + ndim];
        for i in 0..ndim {
            parameters[(ndim + 1) * i] = 1.0;
        }
        for (p, t) in parameters
            .iter_mut()
            .skip(ndim * ndim)
            .zip_eq(translation.iter())
        {
            *p = *t;
        }

        Self {
            parameters,
            dparameters: vec![0f64; ndim * ndim + ndim],
            center: vec![0f64; ndim],
            shape: vec![0usize; ndim],
            ndim,
            dimension: PhantomData,
        }
    }

    pub fn with_shape(mut self, shape: Vec<usize>) -> Self {
        self.shape = shape;
        self
    }

    /// create a transform from a scaling
    pub fn from_scaling(scaling: &[f64]) -> Self {
        let ndim = if let Some(ndim) = D::NDIM {
            ndim
        } else {
            scaling.len()
        };
        debug_assert_eq!(scaling.len(), ndim);
        let mut parameters = vec![0.0; ndim * ndim + ndim];
        parameters[ndim * ndim + ndim - 1] = 1.0;
        for (p, s) in parameters
            .iter_mut()
            .step_by(ndim + 1)
            .zip_eq(scaling.iter())
        {
            *p = *s;
        }

        Self {
            parameters,
            dparameters: vec![0f64; ndim * ndim + ndim],
            center: vec![0f64; ndim],
            shape: vec![0usize; ndim],
            ndim,
            dimension: PhantomData,
        }
    }

    /// add a translation to self
    pub fn with_translation(&self, translation: &[f64]) -> Self {
        Self::from_translation(translation) * self
    }

    /// scale self
    pub fn with_scaling(&self, scaling: &[f64]) -> Self {
        Self::from_scaling(scaling) * self
    }

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

    /// write a transform to a file
    pub fn to_file(&self, path: PathBuf) -> Result<(), Error> {
        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 {
        let n = self.ndim * self.ndim;
        let m = self.ndim + 1;
        self.parameters
            .iter()
            .take(n)
            .enumerate()
            .all(|(i, &x)| if i % m == 0 { x == 1.0 } else { x == 0.0 })
            && self.parameters.iter().skip(n).all(|&x| x == 0.0)
    }

    /// get coordinates resulting from transforming input coordinates, coordinate must have N
    /// columns: x, y, z, ...
    pub fn transform_point<'a, T>(&'a self, point: &'a [T]) -> Vec<f64>
    where
        T: AsPrimitive<f64>,
    {
        transform_point(point, &self.center, &self.parameters)
    }

    /// get coordinates resulting from transforming input coordinates, coordinates must have N
    /// columns: x, y, z, ...
    pub fn transform_points<'a, A, T>(&'a self, points: A) -> Result<Array2<f64>, Error>
    where
        T: AsPrimitive<f64>,
        A: AsArray<'a, T, Ix2>,
    {
        // [a, b, c] x = a*x + b*y + c
        // [d, e, f] y = d*x + e*y + f
        // [0, 0, 1] 1 = 1
        let points = points.into();
        let shape = points.shape();
        let a = points
            .rows()
            .into_iter()
            .flat_map(|p| {
                if let Some(q) = p.as_slice() {
                    transform_point(q, &self.center, &self.parameters)
                } else {
                    transform_point(&p.to_vec(), &self.center, &self.parameters)
                }
            })
            .collect::<Vec<_>>();
        Ok(ArrayD::from_shape_vec(shape, a)?.into_dimensionality()?)
    }

    /// get the matrix defining the transform
    pub fn matrix(&self) -> Array2<f64> {
        let n = self.ndim * self.ndim;
        let mut matrix = Array2::eye(self.ndim + 1);
        for (m, p) in matrix
            .slice_mut(s![..self.ndim, ..self.ndim])
            .iter_mut()
            .zip_eq(self.parameters.iter().take(n))
        {
            *m = *p;
        }
        for (m, p) in matrix
            .slice_mut(s![..self.ndim, self.ndim])
            .iter_mut()
            .zip_eq(self.parameters.iter().skip(n))
        {
            *m = *p;
        }
        matrix
    }

    /// get the matrix describing the error of the transform
    pub fn dmatrix(&self) -> Array2<f64> {
        let n = self.ndim * self.ndim;
        let mut matrix = Array2::zeros([self.ndim + 1, self.ndim + 1]);
        for (m, p) in matrix
            .slice_mut(s![..self.ndim, ..self.ndim])
            .iter_mut()
            .zip_eq(self.dparameters.iter().take(n))
        {
            *m = *p;
        }
        for (m, p) in matrix
            .slice_mut(s![..self.ndim, self.ndim])
            .iter_mut()
            .zip_eq(self.dparameters.iter().skip(n))
        {
            *m = *p;
        }
        matrix
    }

    /// get the inverse transform
    pub fn inverse(&self) -> Result<Self, Error> {
        let matrix = self.matrix();
        let inverse = matrix.inv()?;
        let parameters = inverse
            .slice(s![..self.ndim, ..self.ndim])
            .iter()
            .chain(inverse.slice(s![..self.ndim, self.ndim]))
            .cloned()
            .collect();

        Ok(Self {
            parameters,
            dparameters: vec![0f64; self.ndim * self.ndim + self.ndim],
            center: self.center.clone(),
            shape: self.shape.clone(),
            ndim: self.ndim,
            dimension: PhantomData,
        })
    }

    /// adapt the transform to a new center and shape
    pub fn adapt(&mut self, center: &[f64], shape: &[usize]) {
        self.center = self
            .shape
            .iter()
            .zip_eq(shape.iter())
            .zip_eq(center)
            .map(|((a, b), o)| o + (((a - b) as f64) / 2.0))
            .collect();
        self.shape = shape.to_vec();
    }

    pub fn interpolate<'a, const B: usize, T, A>(&self, image: A) -> Result<Array<f64, D>, Error>
    where
        A: AsArray<'a, T, D>,
        D: Dimension,
        T: 'a + Clone + AsPrimitive<f64>,
        BSpline<B, D>: BSplineTrait<D>,
    {
        let image = image.into();
        let bspline = BSpline::<B, _>::new(&image);
        bspline.interpolate(self)
    }

    pub fn interpolate_par<'a, const B: usize, T, A>(
        &self,
        image: A,
    ) -> Result<Array<f64, D>, Error>
    where
        A: AsArray<'a, T, D>,
        D: Dimension,
        T: 'a + Clone + AsPrimitive<f64>,
        BSpline<B, D>: BSplineTrait<D>,
    {
        let image = image.into();
        let bspline = BSpline::<B, _>::new(&image);
        bspline.interpolate_par(self)
    }
}

impl Transform<Ix2> {
    /// create a transform from a rotation in radians
    pub fn from_rotation(theta: f64, center: &[f64]) -> Self {
        Self {
            parameters: vec![
                theta.cos(),
                -theta.sin(),
                theta.sin(),
                theta.cos(),
                0.0,
                0.0,
            ],
            dparameters: vec![0f64; 6],
            center: center.to_vec(),
            shape: vec![0usize; 2],
            ndim: 2,
            dimension: PhantomData,
        }
    }

    /// multiply self with a rotation matrix
    pub fn with_rotation(&self, theta: f64) -> Self {
        Self::from_rotation(theta, &self.center) * self
    }

    /// multiply self with a rotation matrix
    pub fn with_rotation_around(&self, theta: f64, center: &[f64]) -> Self {
        Self::from_rotation(theta, center) * self
    }
}

#[cfg(test)]
mod tests {
    use crate::julia_image;
    use crate::transform::Transform;
    use itertools::Itertools;
    use ndarray::{Array2, Ix2, s};
    use num::traits::FloatConst;
    use std::fs::File;
    use std::path::Path;
    use tiff::decoder::{Decoder, DecodingResult};
    use tiff::tags::Tag;
    use tiffwrite::IJTiffFile;

    #[test]
    fn interpolate() -> Result<(), Box<dyn std::error::Error>> {
        let image_a = julia_image(
            &[8000, 6000],
            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
            &[3999.5, 2999.5],
            &[-0.8, 0.156],
        );
        let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
        let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?;
        let image_c = transform.interpolate::<3, _, _>(image_a.view())?;

        assert!(
            image_b
                .into_iter()
                .zip_eq(image_c.into_iter())
                .all(|(a, b)| a == b)
        );
        Ok(())
    }

    #[test]
    fn interpolate_serial() -> Result<(), Box<dyn std::error::Error>> {
        let image_a = julia_image(
            &[8000, 6000],
            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
            &[3999.5, 2999.5],
            &[-0.8, 0.156],
        );
        let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
        let image_b = transform.interpolate::<3, _, _>(image_a.view())?;
        assert_ne!(image_b[[4000, 3000]], 0.0);
        Ok(())
    }

    #[test]
    fn interpolate_par() -> Result<(), Box<dyn std::error::Error>> {
        let image_a = julia_image(
            &[8000, 6000],
            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
            &[3999.5, 2999.5],
            &[-0.8, 0.156],
        );
        let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
        let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?;
        assert_ne!(image_b[[4000, 3000]], 0.0);
        Ok(())
    }

    #[test]
    fn grad_check() -> Result<(), Box<dyn std::error::Error>> {
        use crate::bspline::{BSpline, BSplineTrait};
        use crate::metric::{MattesMetric, SamplingArg};
        use algos::ObjectiveFunction;
        let im_a = julia_image(
            &[100, 1],
            &[1.0, 0.0, 0.0, 0.01, 0.0, 0.0],
            &[99.5, 0.5],
            &[-0.8, 0.156],
        )
        .slice(s![.., 0])
        .mapv(|i| i as f64);
        let q = vec![0.85, 4.0];
        let im_b =
            Transform::new(q.clone(), vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;

        // Test with ALL points (like grad_check uses)
        let metric_all = MattesMetric::<ndarray::Ix1>::new(
            BSpline::new(im_b.view()),
            BSpline::new(im_a.view()),
            SamplingArg::FixedAt((0..100).map(|i| vec![i as f64]).collect()),
            3,
            0.05,
        )?
        .with_fixed_mu(crate::metric::FixedMu::new_none(1));

        let identity = vec![1.0, 0.0];
        let val_i = metric_all.evaluate(&identity);
        let grad_i = metric_all.gradient(&identity).unwrap();
        println!("ALL POINTS - identity: val={}, grad={:?}", val_i, grad_i);
        let truth = vec![0.85, 4.0];
        let val_t = metric_all.evaluate(&truth);
        let grad_t = metric_all.gradient(&truth).unwrap();
        println!("ALL POINTS - truth: val={}, grad={:?}", val_t, grad_t);
        let inv = vec![1.0 / 0.85, -4.0 / 0.85];
        let val_inv = metric_all.evaluate(&inv);
        println!("ALL POINTS - inverse: val={}", val_inv);

        // Test with 100 random points (like register1 uses at finest level)
        let metric_rand = MattesMetric::<ndarray::Ix1>::new(
            BSpline::new(im_b.view()),
            BSpline::new(im_a.view()),
            SamplingArg::Fixed(100),
            3,
            0.05,
        )?
        .with_fixed_mu(crate::metric::FixedMu::new_none(1));

        let val_i2 = metric_rand.evaluate(&identity);
        let grad_i2 = metric_rand.gradient(&identity).unwrap();
        println!("RAND 100 - identity: val={}, grad={:?}", val_i2, grad_i2);

        // Test with 100 random points and 32 bins (like level 0)
        let metric_32 = MattesMetric::<ndarray::Ix1>::new(
            BSpline::new(im_b.view()),
            BSpline::new(im_a.view()),
            SamplingArg::Fixed(100),
            32,
            0.05,
        )?
        .with_fixed_mu(crate::metric::FixedMu::new_none(1));

        let val_i3 = metric_32.evaluate(&identity);
        let grad_i3 = metric_32.gradient(&identity).unwrap();
        println!(
            "RAND 100 bins=32 - identity: val={}, grad={:?}",
            val_i3, grad_i3
        );

        let eps = 1e-5;
        for i in 0..2 {
            let mut p_plus = identity.clone();
            let mut p_minus = identity.clone();
            p_plus[i] += eps;
            p_minus[i] -= eps;
            let num_grad =
                (metric_all.evaluate(&p_plus) - metric_all.evaluate(&p_minus)) / (2.0 * eps);
            println!(
                "numerical d/dmu[{}] = {} (analytical: {})",
                i, num_grad, grad_i[i]
            );
        }
        Ok(())
    }

    #[test]
    fn metric_landscape() -> Result<(), Box<dyn std::error::Error>> {
        use crate::bspline::{BSpline, BSplineTrait};
        use crate::metric::{MattesMetric, SamplingArg, Sigma};
        use algos::ObjectiveFunction;
        let im_a = julia_image(
            &[100, 1],
            &[1.0, 0.0, 0.0, 0.01, 0.0, 0.0],
            &[99.5, 0.5],
            &[-0.8, 0.156],
        )
        .slice(s![.., 0])
        .mapv(|i| i as f64);
        let im_b =
            Transform::new(vec![0.85, 4.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;

        // Smooth both images with sigma=8 like level 0
        let sigma = Sigma::Absolute(vec![8.0]);
        let sf = sigma.smooth(im_a.view())?;
        let sm = sigma.smooth(im_b.view())?;

        let points: Vec<Vec<f64>> = (0..100).map(|i| vec![i as f64]).collect();

        let metric = MattesMetric::<ndarray::Ix1>::new(
            BSpline::new(sf.view()),
            BSpline::new(sm.view()),
            SamplingArg::FixedAt(points.clone()),
            32,
            0.05,
        )?
        .with_fixed_mu(crate::metric::FixedMu::new_none(1));

        // Test several points
        let test_points: Vec<(&str, Vec<f64>)> = vec![
            ("identity".into(), vec![1.0, 0.0]),
            ("truth".into(), vec![0.85, 4.0]),
            ("neg_trans".into(), vec![1.0, -4.0]),
            ("scale_0.9".into(), vec![0.9, 0.0]),
            ("scale_1.1".into(), vec![1.1, 0.0]),
        ];

        println!("=== SMOOTHED ALL 100 integer points bins=32 ===");
        for (name, p) in &test_points {
            let val = metric.evaluate(p);
            let grad = metric.gradient(p);
            match grad {
                Some(g) => println!("  {}: val={:.6}, grad={:?}", name, val, g),
                None => println!("  {}: val={:.6}, grad=None", name, val),
            }
        }

        // Now test with Fixed(100) random points
        let metric_rand = MattesMetric::<ndarray::Ix1>::new(
            BSpline::new(sf.view()),
            BSpline::new(sm.view()),
            SamplingArg::Fixed(100),
            32,
            0.05,
        )?
        .with_fixed_mu(crate::metric::FixedMu::new_none(1));

        println!("=== SMOOTHED Random 100 points bins=32 ===");
        for (name, p) in &test_points {
            let val = metric_rand.evaluate(p);
            let grad = metric_rand.gradient(p);
            match grad {
                Some(g) => println!("  {}: val={:.6}, grad={:?}", name, val, g),
                None => println!("  {}: val={:.6}, grad=None", name, val),
            }
        }

        Ok(())
    }

    #[test]
    fn register1() -> Result<(), Box<dyn std::error::Error>> {
        let im_a = julia_image(
            &[100, 1],
            &[1.0, 0.0, 0.0, 0.01, 0.0, 0.0],
            &[99.5, 0.5],
            &[-0.8, 0.156],
        )
        .slice(s![.., 0])
        .mapv(|i| i as f64);
        let q = vec![0.85, 2.0];
        let im_b =
            Transform::new(q.clone(), vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;

        // The registration finds T such that im_b(T(x)) = im_a(x), which is the inverse of q
        let q_inv = Transform::<ndarray::Ix1>::new(q.clone(), vec![im_a.shape()[0]])
            .inverse()?
            .parameters;

        // Use default steps with ASGD — the multi-resolution pyramid converges
        // reliably for this 1D affine problem.
        let t = Transform::register(im_a.view(), im_b.view(), vec![None; 2], None, None)?;
        println!("t: {:?}", t);
        println!("i: {:?}", t.inverse()?);
        println!("q_inv: {:?}", q_inv);
        assert!(
            t.parameters
                .iter()
                .zip(q_inv.iter())
                .map(|(a, b)| (a - b).powi(2))
                .sum::<f64>()
                < 1.0
        );
        Ok(())
    }

    #[test]
    fn register2_random_affine() -> Result<(), Box<dyn std::error::Error>> {
        use rand::prelude::*;
        let mut rng = rand::rngs::StdRng::seed_from_u64(1337);

        let shape = [200, 200];
        let center = [99.5, 99.5];

        let fixed = julia_image(
            &shape,
            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
            &center,
            &[-0.8, 0.156],
        )
        .mapv(|i| i as f64);

        // random rotation ±25°
        let angle: f64 = rng.random_range(-25.0f64..25.0).to_radians();
        let (s, c) = angle.sin_cos();

        // random scale 0.85–1.15
        let sx = rng.random_range(0.85..1.15);
        let sy = rng.random_range(0.85..1.15);

        // small shear
        let shx: f64 = rng.random_range(-0.1..0.1);
        let shy: f64 = rng.random_range(-0.1..0.1);

        // translation ±30 px
        let tx: f64 = rng.random_range(-30.0..30.0);
        let ty: f64 = rng.random_range(-30.0..30.0);

        // [m00, m01, m10, m11, tx, ty]
        let params = vec![c * sx, -s * sy + shx, s * sx + shy, c * sy, tx, ty];

        let transform =
            Transform::<Ix2>::new_with_center(params.clone(), center.to_vec(), shape.to_vec());
        let moving = transform.interpolate::<3, _, _>(fixed.view())?;

        let q_inv = transform.inverse()?.parameters;

        let (t, _steps) = Transform::<Ix2>::register_debug(
            fixed.view(),
            moving.view(),
            vec![None, None, None, None, None, None],
            None,
            None,
        )?;
        println!("params:  {:?}", params);
        println!("result:  {:?}", t.parameters);
        println!("q_inv:   {:?}", q_inv);
        let sse: f64 = t
            .parameters
            .iter()
            .zip(q_inv.iter())
            .map(|(a, b)| (a - b).powi(2))
            .sum();
        println!("sse:     {sse}");
        assert!(sse < 1.0);
        Ok(())
    }

    #[test]
    fn metric_landscape_2d() -> Result<(), Box<dyn std::error::Error>> {
        use crate::bspline::{BSpline, BSplineTrait};
        use crate::filter::gaussian_smooth;
        use crate::metric::{MattesMetric, SamplingArg};
        use algos::ObjectiveFunction;

        let shape = [200, 200];
        let center = [99.5, 99.5];
        let im_a = julia_image(
            &shape,
            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
            &center,
            &[-0.8, 0.156],
        )
        .mapv(|i| i as f64);
        let rotation = Transform::<Ix2>::from_rotation(f64::PI() / 4.0, &center);
        let im_b = rotation.interpolate::<3, _, _>(im_a.view())?;

        let q_inv = rotation.inverse()?.parameters;
        let p = rotation.parameters.clone();
        let identity = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
        let edge = 0.01;

        // Test at full resolution with different sigma levels (matching elastix approach)
        // Elastix does NOT downsample — it smooths at full resolution with
        // sigma = 0.5 * factor where factor is from the pyramid schedule [8, 4, 2, 1]
        println!(
            "=== Full resolution (no downsampling, matching elastix FixedSmoothingImagePyramid) ==="
        );
        for (level, sigma_val) in [4.0, 2.0, 1.0, 0.5].iter().enumerate() {
            let f = gaussian_smooth(im_a.view(), &[*sigma_val; 2])?;
            let m = gaussian_smooth(im_b.view(), &[*sigma_val; 2])?;
            let bf = BSpline::<0, _>::new(f.view());
            let bm = BSpline::<3, _>::new(m.view());
            let metric = MattesMetric::new(bf, bm, SamplingArg::Random(3000), 128, edge)?;

            let mi_id = metric.evaluate(&identity);
            let mi_p = metric.evaluate(&p);
            let mi_qinv = metric.evaluate(&q_inv);

            println!(
                "  Level {} (sigma={:.1}, {}x{}): id={:.4} p={:.4} q_inv={:.4} → {}",
                level,
                sigma_val,
                f.shape()[0],
                f.shape()[1],
                -mi_id,
                -mi_p,
                -mi_qinv,
                if -mi_qinv > -mi_p {
                    "Q_INV correct"
                } else {
                    "P incorrect (landscape inverted!)"
                }
            );
        }

        // Verify: at every sigma level, q_inv should have higher MI than p
        let f4 = gaussian_smooth(im_a.view(), &[4.0, 4.0])?;
        let m4 = gaussian_smooth(im_b.view(), &[4.0, 4.0])?;
        let bf4 = BSpline::<0, _>::new(f4.view());
        let bm4 = BSpline::<3, _>::new(m4.view());
        let metric4 = MattesMetric::new(bf4, bm4, SamplingArg::Random(5000), 128, edge)?;
        let mi_p_coarse = metric4.evaluate(&p);
        let mi_qinv_coarse = metric4.evaluate(&q_inv);
        println!(
            "\nCoarsest (sigma=4.0): p={:.6} q_inv={:.6}",
            -mi_p_coarse, -mi_qinv_coarse
        );
        assert!(
            -mi_qinv_coarse > -mi_p_coarse,
            "MI landscape is inverted at coarsest level! q_inv={} should be > p={}",
            -mi_qinv_coarse,
            -mi_p_coarse
        );

        Ok(())
    }

    #[test]
    fn angle_sweep() -> Result<(), Box<dyn std::error::Error>> {
        use crate::bspline::{BSpline, BSplineTrait};
        use crate::filter::gaussian_smooth;
        use crate::metric::{MattesMetric, SamplingArg};
        use algos::ObjectiveFunction;
        let shape = [200, 200];
        let center = [99.5, 99.5];
        let im_a = julia_image(
            &shape,
            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
            &center,
            &[-0.8, 0.156],
        )
        .mapv(|i| i as f64);
        let edge = 0.01;
        // Full resolution, sigma=4.0 — matching elastix FixedSmoothingImagePyramid (no downsampling)
        for angle_deg in [5.0f64, 10.0, 15.0, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0] {
            let rotation = Transform::<Ix2>::from_rotation(angle_deg.to_radians(), &center);
            let im_b = rotation.interpolate::<3, _, _>(im_a.view())?;
            let q_inv = rotation.inverse()?.parameters;
            let p = rotation.parameters.clone();
            let f = gaussian_smooth(im_a.view(), &[4.0, 4.0])?;
            let m = gaussian_smooth(im_b.view(), &[4.0, 4.0])?;
            let metric = MattesMetric::new(
                BSpline::<0, _>::new(f.view()),
                BSpline::<3, _>::new(m.view()),
                SamplingArg::Random(3000),
                128,
                edge,
            )?;
            let id = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
            let mi_id = metric.evaluate(&id);
            let mi_p = metric.evaluate(&p);
            let mi_qi = metric.evaluate(&q_inv);
            println!(
                "{:5.1}°: id={:.3} p={:.3} q_inv={:.3} → {}",
                angle_deg,
                -mi_id,
                -mi_p,
                -mi_qi,
                if mi_p > mi_qi { "P wins" } else { "Q_INV wins" }
            );
        }
        Ok(())
    }

    fn read_tiff<P: AsRef<Path>>(path: P) -> Result<Array2<f64>, Box<dyn std::error::Error>> {
        let mut reader = Decoder::new(File::open(path)?)?;
        reader.seek_to_image(0)?;
        let bytes = match reader.read_image()? {
            DecodingResult::U8(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
            DecodingResult::U16(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
            DecodingResult::U32(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
            DecodingResult::U64(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
            DecodingResult::I8(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
            DecodingResult::I16(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
            DecodingResult::I32(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
            DecodingResult::I64(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
            DecodingResult::F16(data) => data.into_iter().map(f64::from).collect::<Vec<_>>(),
            DecodingResult::F32(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
            DecodingResult::F64(data) => data,
        };
        let width = reader.get_tag(Tag::ImageWidth)?.into_u32()? as usize;
        let height = reader.get_tag(Tag::ImageLength)?.into_u32()? as usize;
        Ok(Array2::from_shape_vec((width, height), bytes)?)
    }

    #[test]
    fn register_real_images() -> Result<(), Box<dyn std::error::Error>> {
        let fixed = read_tiff("test_files/fixed.tif")?;
        let moving = read_tiff("test_files/moving.tif")?;

        // SimpleElastix expected — the registration result is in TIFF-native coordinates.
        let expected = [
            0.9899559376493817,
            0.011269992506480442,
            0.017048860489651384,
            0.8533772512806085,
            -12.877028775909979,
            4.143118928117275,
        ];

        let t = Transform::<Ix2>::register(
            fixed.view(),
            moving.view(),
            vec![None, None, None, None, None, None],
            None,
            None,
        )?;

        let expected_transform = Transform::<Ix2>::new(expected.to_vec(), fixed.shape().to_vec());

        // Compare transformed coordinates of all pixels
        let shape = fixed.shape();
        let mut sum_diff = 0.0;
        let mut count = 0;
        for row in 0..shape[0] {
            for col in 0..shape[1] {
                let point = [row as f64, col as f64];
                let t_point = t.transform_point(&point);
                let e_point = expected_transform.transform_point(&point);
                let diff_sq = (t_point[0] - e_point[0]).powi(2) + (t_point[1] - e_point[1]).powi(2);
                sum_diff += diff_sq.sqrt();
                count += 1;
            }
        }
        let mean_diff = sum_diff / count as f64;
        println!("Our: {:?}  mean_coord_diff: {:.6}", t, mean_diff);

        let mut tif = IJTiffFile::new(
            std::env::home_dir()
                .unwrap()
                .join("tmp/register_real_images.tif"),
        )?;
        tif.save(fixed.mapv(|i| i as u16), 0, 0, 0)?;
        tif.save(
            t.interpolate_par::<1, _, _>(moving.view())?
                .mapv(|i| i as u16),
            1,
            0,
            0,
        )?;
        tif.save(moving.mapv(|i| i as u16), 2, 0, 0)?;

        assert!(mean_diff < 0.1);

        Ok(())
    }

    #[test]
    fn register_real_images2() -> Result<(), Box<dyn std::error::Error>> {
        let fixed = read_tiff("test_files/fixed.tif")?;
        let e = Transform::<Ix2>::new(vec![0.8, 0.0, 0.0, 1.0, 0.0, 0.0], fixed.shape().to_vec())
            .inverse()?;
        let moving = e.interpolate::<3, _, _>(fixed.view())?;

        let t = Transform::<Ix2>::register(
            fixed.view(),
            moving.view(),
            vec![None, None, None, None, None, None],
            None,
            None,
        )?;

        let e_inv = e.inverse()?;

        // Compare transformed coordinates of all pixels
        let shape = fixed.shape();
        let mut sum_diff = 0.0;
        let mut count = 0;
        for row in 0..shape[0] {
            for col in 0..shape[1] {
                let point = [row as f64, col as f64];
                let t_point = t.transform_point(&point);
                let e_point = e_inv.transform_point(&point);
                let diff_sq = (t_point[0] - e_point[0]).powi(2) + (t_point[1] - e_point[1]).powi(2);
                sum_diff += diff_sq.sqrt();
                count += 1;
            }
        }
        let mean_diff = sum_diff / count as f64;
        println!("Our: {:?}  mean_coord_diff: {:.6}", t, mean_diff);

        let mut tif = IJTiffFile::new(
            std::env::home_dir()
                .unwrap()
                .join("tmp/register_real_images2.tif"),
        )?;
        tif.save(fixed.mapv(|i| i as u16), 0, 0, 0)?;
        tif.save(
            t.interpolate_par::<1, _, _>(moving.view())?
                .mapv(|i| i as u16),
            1,
            0,
            0,
        )?;
        tif.save(moving.mapv(|i| i as u16), 2, 0, 0)?;

        assert!(mean_diff < 0.1);

        Ok(())
    }
}