use crate::bspline::{
    BSpline, BSplineTrait, cubic_bspline, square_bspline, square_bspline_integral,
};
use crate::error::Error;
use crate::filter::gaussian_smooth;
use crate::transform::transform_point;
use algos::ObjectiveFunction;
use itertools::Itertools;
use ndarray::{Array, Array1, Array2, Array3, AsArray, Axis, Dimension, s};
use num::cast::AsPrimitive;
use rand::RngExt;
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use serde::{Deserialize, Serialize};
use std::cell::RefCell;
use std::ops::Deref;

#[derive(Default)]
struct IntMut {
    mu: Vec<f64>,
    metric: f64,
    derivative: Vec<f64>,
}

pub struct FixedMu {
    mu: Vec<Option<f64>>,
    len_variable: usize,
}

impl Deref for FixedMu {
    type Target = Vec<Option<f64>>;

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

impl From<Vec<Option<f64>>> for FixedMu {
    fn from(mu: Vec<Option<f64>>) -> Self {
        let len_variable = mu.iter().filter(|v| v.is_none()).count();
        FixedMu { mu, len_variable }
    }
}

impl FixedMu {
    pub fn new_none(n: usize) -> Self {
        Self {
            mu: vec![None; n * n + n],
            len_variable: n * n + n,
        }
    }

    /// combine mu and fixed_mu into complete parameter set
    pub fn combine(&self, mu: &[f64]) -> Vec<f64> {
        debug_assert_eq!(mu.len(), self.len_variable);
        let mut mu_iter = mu.iter();
        self.iter()
            .map(|i| {
                if let Some(i) = i {
                    *i
                } else {
                    *mu_iter.next().unwrap()
                }
            })
            .collect()
    }

    /// combine the dparameters of the variable parameters into a complete
    /// parameter set, with no uncertainty (0) on the fixed parameters
    pub fn combine_dparameters(&self, dparameters: &[f64]) -> Vec<f64> {
        debug_assert_eq!(dparameters.len(), self.len_variable);
        let mut dparameters_iter = dparameters.iter();
        self.iter()
            .map(|i| {
                if i.is_some() {
                    0.0
                } else {
                    *dparameters_iter.next().unwrap()
                }
            })
            .collect()
    }

    /// extract mu from complete parameter set
    pub fn extract_variable(&self, mu: &[f64]) -> Vec<f64> {
        self.iter()
            .zip_eq(mu)
            .filter_map(|(i, j)| if i.is_none() { Some(*j) } else { None })
            .collect()
    }

    pub fn len_variable(&self) -> usize {
        self.len_variable
    }
}

/// first_bin, last_bin: center of bin
/// n_bins >= 2
#[expect(clippy::too_many_arguments)]
fn parzen(
    alpha: f64,
    dalpha: &[f64],
    fixed: &[f64],
    moving: &[f64],
    dmoving: &[Vec<f64>],
    weight: &[f64],
    dweight: &[Vec<f64>],
    first_bin: f64,
    last_bin: f64,
    n_bins: usize,
) -> (Array2<f64>, Array3<f64>) {
    let ndim = dalpha.len();
    let width = (last_bin - first_bin) / ((n_bins - 1) as f64);
    let v = 1.0 / width;
    // add two bins left and right
    let first_bin = first_bin - 2.0 * width;
    let last_bin = last_bin + 2.0 * width;
    let n_bins = n_bins + 4;
    let n_bins_i = n_bins as isize;
    let mut hist = Array2::zeros([n_bins, n_bins]);
    let mut grad_b = Array3::zeros([n_bins, n_bins, ndim]);
    let left = first_bin - width / 2.0;
    let bins = Array1::linspace(first_bin, last_bin, n_bins);
    let mut nf;
    let mut nm;
    let mut uf;
    let mut um;
    let mut kf;
    let mut km;
    let mut csf;
    for ((f, m), (dm, (w, dw))) in fixed
        .iter()
        .zip_eq(moving)
        .zip_eq(dmoving.iter().zip_eq(weight.iter().zip_eq(dweight)))
    {
        nf = (v * (f - left)) as isize;
        nm = (v * (m - left)) as isize;
        for if_ in (nf - 2)..=(nf + 2) {
            if (0 <= if_) && (if_ < n_bins_i) {
                uf = if_ as usize;
                kf = (f - bins[uf]) * v;
                csf = cubic_bspline(kf);
                for im in (nm - 2)..=(nm + 2) {
                    if (0 <= im) && (im < n_bins_i) {
                        um = im as usize;
                        km = (m - bins[um]) * v;
                        hist[[uf, um]] += w * csf * cubic_bspline(km);
                        for ((dmi, dwi), d) in
                            dm.iter()
                                .zip_eq(dw)
                                .zip_eq(grad_b.slice_mut(s![uf, um, ..]))
                        {
                            *d += w
                                * csf
                                * (square_bspline(km + 0.5) - square_bspline(km - 0.5))
                                * v
                                * dmi
                                + dwi * csf * cubic_bspline(km);
                        }
                    }
                }
            }
        }
    }
    let hist = hist / alpha;
    let mut hist_bc = hist
        .broadcast([dalpha.len(), hist.shape()[0], hist.shape()[1]])
        .unwrap();
    hist_bc.permute_axes([1, 2, 0]);
    let dhist = (grad_b - (&hist_bc * &Array1::from_iter(dalpha))) / alpha;
    (hist, dhist)
}

enum Sampling {
    Fixed(Vec<Vec<f64>>),
    Random((usize, Vec<f64>)),
}

impl Sampling {
    fn fixed(index: Vec<Vec<f64>>) -> Self {
        Self::Fixed(index)
    }

    fn random(n_samples: usize, shape: Vec<f64>) -> Self {
        Self::Random((n_samples, shape))
    }

    fn index(&self) -> Vec<Vec<f64>> {
        match self {
            Self::Fixed(index) => index.clone(),
            Self::Random((n_samples, shape)) => rand::rng()
                .random_iter::<f64>()
                .take(n_samples * shape.len())
                .chunks(shape.len())
                .into_iter()
                .map(|c| c.zip_eq(shape.iter()).map(|(i, s)| i * s - 0.5).collect())
                .collect(),
        }
    }
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub enum SamplingArg {
    Fixed(usize),
    FixedAt(Vec<Vec<f64>>),
    Random(usize),
}

/// (dI / dT(x)) * (dT(x) / dmu)
fn image_jacobian(point: &[f64], center: &[f64], image_derivative: &[f64]) -> Vec<f64> {
    image_derivative
        .iter()
        .flat_map(|d| {
            point
                .iter()
                .zip_eq(center.iter())
                .map(|(p, c)| *d * (p - c))
        })
        .chain(image_derivative.iter().cloned())
        .collect()
}

/// Local Gauss-Newton statistics of the intensity fit
/// `fixed(x) = moving(T(x)) + noise` at a parameter point.
///
/// `hessian` is the weighted normal matrix `JᵀWJ` and `variance` is the
/// weighted residual variance `Σ w r² / (Σ w − q − 1)`, with `q` the number of
/// variable parameters. The extra degree of freedom accounts for the weighted
/// mean intensity offset that is removed from the residuals. The parameter
/// covariance of the fit is `variance * hessian⁻¹`.
#[derive(Clone, Debug)]
pub struct ResidualFit {
    /// weighted normal matrix `JᵀWJ`
    pub hessian: Vec<Vec<f64>>,
    /// residual variance (mean squared residual, corrected for the dof)
    pub variance: f64,
    /// number of sample points that contributed
    pub n_samples: usize,
}

pub struct MattesMetric<D>
where
    D: Dimension,
{
    shape: Vec<f64>,
    center: Vec<f64>,
    fixed: BSpline<0, D>,
    moving: BSpline<3, D>,
    fixed_mu: FixedMu,
    minmax: [f64; 2],
    sampling: Sampling,
    n_bins: usize,
    /// memory for (point, value, gradient) combination, refcell gives interior mutability
    metric: RefCell<IntMut>,
    edge: f64,
}

impl<D> MattesMetric<D>
where
    D: Dimension,
{
    pub fn new(
        fixed: BSpline<0, D>,
        moving: BSpline<3, D>,
        sampling: SamplingArg,
        n_bins: usize,
        edge: f64,
    ) -> Result<Self, Error> {
        let shape = fixed.shape().iter().map(|&i| i as f64).collect::<Vec<_>>();
        let center = shape.iter().map(|i| (i - 1.0) / 2.0).collect();
        let (min, max) = fixed
            .iter()
            .chain(moving.iter())
            .minmax()
            .into_option()
            .expect("m and f cannot be empty");
        let minmax = [*min, *max];
        let ndim = shape.len();
        let sampling = match sampling {
            SamplingArg::Fixed(n_samples) => {
                Sampling::fixed(Sampling::random(n_samples, shape.clone()).index())
            }
            SamplingArg::FixedAt(points) => Sampling::fixed(points),
            SamplingArg::Random(n_samples) => Sampling::random(n_samples, shape.clone()),
        };

        Ok(Self {
            shape,
            center,
            fixed,
            moving,
            fixed_mu: FixedMu::new_none(ndim),
            minmax,
            sampling,
            n_bins,
            metric: RefCell::default(),
            edge,
        })
    }

    pub fn new_from_arrays<'a, T, A>(
        fixed: A,
        moving: A,
        sampling: SamplingArg,
        n_bins: usize,
        edge: f64,
    ) -> Result<Self, Error>
    where
        A: AsArray<'a, T, D>,
        T: 'a + Clone + AsPrimitive<f64>,
    {
        Self::new(
            BSpline::<0, _>::new(fixed),
            BSpline::<3, _>::new(moving),
            sampling,
            n_bins,
            edge,
        )
    }

    pub fn with_fixed_mu<F: Into<FixedMu>>(mut self, fixed_mu: F) -> Self {
        self.fixed_mu = fixed_mu.into();
        self
    }

    pub fn fixed_mu(&self) -> &FixedMu {
        &self.fixed_mu
    }

    pub fn eval(&self, mu: &[f64]) -> Result<(), Error> {
        assert!(mu.iter().all(|i| i.is_finite()));
        if self.metric.borrow().mu != mu {
            let complete_mu = self.fixed_mu.combine(mu);
            let idx = self.sampling.index();
            let shape = self.shape.clone();
            let shape_f = self.shape.iter().map(|s| s - 0.5).collect::<Vec<_>>();
            let center = self.center.as_slice();
            let fixed = &self.fixed;
            let mem_f = fixed.get_mem();
            let moving = &self.moving;
            let mem_m = moving.get_mem();
            let fixed_mu = &self.fixed_mu;
            let e = 1.0 / self.edge;
            let (f, ((m, j), (w, dw))) = idx
                .into_par_iter()
                .map_init(
                    || (mem_f.clone(), mem_m.clone()),
                    |(mem_f, mem_m), u| {
                        let v = transform_point(&u, center, &complete_mu);
                        if shape_f
                            .iter()
                            .zip_eq(v.iter())
                            .all(|(s, x)| (-0.5 <= *x) && (x <= s))
                        {
                            let f = fixed.evaluate_at_continuous_index(&u, mem_f)?;
                            let (m, d) = moving
                                .evaluate_value_and_derivative_at_continuous_index(&v, mem_m)?;
                            let j = fixed_mu.extract_variable(&image_jacobian(&u, center, &d));
                            let mut k;
                            let mut w = 1.0;
                            let mut ws = Vec::with_capacity(shape.len());
                            let mut dw = Vec::with_capacity(shape.len());
                            // the metric wouldn't be smooth if we'd drop a point abruptly,
                            // so apply a 0 <= weight <= 1 to points within 5% of the edge
                            for (s, x) in shape.iter().zip_eq(&v) {
                                if e * (x + 0.5) < *s {
                                    k = 3.0 * e * (x + 0.5) / s - 1.5;
                                    let b = square_bspline_integral(k);
                                    w *= b;
                                    ws.push(b);
                                    dw.push(square_bspline(k) * 3.0 * e / s);
                                } else if e * (x + 0.5) > (e - 1.0) * s {
                                    k = 3.0 * e * (s - x - 0.5) / s - 1.5;
                                    let b = square_bspline_integral(k);
                                    w *= b;
                                    ws.push(b);
                                    dw.push(-square_bspline(k) * 3.0 * e / s);
                                } else {
                                    ws.push(1.0);
                                    dw.push(0.0);
                                }
                            }
                            for (dwi, wi) in dw.iter_mut().zip(ws) {
                                *dwi *= w / wi;
                            }
                            let dw = fixed_mu.extract_variable(&image_jacobian(&u, center, &dw));
                            Ok(Some((f, ((m, j), (w, dw)))))
                        } else {
                            Ok(None)
                        }
                    },
                )
                .filter_map(|i| i.transpose())
                .collect::<Result<
                    (
                        Vec<f64>,
                        ((Vec<f64>, Vec<Vec<f64>>), (Vec<f64>, Vec<Vec<f64>>)),
                    ),
                    Error,
                >>()?;
            let alpha = w.iter().sum::<f64>();
            let mut dalpha = vec![0.0; mu.len()];
            for i in dw.iter() {
                for (j, a) in i.iter().zip_eq(dalpha.iter_mut()) {
                    *a += j;
                }
            }

            let (jpdf, d_jpdf_m) = parzen(
                alpha,
                &dalpha,
                &f,
                &m,
                j.as_slice(),
                &w,
                &dw,
                self.minmax[0],
                self.minmax[1],
                self.n_bins,
            );

            let pdf_f = jpdf.axis_iter(Axis(0)).map(|i| i.sum()).collect::<Vec<_>>();
            let pdf_m = jpdf.axis_iter(Axis(1)).map(|i| i.sum()).collect::<Vec<_>>();
            let d_pdf_m = d_jpdf_m
                .axis_iter(Axis(1))
                .zip(pdf_m.iter())
                .map(|(di, i)| di.sum() / i)
                .collect::<Vec<_>>();

            let mut metric = 0.0;
            let mut dmetric = vec![0.0; self.fixed_mu.len_variable];
            let mut n;
            let mut m;
            let mut ik;
            let mut ln;

            for ((row, d_row), k) in jpdf
                .axis_iter(Axis(0))
                .zip_eq(d_jpdf_m.axis_iter(Axis(0)))
                .zip_eq(pdf_f)
            {
                debug_assert_eq!(row.len(), d_row.shape()[0]);
                debug_assert_eq!(row.len(), pdf_m.len());

                for ((&j, dj), (i, di)) in row
                    .into_iter()
                    .zip_eq(d_row.axis_iter(Axis(0)))
                    .zip_eq(pdf_m.iter().zip(d_pdf_m.iter()))
                {
                    ik = i * k;
                    // check for non-zero bin contribution
                    if (j > 1e-16) && (ik > 1e-16) {
                        ln = (j / (ik)).ln();
                        n = 1.0 + ln;
                        m = di * j;
                        metric -= j * ln;
                        for (d, g) in dmetric.iter_mut().zip_eq(dj) {
                            *d += m - n * g;
                        }
                    }
                }
            }

            debug_assert!(mu.iter().all(|i| i.is_finite()));
            debug_assert!(alpha.is_finite());
            debug_assert!(metric.is_finite());
            debug_assert!(dmetric.iter().all(|i| i.is_finite()));

            self.metric.replace(IntMut {
                mu: mu.to_vec(),
                metric,
                derivative: dmetric.clone(),
            });
        }
        Ok(())
    }

    /// Compute the Gauss-Newton normal matrix and the residual variance of the
    /// intensity fit at `point`, using the same sample points and edge
    /// weighting as [`Self::eval`].
    ///
    /// Residuals are `fixed(x) - moving(T(x))`, so a well-registered image pair
    /// gives a small variance. Returns `None` when interpolation fails or when
    /// too few samples contribute to estimate the residual variance.
    pub fn residual_fit(&self, point: &[f64]) -> Option<ResidualFit> {
        let complete_mu = self.fixed_mu.combine(point);
        let idx = self.sampling.index();
        let shape = self.shape.as_slice();
        let shape_f = self.shape.iter().map(|s| s - 0.5).collect::<Vec<_>>();
        let center = self.center.as_slice();
        let e = 1.0 / self.edge;
        let n_params = point.len();
        let mut hessian = vec![vec![0.0; n_params]; n_params];
        let mut sum_w = 0.0;
        let mut sum_wr = 0.0;
        let mut sum_wr2 = 0.0;
        let mut n_samples = 0;
        let mut mem_f = self.fixed.get_mem();
        let mut mem_m = self.moving.get_mem();
        for u in idx {
            let v = transform_point(&u, center, &complete_mu);
            if !shape_f
                .iter()
                .zip_eq(v.iter())
                .all(|(s, x)| (-0.5 <= *x) && (x <= s))
            {
                continue;
            }
            let f = self
                .fixed
                .evaluate_at_continuous_index(&u, &mut mem_f)
                .ok()?;
            let (m, d) = self
                .moving
                .evaluate_value_and_derivative_at_continuous_index(&v, &mut mem_m)
                .ok()?;
            // same smooth edge weighting as in `eval`
            let mut w = 1.0;
            for (s, x) in shape.iter().zip_eq(&v) {
                if e * (x + 0.5) < *s {
                    let k = 3.0 * e * (x + 0.5) / s - 1.5;
                    w *= square_bspline_integral(k);
                } else if e * (x + 0.5) > (e - 1.0) * s {
                    let k = 3.0 * e * (s - x - 0.5) / s - 1.5;
                    w *= square_bspline_integral(k);
                }
            }
            if w <= 0.0 || !w.is_finite() {
                continue;
            }
            let r = f - m;
            sum_w += w;
            sum_wr += w * r;
            sum_wr2 += w * r * r;
            n_samples += 1;
            let j = self
                .fixed_mu
                .extract_variable(&image_jacobian(&u, center, &d));
            for (a, ja) in j.iter().enumerate() {
                for (b, jb) in j.iter().enumerate() {
                    hessian[a][b] += w * ja * jb;
                }
            }
        }
        // the residual variance needs at least one degree of freedom
        if n_samples == 0 || sum_w <= n_params as f64 + 1.0 {
            return None;
        }
        // subtract the weighted mean intensity mismatch (e.g. a constant
        // offset between the images) before estimating the variance
        let mean_r = sum_wr / sum_w;
        let sse = (sum_wr2 - sum_w * mean_r * mean_r).max(0.0);
        // q parameters plus the mean intensity offset that was removed
        let variance = sse / (sum_w - n_params as f64 - 1.0);
        if !variance.is_finite() {
            return None;
        }
        Some(ResidualFit {
            hessian,
            variance,
            n_samples,
        })
    }
}

impl<D> ObjectiveFunction<f64> for MattesMetric<D>
where
    D: Dimension,
{
    fn evaluate(&self, point: &[f64]) -> f64 {
        self.eval(point).unwrap();
        self.metric.borrow().metric
    }

    fn gradient(&self, point: &[f64]) -> Option<Vec<f64>> {
        self.eval(point).unwrap();
        Some(self.metric.borrow().derivative.clone())
    }

    /// The hessian is computed numerically: central finite differences of
    /// `gradient` with respect to each parameter (2 * n full evaluations).
    /// `MattesMetric` only accumulates first-order terms, so an analytic hessian
    /// is not available. The result is symmetrized, which also damps the small
    /// asymmetry introduced by finite differencing.
    fn hessian(&self, point: &[f64]) -> Option<Vec<Vec<f64>>> {
        let n = point.len();
        let mut hessian = vec![vec![0.0; n]; n];
        let mut p = point.to_vec();
        for j in 0..n {
            let h = 1e-5 * f64::max(1.0, point[j].abs());
            p[j] += h;
            self.eval(&p).unwrap();
            let g_plus = self.metric.borrow().derivative.clone();
            p[j] -= 2.0 * h;
            self.eval(&p).unwrap();
            let g_minus = self.metric.borrow().derivative.clone();
            p[j] += h;
            for i in 0..n {
                hessian[i][j] = (g_plus[i] - g_minus[i]) / (2.0 * h);
            }
        }
        #[expect(clippy::needless_range_loop)]
        for i in 0..n {
            for j in (i + 1)..n {
                let w = (hessian[i][j] + hessian[j][i]) / 2.0;
                hessian[i][j] = w;
                hessian[j][i] = w;
            }
        }
        Some(hessian)
    }
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub enum Sigma {
    None,
    Absolute(Vec<f64>),
    Relative(Vec<f64>),
}

impl Sigma {
    pub fn smooth<'a, A, D>(&self, array: A) -> Result<Array<f64, D>, Error>
    where
        A: AsArray<'a, f64, D>,
        D: Dimension,
    {
        match self {
            Self::None => Ok(array.into().to_owned()),
            Self::Absolute(sigma) => gaussian_smooth(array, sigma.as_slice()),
            Self::Relative(sigma) => {
                let array = array.into();
                let sigma = sigma
                    .iter()
                    .zip(array.shape())
                    .map(|(i, j)| i * (*j as f64))
                    .collect::<Vec<_>>();
                gaussian_smooth(array, sigma.as_slice())
            }
        }
    }

    pub fn sigma(&self, shape: &[usize]) -> Option<Vec<f64>> {
        match self {
            Self::None => None,
            Self::Absolute(sigma) => Some(sigma.to_vec()),
            Self::Relative(sigma) => Some(
                sigma
                    .iter()
                    .zip(shape)
                    .map(|(i, j)| i * (*j as f64))
                    .collect::<Vec<_>>(),
            ),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::filter::gaussian_smooth;
    use algos::OptimizationConfig;
    use algos::optimization::{bfgs_minimize, gradient_descent_minimize};
    use ndarray::array;

    #[test]
    fn derivative() -> Result<(), Box<dyn std::error::Error>> {
        let im = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0];
        let b = BSpline::<3, _>::new(im.view());
        let mut mem = b.get_mem();
        let mut value = Vec::new();
        let mut derivative = Vec::new();
        let x = (0..100 * im.len())
            .map(|i| i as f64 / 100.0)
            .collect::<Vec<_>>();
        for i in x.iter() {
            let (v, d) = b.evaluate_value_and_derivative_at_continuous_index(&[*i], &mut mem)?;
            value.push(v);
            derivative.push(d[0]);
        }
        println!("x = {:?}\nv = {:?}\nd = {:?}", x, value, derivative);
        Ok(())
    }

    #[test]
    fn metric2() -> Result<(), Box<dyn std::error::Error>> {
        let im_a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0];
        let im_b = array![1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0];
        let a = gaussian_smooth(im_a.view(), &[2.0])?;
        let b = gaussian_smooth(im_b.view(), &[2.0])?;

        let points = vec![
            vec![3.0],
            vec![4.0],
            vec![5.0],
            vec![6.0],
            vec![7.0],
            vec![8.0],
        ];

        let m = MattesMetric::new_from_arrays(
            a.view(),
            b.view(),
            SamplingArg::FixedAt(points),
            5,
            0.05,
        )?
        .with_fixed_mu(vec![Some(1.0), None]);
        // let mu = vec![-3.52, -3.51, -3.50, -3.49, -3.48];
        let mu = vec![-3.51, -3.49];
        // let mu = vec![-2.0, -3.0];
        let mut v = Vec::new();
        let mut d = Vec::new();
        for x in mu.iter() {
            v.push(m.evaluate(&[*x]));
            d.push(m.gradient(&[*x]).unwrap()[0]);
        }
        println!("a = {:?}\nb = {:?}", a.to_vec(), b.to_vec());
        println!("mu = {:?}\nv = {:?}\nd = {:?}", mu, v.to_vec(), d.to_vec());
        Ok(())
    }

    #[test]
    fn optimize() -> Result<(), Box<dyn std::error::Error>> {
        let im_a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0];
        let im_b = array![1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0];
        let a = gaussian_smooth(im_a.view(), &[4.0])?;
        let b = gaussian_smooth(im_b.view(), &[4.0])?;

        let m =
            MattesMetric::new_from_arrays(a.view(), b.view(), SamplingArg::Random(100), 5, 0.05)?
                .with_fixed_mu(vec![Some(1.0), None]);
        let optimization_config = OptimizationConfig {
            max_iterations: 1000,
            tolerance: 1e-6,
            learning_rate: 0.001,
        };

        let r_grad = gradient_descent_minimize(&m, &[-0.0], &optimization_config);
        let r_bfgs = bfgs_minimize(&m, &[-0.0], &optimization_config);
        println!("r_grad = {:?}", r_grad);
        println!("r_bfgs = {:?}", r_bfgs);
        Ok(())
    }

    #[test]
    fn residual_fit_1d() -> Result<(), Box<dyn std::error::Error>> {
        let a = array![1.0, 2.0, 3.0, 2.0, 1.0, 2.0, 3.0, 4.0, 2.0, 0.0, 1.0];
        // b is a shifted right by 3 (padded on the left with the first value)
        let b = array![1.0, 1.0, 1.0, 1.0, 2.0, 3.0, 2.0, 1.0, 2.0, 3.0, 4.0];
        let points = (0..11).map(|i| vec![i as f64]).collect::<Vec<_>>();
        let m = MattesMetric::new_from_arrays(
            a.view(),
            b.view(),
            SamplingArg::FixedAt(points),
            16,
            0.05,
        )?
        .with_fixed_mu(vec![Some(1.0), None]);

        // fixed(x) = moving(x + 3): no residual
        let fit = m.residual_fit(&[3.0]).expect("residual fit");
        assert!(fit.n_samples > 0);
        assert!(fit.hessian[0][0] > 0.0);
        assert!(fit.variance < 1e-12);

        // away from the optimum the residual variance must grow
        let worse = m.residual_fit(&[0.0]).expect("residual fit");
        assert!(worse.variance > fit.variance);
        Ok(())
    }

    #[test]
    fn hessian_1d() -> Result<(), Box<dyn std::error::Error>> {
        let im_a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0];
        let im_b = array![1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0];
        let a = gaussian_smooth(im_a.view(), &[2.0])?;
        let b = gaussian_smooth(im_b.view(), &[2.0])?;

        let points = vec![
            vec![3.0],
            vec![4.0],
            vec![5.0],
            vec![6.0],
            vec![7.0],
            vec![8.0],
        ];

        let m = MattesMetric::new_from_arrays(
            a.view(),
            b.view(),
            SamplingArg::FixedAt(points),
            5,
            0.05,
        )?
        .with_fixed_mu(vec![Some(1.0), None]);

        let p = -3.49;
        let h = m.hessian(&[p]).unwrap();
        assert_eq!(h.len(), 1);
        assert_eq!(h[0].len(), 1);

        // cross-check against the second central difference of `evaluate`
        let step = 1e-4;
        let fd = (m.evaluate(&[p + step]) - 2.0 * m.evaluate(&[p]) + m.evaluate(&[p - step]))
            / (step * step);
        println!("hessian = {:?}\nfinite diff = {:?}", h[0][0], fd);
        assert!((h[0][0] - fd).abs() < 1e-3 + 1e-2 * fd.abs());
        Ok(())
    }

    #[test]
    fn hessian_2d() -> Result<(), Box<dyn std::error::Error>> {
        let a = Array2::from_shape_fn((12, 12), |(i, j)| i as f64 + 0.7 * j as f64);
        let b = Array2::from_shape_fn((12, 12), |(i, j)| 0.5 * i as f64 + 24.0 - 0.7 * j as f64);
        let a = gaussian_smooth(a.view(), &[2.0, 2.0])?;
        let b = gaussian_smooth(b.view(), &[2.0, 2.0])?;

        let points = (0..12)
            .flat_map(|i| (0..12).map(move |j| vec![i as f64, j as f64]))
            .step_by(2)
            .collect::<Vec<_>>();

        let m = MattesMetric::new_from_arrays(
            a.view(),
            b.view(),
            SamplingArg::FixedAt(points),
            7,
            0.05,
        )?;
        let n = m.fixed_mu().len_variable();
        assert_eq!(n, 2 * 2 + 2);

        let p = vec![0.5, -0.3, 0.2, 0.7, -0.4, 0.3];
        let h = m.hessian(&p).unwrap();
        assert_eq!(h.len(), n);
        assert!(h.iter().all(|row| row.len() == n));
        // symmetrized by construction
        for (i, row) in h.iter().enumerate() {
            for (j, v) in row.iter().enumerate() {
                assert_eq!(*v, h[j][i]);
            }
        }

        // cross-check each diagonal entry against the 2nd central difference of `evaluate`
        let step = 1e-4;
        for j in 0..n {
            let mut pp = p.clone();
            let mut pm = p.clone();
            pp[j] += step;
            pm[j] -= step;
            let fd = (m.evaluate(&pp) - 2.0 * m.evaluate(&p) + m.evaluate(&pm)) / (step * step);
            println!("diagonal[{j}] = {:?}\nfinite diff   = {:?}", h[j][j], fd);
            assert!(
                (h[j][j] - fd).abs() < 1e-3 + 1e-2 * fd.abs(),
                "diagonal entry {j}"
            );
        }
        Ok(())
    }

    #[test]
    fn test_square_bspline_integral() -> Result<(), Box<dyn std::error::Error>> {
        let x = Array1::linspace(-2.0, 2.0, 500);
        let y = x.mapv(square_bspline);
        let z = x.mapv(square_bspline_integral);
        let mut w = vec![0.0];
        let mut j = 0.0;
        let dx = (x[x.len() - 1] - x[0]) / (x.len() - 1) as f64;
        for i in y.windows([2]) {
            j += dx * (i[0] + i[1]) / 2.0;
            w.push(j);
        }
        let w = Array1::from(w);
        assert!((w - z).abs().iter().all(|i| *i < 1e-5));
        Ok(())
    }

    #[test]
    fn test_cubic_bspline_derivative() -> Result<(), Box<dyn std::error::Error>> {
        let x = Array1::linspace(-2.0, 2.0, 500);
        let y = x.mapv(cubic_bspline);
        let z = x.mapv(|i| square_bspline(i + 0.5) - square_bspline(i - 0.5));
        let mut w = vec![0.0];
        let mut j = 0.0;
        let dx = (x[x.len() - 1] - x[0]) / (x.len() - 1) as f64;
        for i in z.windows([2]) {
            j += dx * (i[0] + i[1]) / 2.0;
            w.push(j);
        }
        let w = Array1::from(w);
        assert!((w - y).abs().iter().all(|i| *i < 1e-4));
        Ok(())
    }
}