plotpoint 0.1.9

use levenberg_marquardt::LeastSquaresProblem;
use nalgebra::{Affine2, DMatrix, DVector, Dyn, Owned, Point2, RealField};
use num_dual::{DualNum, DualNumFloat, DualVec, try_jacobian};
use num_traits::{Float, FloatConst, NumCast};

use crate::image::Image;
pub type DV<T> = DualVec<T, T, Dyn>;
//suggestion: use U: IntoIterator<Item=(Point2<DV<T>>, DV<T>)>
pub trait ParamsToPoints<T: DualNum<T> + DualNumFloat> =
    Fn(&[DV<T>]) -> Option<Vec<(Point2<DV<T>>, DV<T>)>>;
pub struct PointProblem<T, F>
where
    T: DualNum<T, Inner = T> + RealField + Float + FloatConst, //T: f32 or f64
    F: ParamsToPoints<T>,                                      //todo: swith to Result + ?operator
{
    pub params: DVector<T>,
    target_image: Image<T, T>,
    pix_to_point: Affine2<DV<T>>,
    params_to_points: Box<F>,

    cached_f: Option<DVector<T>>,
    cached_jac: Option<DMatrix<T>>,
}

impl<T, F> PointProblem<T, F>
where
    T: DualNum<T, Inner = T> + DualNumFloat + RealField,
    F: ParamsToPoints<T>, //todo: define type
{
    pub fn new(
        params: &DVector<T>,
        image: Image<T, T>,
        params_to_points: F,
        pix_to_point: Affine2<T>,
    ) -> Self {
        let pix_to_point =
            Affine2::from_matrix_unchecked(pix_to_point.matrix().map(|f| DualVec::from_inner(f)));
        let mut s = Self {
            target_image: image,
            cached_f: None,
            cached_jac: None,
            params: params.clone(),
            params_to_points: Box::new(params_to_points),
            pix_to_point,
        };
        s.set_jacobian(params);
        s
    }
    fn set_jacobian(&mut self, param: &DVector<T>) {
        if let Ok((f, j)) = try_jacobian(
            |param| {
                let res = self.residuals(param.as_slice());
                if res.is_none() {
                    return Err("residuals failed".to_string());
                }
                Ok(DVector::<DV<T>>::from_column_slice(
                    res.unwrap().0.as_slice(),
                ))
            },
            param.clone(),
        ) {
            // println!( //todo: member function?
            //     "obj: {}",
            //     f.clone().norm_squared() * T::from_f64(0.5).unwrap()
            // );
            self.cached_jac = Some(j);
            self.cached_f = Some(f);
        } else {
            self.cached_jac = None;
            self.cached_f = None;
        }
    }
    fn residuals(&self, params: &[DV<T>]) -> Option<Image<T, DV<T>>> {
        let point_w = (self.params_to_points)(params)?;
        let background = DMatrix::from_element(
            //todo: background should be a parameter
            self.target_image.0.nrows(),
            self.target_image.0.ncols(),
            <T as NumCast>::from(1.0).unwrap(),
        );
        let mut im = Image::from_points(&point_w, &self.pix_to_point, &background);
        im.0.zip_apply(&self.target_image.0, |d, f| *d -= f);
        Some(im)
    }
}

impl<T, F> LeastSquaresProblem<T, Dyn, Dyn> for PointProblem<T, F>
where
    T: DualNum<T, Inner = T>
        + std::marker::Copy
        + nalgebra::ComplexField
        + RealField
        + Float
        + DualNumFloat,
    F: ParamsToPoints<T>,
{
    type ResidualStorage = Owned<T, Dyn>; //todo: try borrowed

    type JacobianStorage = Owned<T, Dyn, Dyn>;

    type ParameterStorage = Owned<T, Dyn>;

    fn set_params(&mut self, x: &DVector<T>) {
        // called last int the loop, compute residuals and jacobian for next loop
        self.params = x.clone();
        self.set_jacobian(x);
    }

    fn params(&self) -> DVector<T> {
        // first to get called, and only once.
        self.params.clone()
    }

    fn residuals(&self) -> Option<DVector<T>> {
        // first called in the loop, values first computed in new(), then in set_params()
        // always called first after set_params()
        self.cached_f.clone()
    }

    fn jacobian(&self) -> Option<nalgebra::DMatrix<T>> {
        // called after residuals(), but sometimes skipped
        self.cached_jac.clone()
    }
}

#[cfg(test)]
mod tests {
    use std::ops::Mul;

    use super::*;
    use approx::assert_abs_diff_eq;
    use levenberg_marquardt::LevenbergMarquardt;
    use nalgebra::{DVector, Point2};

    fn make_point_problem<T>() -> PointProblem<T, impl ParamsToPoints<T>>
    where
        T: DualNum<T, Inner = T> + DualNumFloat + RealField + NumCast,
        for<'a> &'a T: Mul<&'a T, Output = T>,
    {
        let n = T::one();
        let z = T::zero();
        let point = Point2::new(n, n);
        let image = Image::new(DMatrix::from_column_slice(
            3,
            3,
            &[n, n, n, n, z, n, n, n, n],
        ));

        let w = DV::from_inner(T::from_f64(0.01).unwrap());
        PointProblem::new(
            &DVector::from_column_slice(&[point.x, point.y]),
            image,
            move |p: &[DV<T>]| Some(vec![(Point2::new(p[0].clone(), p[1].clone()), w.clone())]),
            Affine2::identity(),
        )
    }

    #[test]
    fn test_point_problem_f64() {
        let problem = make_point_problem::<f64>();
        let solver = LevenbergMarquardt::new();
        let (result, report) = solver.minimize(problem);
        assert_abs_diff_eq!(result.params()[0], 1.5, epsilon = 1e-6); //center of the image
        assert_abs_diff_eq!(result.params()[1], 1.5, epsilon = 1e-6);
        assert!(matches!(
            report.termination,
            levenberg_marquardt::TerminationReason::Converged { .. }
        ));
    }

    #[test]
    fn test_point_problem_f32() {
        let problem = make_point_problem::<f32>();
        let solver = LevenbergMarquardt::new();
        let (result, report) = solver.minimize(problem);
        assert_abs_diff_eq!(result.params()[0], 1.5, epsilon = 1e-5); //center of the image
        assert_abs_diff_eq!(result.params()[1], 1.5, epsilon = 1e-5);
        assert!(matches!(
            report.termination,
            levenberg_marquardt::TerminationReason::Converged { .. }
        ));
    }
}