// 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::{Matrix6, DMatrix, DVector, U6};
use crate::geometry::{Twist};
use crate::chains::{Chain};
use crate::jacobian::{Jacobian, JacobianSolver};
use crate::svd_eigen::{SvdResult};

/// Matrix for holding task space weights
pub type WeightTaskSpace = Matrix6<f64>;
/// Matrix for holding joint space weights
pub type WeightJointSpace = DMatrix<f64>;
type JointQdots = DVector<f64>;


fn transform_vectors(in_vector: JointQdots) -> Vec<f64> {
    let mut result = Vec::with_capacity(in_vector.nrows());
    for i in 0..in_vector.nrows() {
        result.push(in_vector[i]);
    }
    result
}


/// Solver for inverse differential kinematcs. Calculate the joints
/// rates from their current position and desired velocity at end-effector
///
/// `chain`: chain to use by the solver\
/// `weight_task_space`: Weights in the task space\
/// `weight_joint_space`: How much to use each joint\
/// `jacobian`: Jacobian of the current chain\
/// `temp_qdots`: Temporal storage of qdot values\
/// `svd_result`: For calculating SVD\
/// `lambda`: parameter for tresholding singular values\
/// `epsilon`: (1e-300) allowed error in convergence of singular values\
/// `maxiter`: (150)Max allowed iteration for convergence of SV\
/// `n_joints`: Number of joints in the kinematic chain\
pub struct InverseDiffKinematicsSolver {
    chain: Chain,
    weight_task_space: WeightTaskSpace,
    weight_joint_space: WeightJointSpace,
    jacobian: Jacobian,
    temp_qdot: JointQdots,
    svd_result: SvdResult,
    lambda: f64,
    epsilon: f64,
    maxiter: usize,
    n_joints: usize,
}

impl InverseDiffKinematicsSolver{
    /// Create a new solver from a given chain
    pub fn new(in_chain: Chain) -> Self{
        let num_joints = in_chain.get_num_joints();
        Self{
            chain: in_chain,
            weight_task_space: WeightTaskSpace::identity(),
            weight_joint_space: WeightJointSpace::identity(
                num_joints, num_joints),
            jacobian: Jacobian::zeros(num_joints),
            temp_qdot: JointQdots::zeros(num_joints),
            svd_result: SvdResult::new(num_joints),
            lambda: 0.0,
            epsilon: 1e-300,
            maxiter: 150,
            n_joints: num_joints,
        }
    }

    /// Set the weight matrices for inverse kinematics
    pub fn set_weights(&mut self,
                       task_space: WeightTaskSpace,
                       joint_space: WeightJointSpace) {
        if joint_space.ncols() != joint_space.nrows() {
            panic!("Joint space matrix must be square. Got {}x{}",
            joint_space.nrows(), joint_space.ncols());
        }
        if joint_space.ncols() != self.n_joints {
            panic!("Wrong joint space matrix size. Expected {},\
                    got {}", self.n_joints, joint_space.nrows());
        }
        self.weight_task_space = task_space;
        self.weight_joint_space = joint_space;
    }

    /// Set the Lambda parameter
    pub fn set_lambda(&mut self, new_lambda: f64) {
        self.lambda = new_lambda;
    }

    /// Set parameters for convergence algorithm
    pub fn set_convergence(&mut self, new_epsilon: f64, new_maxiter: usize){
        if new_epsilon <= 0.0 {
            panic!("Epsilon must be positive! got {}", new_epsilon);
        }
        self.epsilon = new_epsilon;
        self.maxiter = new_maxiter;
    }
    
    /// Solver for inverse differential kinematics
    ///
    /// This function uses the weighted damped least square inverse kinematics
    /// algorithm. From a Twist and current joint state, it calculates the
    /// velocities at each joint.
    ///
    /// `twist`: desired twist at the end effector\
    /// `init_angles`: current state of each joint\
    /// Returns a vector with the velocities at each joint.
    pub fn solve(
        &mut self, twist: &Twist, init_angles: &Vec<f64>) -> Vec<f64> {
        
        let mut num_small_sigmas = 0;

        self.jacobian.solve_from_chain(&self.chain, init_angles);
        
        // Get the weighted Jacobian.
        let weight_jac = self.weight_task_space * &self.jacobian *
            &self.weight_joint_space;
        
        self.svd_result.compute(&weight_jac, self.epsilon, self.maxiter);
        let u_matrix = &self.svd_result.u_matrix;
        let v_matrix = &self.svd_result.v_matrix;
        let singular_values = &self.svd_result.s_vector;
        // Premultiply the singular vectors by the weights
        let temp_ts = self.weight_task_space *
            u_matrix.fixed_slice::<U6, U6>(0,0);
        let temp_js = &self.weight_joint_space * v_matrix;
        
        // Sigma is the name that is used to describe a particular singular value
        // In this case sigma_min is the min of all singular values
        let sigma_min = if self.n_joints >= 6 {
            singular_values[5]
        } else {
            0.0
        };
        
        for col in 0..self.n_joints {
            let mut sum = 0.0;
            for row in 0..6 {
                if col < 6 {
                    sum += temp_ts[(row, col)]*twist[row]
                }
            }
            // If sigmaMin > eps, then wdls is not active and lambda_scaled = 0
            // If sigmaMin < eps, then wdls is active and lambda_scaled is scaled
            // from 0 to lambda
            let lambda_scaled = if sigma_min < self.epsilon {
                (1.0-(sigma_min/self.epsilon)*(sigma_min/self.epsilon))
                    .sqrt()*self.lambda
            } else {0.0};
            
            self.temp_qdot[col] = if singular_values[col].abs() < self.epsilon {
                num_small_sigmas += 1;
                if col < 6 {
                    // Re-Scale the singular value
                    let sigma = singular_values[col];
                    sum*((sigma/(sigma*sigma+lambda_scaled*lambda_scaled)))
                } else {
                    0.0
                }
            } else {
                sum/singular_values[col]
            };
        }
        
        let qdot_out = self.chain.get_coupling_matrix() * temp_js * &self.temp_qdot;
        if num_small_sigmas > (self.n_joints - 6) {
            panic!{"The pseudo inverse converged but it is Singular"};
        }
        transform_vectors(qdot_out)
    }
}

#[cfg(test)]
mod test {
    use crate::inverse_diff_kinematics::{InverseDiffKinematicsSolver};
    use crate::geometry::{Twist, get_twist_error};
    use crate::chains::tests::{create_testing_chain};
    use crate::forward_diff_kinematics::{ForwardDiffKinematicsSolver};

    #[test]
    fn inverse_kinematics_correct() {
        let chain = create_testing_chain();
        let for_solver = ForwardDiffKinematicsSolver::new(chain.clone());
        let angles = vec![0.1, -0.95, 0.57, 0.68, -0.27, 0.39, 0.47];
        let twist = Twist::new(
            -0.423878, -1.01074, 0.123592, -2.28297, 0.209941, 1.80897);
        let mut inv_solver = InverseDiffKinematicsSolver::new(chain.clone());

        let inv_result = inv_solver.solve(
            &twist, &angles);
        assert_eq!(7, inv_result.len());

        let forward_result = for_solver.solve(
            &angles, &inv_result);
        assert!(get_twist_error(forward_result, twist) < 0.0001);
    }
}