sophus_sensor 0.15.0

use core::{
    borrow::Borrow,
    marker::PhantomData,
};

use log::warn;
use sophus_autodiff::{
    linalg::EPS_F64,
    params::IsParamsImpl,
};

use crate::{
    distortions::affine::AffineDistortionImpl,
    prelude::*,
    traits::IsCameraDistortionImpl,
};

extern crate alloc;

/// Kannala-Brandt distortion implementation
#[derive(Debug, Clone, Copy)]
pub struct BrownConradyDistortionImpl<
    S: IsScalar<BATCH, DM, DN>,
    const BATCH: usize,
    const DM: usize,
    const DN: usize,
> {
    phantom: PhantomData<S>,
}

impl<S: IsScalar<BATCH, DM, DN>, const BATCH: usize, const DM: usize, const DN: usize>
    IsParamsImpl<S, 12, BATCH, DM, DN> for BrownConradyDistortionImpl<S, BATCH, DM, DN>
{
    fn are_params_valid(_params: S::Vector<12>) -> S::Mask {
        S::Mask::all_true()
    }

    fn params_examples() -> alloc::vec::Vec<S::Vector<12>> {
        alloc::vec![S::Vector::<12>::from_f64_array([
            286.0,
            286.0,
            424.0,
            400.0,
            0.726405,
            -0.0148413,
            1.38447e-05,
            0.000419742,
            -0.00514224,
            1.06774,
            0.128429,
            -0.019901,
        ])]
    }

    fn invalid_params_examples() -> alloc::vec::Vec<S::Vector<12>> {
        alloc::vec![]
    }
}

impl<S: IsScalar<BATCH, DM, DN>, const BATCH: usize, const DM: usize, const DN: usize>
    BrownConradyDistortionImpl<S, BATCH, DM, DN>
{
    fn distortion_impl(
        params: &S::Vector<8>,
        proj_point_in_camera_z1_plane: &S::Vector<2>,
    ) -> S::Vector<2> {
        let x = proj_point_in_camera_z1_plane.elem(0);
        let y = proj_point_in_camera_z1_plane.elem(1);
        let d0 = params.elem(0);
        let d1 = params.elem(1);
        let d2 = params.elem(2);
        let k3 = params.elem(3);
        let d4 = params.elem(4);
        let d5 = params.elem(5);
        let d6 = params.elem(6);
        let d7 = params.elem(7);

        // From:
        // https://github.com/opencv/opencv/blob/63bb2abadab875fc648a572faccafee134f06fc8/modules/calib3d/src/calibration.cpp#L791

        let one = S::from_f64(1.0);
        let two = S::from_f64(2.0);

        let r2 = proj_point_in_camera_z1_plane.squared_norm();
        let r4 = r2 * r2;
        let r6 = r4 * r2;
        let a1 = two * x * y;
        let a2 = r2 + two * x * x;
        let a3 = r2 + two * y * y;
        let cdist = one + d0 * r2 + d1 * r4 + d4 * r6;
        let icdist2 = one / (one + d5 * r2 + d6 * r4 + d7 * r6);
        let xd0 = x * cdist * icdist2 + d2 * a1 + k3 * a2;
        let yd0 = y * cdist * icdist2 + d2 * a3 + k3 * a1;

        S::Vector::<2>::from_array([xd0, yd0])
    }

    fn undistort_impl(
        distortion_params: &S::Vector<8>,
        distorted_point: &S::Vector<2>,
        dbg_info_distorted_point: &S::Vector<2>,
    ) -> S::Vector<2> {
        // from https://github.com/farm-ng/farm-ng-core/blob/main/cpp/sophus/sensor/camera_distortion/brown_conrady.h

        // We had no luck with OpenCV's undistort. It seems not to be accurate if
        // "icdist" is close to 0.
        // https://github.com/opencv/opencv/blob/63bb2abadab875fc648a572faccafee134f06fc8/modules/calib3d/src/undistort.dispatch.cpp#L365
        //
        // Hence, we derive the inverse approximation scheme from scratch.
        //
        //
        // Objective: find that xy such that proj_impl(xy) = uv
        //
        // Using multivariate Newton scheme, by defining f and find the root of it:
        //
        //  f: R^2 -> R^2
        //  f(xy) :=  proj_impl(xy) - uv
        //
        //  xy_i+1 = xy_i + J^{-1} * f(xy)   with J being the Jacobian of f.
        //
        // TODO(hauke): There is most likely a 1-dimensional embedding and one only
        // need to solve a less computational heavy newton iteration...

        // initial guess
        let mut xy = *distorted_point;

        let d0 = distortion_params.elem(0);
        let d1 = distortion_params.elem(1);
        let d2 = distortion_params.elem(2);
        let k3 = distortion_params.elem(3);
        let d4 = distortion_params.elem(4);
        let d5 = distortion_params.elem(5);
        let d6 = distortion_params.elem(6);
        let d7 = distortion_params.elem(7);

        for _i in 0..50 {
            let x = xy.elem(0);
            let y = xy.elem(1);

            let f_xy = Self::distortion_impl(distortion_params, &xy.clone()) - *distorted_point;

            let du_dx;
            let du_dy;
            let dv_dx;
            let dv_dy;

            {
                // Generated by brown_conrady_camera.py
                let a = x;
                let b = y;

                let c0 = a * a; // pow(a, 2);
                let c1 = b * b; // pow(b, 2);
                let c2 = c0 + c1;
                let c3 = c2 * c2; // pow(c2, 2);
                let c4 = c3 * c2; // pow(c2, 3);
                let c5 = c2 * d5 + c3 * d6 + c4 * d7 + S::from_f64(1.0);
                let c6 = c5 * c5; // pow(c5, 2);
                let c7 = S::from_f64(1.0) / c6;
                let c8 = a * k3;
                let c9 = S::from_f64(2.0) * d2;
                let c10 = S::from_f64(2.0) * c2;
                let c11 = S::from_f64(3.0) * c3;
                let c12 = c2 * d0;
                let c13 = c3 * d1;
                let c14 = c4 * d4;
                let c15 = S::from_f64(2.0)
                    * (c10 * d6 + c11 * d7 + d5)
                    * (c12 + c13 + c14 + S::from_f64(1.0));
                let c16 = S::from_f64(2.0) * c10 * d1
                    + S::from_f64(2.0) * c11 * d4
                    + S::from_f64(2.0) * d0;
                let c17 = c12 + c13 + c14 + S::from_f64(1.0);
                let c18 = b * k3;
                let c19 = a * b;
                let c20 = -c15 * c19 + c16 * c19 * c5;
                du_dx = c7
                    * (-c0 * c15 + c5 * (c0 * c16 + c17) + c6 * (b * c9 + S::from_f64(6.0) * c8));
                du_dy = c7 * (c20 + c6 * (a * c9 + S::from_f64(2.0) * c18));
                dv_dx = c7 * (c20 + c6 * (S::from_f64(2.0) * a * d2 + S::from_f64(2.0) * c18));
                dv_dy = c7
                    * (-c1 * c15
                        + c5 * (c1 * c16 + c17)
                        + c6 * (S::from_f64(6.0) * b * d2 + S::from_f64(2.0) * c8));
            }

            //     | du_dx  du_dy |      | a  b |
            // J = |              |  =:  |      |
            //     | dv_dx  dv_dy |      | c  d |

            let a = du_dx;
            let b = du_dy;
            let c = dv_dx;
            let d = dv_dy;

            // | a  b | -1       1   |  d  -b |
            // |      |     =  ----- |        |
            // | c  d |        ad-bc | -c   a |

            let m: S::Matrix<2, 2> = S::Matrix::from_array2([[d, -b], [-c, a]]);

            let j_inv: S::Matrix<2, 2> = m.scaled(S::from_f64(1.0) / (a * d - b * c));
            let step: S::Vector<2> = j_inv * f_xy;

            if step
                .norm()
                .real_part()
                .less_equal(&S::RealScalar::from_f64(EPS_F64))
                .all()
            {
                break;
            }

            *xy.elem_mut(0) = xy.elem(0) - step.elem(0);
            *xy.elem_mut(1) = xy.elem(1) - step.elem(1);
        }

        let f_xy = Self::distortion_impl(distortion_params, &xy.clone()) - *distorted_point;
        if !f_xy
            .norm()
            .real_part()
            .less_equal(&S::RealScalar::from_f64(EPS_F64))
            .any()
        {
            warn!("Newton did not converge: f_xy: {f_xy:?}: pixel {dbg_info_distorted_point:?}");
        }
        xy
    }
}

impl<S: IsScalar<BATCH, DM, DN>, const BATCH: usize, const DM: usize, const DN: usize>
    IsCameraDistortionImpl<S, 8, 12, BATCH, DM, DN>
    for BrownConradyDistortionImpl<S, BATCH, DM, DN>
{
    fn distort<PA, PO>(params: PA, proj_point_in_camera_z1_plane: PO) -> S::Vector<2>
    where
        PA: Borrow<S::Vector<12>>,
        PO: Borrow<S::Vector<2>>,
    {
        let params = params.borrow();
        let proj_point_in_camera_z1_plane = proj_point_in_camera_z1_plane.borrow();

        let distorted_point_in_camera_z1_plane =
            Self::distortion_impl(&params.get_fixed_subvec(4), proj_point_in_camera_z1_plane);

        AffineDistortionImpl::<S, BATCH, DM, DN>::distort(
            params.get_fixed_subvec(0),
            distorted_point_in_camera_z1_plane,
        )
    }

    fn undistort<PA, PO>(params: PA, distorted_point: PO) -> S::Vector<2>
    where
        PA: Borrow<S::Vector<12>>,
        PO: Borrow<S::Vector<2>>,
    {
        let params = params.borrow();
        let distorted_point = distorted_point.borrow();
        let undistorted_point_in_camera_z1_plane =
            AffineDistortionImpl::<S, BATCH, DM, DN>::undistort(
                params.get_fixed_subvec(0),
                distorted_point,
            );

        Self::undistort_impl(
            &params.get_fixed_subvec(4),
            &undistorted_point_in_camera_z1_plane,
            distorted_point,
        )
    }

    fn dx_distort_x<PA, PO>(params: PA, proj_point_in_camera_z1_plane: PO) -> S::Matrix<2, 2>
    where
        PA: Borrow<S::Vector<12>>,
        PO: Borrow<S::Vector<2>>,
    {
        let params = params.borrow();
        let proj_point_in_camera_z1_plane = proj_point_in_camera_z1_plane.borrow();

        const OFFSET: usize = 4;

        let fx = params.elem(0);
        let fy = params.elem(1);
        let d0 = params.elem(OFFSET);
        let d1 = params.elem(1 + OFFSET);
        let d2 = params.elem(2 + OFFSET);
        let d3 = params.elem(3 + OFFSET);
        let d4 = params.elem(4 + OFFSET);
        let d5 = params.elem(5 + OFFSET);
        let d6 = params.elem(6 + OFFSET);
        let d7 = params.elem(7 + OFFSET);

        let one = S::from_f64(1.0);
        let two = S::from_f64(2.0);

        let a = proj_point_in_camera_z1_plane.elem(0);
        let b = proj_point_in_camera_z1_plane.elem(1);

        // Generated by brown_conrady_camera.py
        let c0 = a * d3;
        let c1 = two * d2;
        let c2 = a * a; // pow(a, 2);
        let c3 = b * b; // pow(b, 2);
        let c4 = c2 + c3;
        let c5 = c4 * c4; // pow(c4, 2);
        let c6 = c5 * c4; // pow(c4, 3);
        let c7 = c4 * d5 + c5 * d6 + c6 * d7 + one;
        let c8 = c7 * c7; // pow(c7, 2);
        let c9 = two * c4;
        let c10 = S::from_f64(3.0) * c5;
        let c11 = c4 * d0;
        let c12 = c5 * d1;
        let c13 = c6 * d4;
        let c14 = two * (c10 * d7 + c9 * d6 + d5) * (c11 + c12 + c13 + one);
        let c15 = two * c10 * d4 + two * c9 * d1 + two * d0;
        let c16 = one * c11 + c12 + c13 + one;
        let c17 = one / c8;
        let c18 = c17 * fx;
        let c19 = b * d3;
        let c20 = a * b;
        let c21 = -c14 * c20 + c15 * c20 * c7;
        let c22 = c17 * fy;

        let dfx_dfx =
            c18 * (-c14 * c2 + c7 * (c15 * c2 + c16) + c8 * (b * c1 + S::from_f64(6.0) * c0));
        let dfy_dfx = c22 * (c21 + c8 * (two * a * d2 + two * c19));

        let dfx_dfy = c18 * (c21 + c8 * (a * c1 + two * c19));
        let dfy_dfy =
            c22 * (-c14 * c3 + c7 * (c15 * c3 + c16) + c8 * (S::from_f64(6.0) * b * d2 + two * c0));

        S::Matrix::from_array2([[dfx_dfx, dfx_dfy], [dfy_dfx, dfy_dfy]])
    }

    fn dx_distort_params<PA, PO>(
        _params: PA,
        _proj_point_in_camera_z1_plane: PO,
    ) -> S::Matrix<2, 12>
    where
        PA: Borrow<S::Vector<12>>,
        PO: Borrow<S::Vector<2>>,
    {
        todo!()
    }
}