// Copyright (c) 2019 Autonomous Robots and Cognitive Systems Laboratory
// Author: Daniel Garcia-Vaglio <degv364@gmail.com>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

use nalgebra::{MatrixMN, U6, Dynamic, Matrix3};
use crate::geometry::{Vector, Twist, Frame, new_ref_point, new_ref_base,
                   new_ref_frame, Introspection};
use crate::joint::{JointType};
use crate::chains::{Chain};

/// The jacobian stores a twist for each joint of a kinematic chain
///
/// Jacobians are matrices of the form 6xn where n is the number of
/// joints of a certain kinematic chain.
pub type Jacobian = MatrixMN<f64, U6, Dynamic>;

/// Obtain a jacobian from a vector.
pub fn vec_to_jacobian(ncols: usize, vector: Vec<f64>) -> Jacobian {
    // The vector is row major and the Jacobian is column mayor.
    // So the transpose of the vector has to be taken.
    Jacobian::from_fn(ncols, |i, j| vector[(i*ncols+j) % (ncols * 6)])
}

/// Operations that are computed per-joint.
pub trait JointOperations {
    /// Get the twist represented by the n-th joint
    fn get_joint_twist(&self, joint_index: usize) -> Twist;
    /// Change the reference point of each twist
    fn new_ref_point(&mut self, new_ref: Vector);
    /// Change the reference rotation of each twist
    fn new_ref_base(&mut self, new_ref: Matrix3<f64>);
    /// Change the reference frame of each twist
    fn new_ref_frame(&mut self, new_ref: Frame);
}

impl JointOperations for Jacobian {
    fn get_joint_twist(&self, joint_index: usize) -> Twist{
        let num_cols = self.ncols();
        if num_cols <= joint_index {
            panic!(
                "Requested Joint {}, but Jacobian only has {}",
                joint_index, num_cols);
        }
        let col = self.column(joint_index);
        Twist::from_columns(&[col])
    }

    fn new_ref_point(&mut self, new_ref: Vector) {
        let num_cols = self.ncols();
        for col in 0..num_cols {
            self.set_column(
                col, &new_ref_point(self.get_joint_twist(col), new_ref));
        }
    }
    fn new_ref_base(&mut self, new_ref: Matrix3<f64>){
        let num_cols = self.ncols();
        for col in 0..num_cols {
            self.set_column(
                col, &new_ref_base(self.get_joint_twist(col), new_ref));
        }
    }
    fn new_ref_frame(&mut self, new_ref: Frame){
        let num_cols = self.ncols();
        for col in 0..num_cols {
            self.set_column(
                col, &new_ref_frame(self.get_joint_twist(col), new_ref));
        }
    }
}

/// Operations specific to solvers
pub trait JacobianSolver {
    /// Calculate the jacobian of a chain at a specific joint-space configuration
    ///
    /// `chain`: chain to use by the solver\
    /// `angles`: vector holding the current state of each joint in the chain
    fn solve_from_chain(&mut self, chain: &Chain, angles: &Vec<f64>);
}

impl JacobianSolver for Jacobian{
    fn solve_from_chain(&mut self, chain: &Chain, angles: &Vec<f64>){
        let num_joints = chain.get_num_ind_joints();
        if angles.len() != num_joints {
            panic!("Unable to compute Jacobian of a chain with {} \
                    joints from {} angles",
                   num_joints, angles.len());
        }
        if num_joints != self.ncols() {
            panic!("Unable to store jacobian of a chain with {} \
                    joints in a matrix with {} columns",
                   num_joints, self.ncols());
        }
        let mut pose_tmp = Frame::identity();
        let mut pose_total;
        
        let mut jac_index = 0;
        let mut joint_index = 0;
        
        for segment_index in 0..chain.get_num_segments() {
            let segment = chain.get_segment(segment_index);
            // Calculate forward kinematics to change jacobian reference
            pose_total = pose_tmp * segment.pose(angles[joint_index]);
            
            self.new_ref_point(
                pose_total.get_translation() - pose_tmp.get_translation());
            
            //Update the jacobian
            if segment.get_joint_type() != JointType::NoJoint {
                if !chain.get_locked_joints()[joint_index] {
                    let twist = new_ref_base(
                        segment.twist(angles[joint_index], 1.0),
                        pose_tmp.get_rotation());
                    self.set_column(jac_index, &twist);
                    jac_index += 1;
                }
                joint_index += 1;
            }
            pose_tmp = pose_total;
        }
        self.clone().mul_to(&chain.get_coupling_matrix(), self);
    }
}

/// Calculate the absolute difference between two jacobians
///
/// Because jacobians can be very big, we calculate the mean of the
/// per-twist error
pub fn get_jacobian_error(left: Jacobian, right: Jacobian) -> f64 {
    let num_joints = left.ncols();
    if num_joints != right.ncols() {
        panic!("Unable to comapre Jacobians of different sizes. Got {} and {}",
        num_joints, right.ncols());
    }
    let error = (left-right).abs();
    let total = num_joints as f64;
    // Because Jacobians can be very big, we use the mean error (per twist),
    // not the total error.
    error.fold(0.0, |a:f64, b:f64| a+b)/total
}

#[cfg(test)]
mod tests {
    use nalgebra::{Matrix3};
    use crate::geometry::{Vector, Frame, Twist};
    use crate::jacobian::{Jacobian, JointOperations, JacobianSolver,
                          vec_to_jacobian, get_jacobian_error};
    use crate::chains::tests::{create_testing_chain};
    #[test]
    fn get_twist_from_jacobian() {
        let jacobian = vec_to_jacobian(3, vec![
            0.419164, 0.818984, 0.707018,
            0.41915, 0.421392, 0.753147,
            0.865399, 0.119757, 0.30401,
            0.65547, 0.674511, 0.386427,
            0.647384, 0.251815, 0.908477,
            0.50629, 0.060354, 0.482741]);
        let first_twist = Twist::new(
            0.419164, 0.41915, 0.865399, 0.65547, 0.647384, 0.50629);
        let second_twist = Twist::new(
            0.818984, 0.421392, 0.119757, 0.674511, 0.251815, 0.060354);
        let third_twist = Twist::new(
            0.707018, 0.753147, 0.30401, 0.386427, 0.908477, 0.482741);
        assert_eq!(first_twist, jacobian.get_joint_twist(0));
        assert_eq!(second_twist, jacobian.get_joint_twist(1));
        assert_eq!(third_twist, jacobian.get_joint_twist(2));
    }

    #[test]
    fn change_ref_point() {
        let point = Vector::new(0.1, 0.5, 6.0);
        let mut jacobian = vec_to_jacobian(4, vec![
            0.419164, 0.818984, 0.707018, 0.01,
            0.41915, 0.421392, 0.753147, 0.01,
            0.865399, 0.119757, 0.30401, 0.01,
            0.65547, 0.674511, 0.386427, 0.01,
            0.647384, 0.251815, 0.908477, 0.01,
            0.50629, 0.060354, 0.482741, 0.01]);
        let expected = vec_to_jacobian(4, vec![
            4.05032, 2.2997, 5.91651, 0.065,
            -3.46304, -3.61964, -1.51714, -0.049,
            1.1284, 0.431831, 0.406376, 0.014,
            0.65547, 0.674511, 0.386427, 0.01,
            0.647384, 0.251815, 0.908477, 0.01,
            0.50629, 0.060354, 0.482741, 0.01]);
        jacobian.new_ref_point(point);
        let error = get_jacobian_error(jacobian, expected);
        assert!(error < 0.0001)
    }
    #[test]
    fn change_ref_base() {
        let rotation = Matrix3::new(
            0.842629, 0.323976, 0.430135,
            -0.24521, 0.942, -0.229147,
            -0.479426, 0.0876121, 0.873198);
        let mut jacobian = vec_to_jacobian(4, vec![
            0.419164, 0.818984, 0.707018, 0.01,
            0.41915, 0.421392, 0.753147, 0.01,
            0.865399, 0.119757, 0.30401, 0.01,
            0.65547, 0.674511, 0.386427, 0.01,
            0.647384, 0.251815, 0.908477, 0.01,
            0.50629, 0.060354, 0.482741, 0.01]);
        let expected = vec_to_jacobian(4, vec![
            0.861233, 0.878132, 0.970521, 0.0159674,
            0.0937526, 0.168686, 0.466433, 0.00467643,
            0.59143, -0.251151, -0.0075167, 0.00481385,
            0.979828, 0.675905, 0.827583, 0.0159674,
            0.333093, 0.0579828, 0.650411, 0.00467643,
            0.184561, -0.248615, 0.315859, 0.00481385]);
        jacobian.new_ref_base(rotation);
        let error = get_jacobian_error(jacobian, expected);
        assert!(error < 0.0001);
    }
    #[test]
    fn change_ref_frame() {
        let frame = Frame::new(
            0.842629, 0.323976, 0.430135, 0.1,
            -0.24521, 0.942, -0.229147, 0.5,
            -0.479426, 0.0876121, 0.873198, 6.0,
            0.0, 0.0, 0.0, 1.0);
        let mut jacobian = vec_to_jacobian(4, vec![
            0.419164, 0.818984, 0.707018, 0.01,
            0.41915, 0.421392, 0.753147, 0.01,
            0.865399, 0.119757, 0.30401, 0.01,
            0.65547, 0.674511, 0.386427, 0.01,
            0.647384, 0.251815, 0.908477, 0.01,
            0.50629, 0.060354, 0.482741, 0.01]);
        let expected = vec_to_jacobian(4, vec![
            -1.04504, 0.405928, -2.77401, -0.00968425,
            5.95426, 4.24898, 5.40035, 0.0999994,
            0.134825, -0.583305, -0.356267, -0.00270221,
            0.979828, 0.675905, 0.827583, 0.0159674,
            0.333093, 0.0579828, 0.650411, 0.00467643,
            0.184561, -0.248615, 0.315859, 0.00481385]);
        jacobian.new_ref_frame(frame);
        let error = get_jacobian_error(jacobian, expected);
        assert!(error < 0.0001);
    }

    #[test]
    fn jacobian_from_chain_correct() {
        let chain = create_testing_chain();
        let angles = vec![0.1, -0.95, 0.57, 0.68, -0.27, 0.39, 0.47];
        let mut result = Jacobian::zeros(angles.len());
        result.solve_from_chain(&chain, &angles);
        let expected = vec_to_jacobian(7, vec![
            -0.330827, 0.261285, -0.0634603, -0.640334, -0.44059, -0.00594982,
            -0.20031,
            -1.48783, -1.80195, -0.39123, -0.911489, -0.717405, 0.675781,
            0.0976329,
            -0.0176899, 0.160687, -0.0226411, -0.952894, -0.85285, 0.0807787,
            0.0185394,
            -0.218351, 0.852363, -0.387827, 0.629791, 0.632632, 0.971581,
            0.133574,
            0.036957, 0.167762, 0.00956914, -0.727376, -0.721899, 0.0365101,
            0.0865809,
            0.97517, 0.495311, 0.921682, 0.272557, 0.280427, -0.233875, 0.98725
        ]);
        let error = get_jacobian_error(result, expected);
        assert!(error < 0.0001);
    }
}