//
// Created by june on 19. 12. 5..
//

// from ros-control meta packages
#include <controller_interface/controller.h>
#include <hardware_interface/joint_command_interface.h>
#include <pluginlib/class_list_macros.h>
#include <urdf/model.h>
#include <realtime_tools/realtime_buffer.h>
#include <realtime_tools/realtime_publisher.h>
#include <geometry_msgs/WrenchStamped.h>
#include <nav_msgs/Odometry.h>
#include <ros_myo/EmgArray.h>

#include "utils.h"
#include "dualarm_controller/TaskDesiredState.h"
#include "dualarm_controller/TaskCurrentState.h"

//manipulability
#include <ellipsoid.h>
#include <similarity.h>

// from kdl packages
#include <kdl/tree.hpp>
#include <kdl/chain.hpp>
#include <kdl_parser/kdl_parser.hpp>
#include <kdl/chaindynparam.hpp>                // inverse dynamics
#include <kdl/chainjnttojacsolver.hpp>          // jacobian
#include <kdl/chainjnttojacdotsolver.hpp>       // jcobian dot
#include <kdl/chainfksolverpos_recursive.hpp>   // forward kinematics

#include <boost/scoped_ptr.hpp>

#define _USE_MATH_DEFINES
#include <cmath>

#include <SerialManipulator.h>
#include <Controller.h>
#include <Motion.h>

#define D2R M_PI/180.0
#define R2D 180.0/M_PI

#define A 0.10
#define b1 0.45
#define b2 -0.30
#define b3 0.45
#define f 0.2

#define l_p1 0.40
#define l_p2 0.30
#define l_p3 0.39

#define Deg_A 70
#define Deg_f 0.5
#define aaa 0
namespace dualarm_controller
{
    class Impedance_Control : public controller_interface::Controller<hardware_interface::EffortJointInterface>
    {
    public:
        bool init(hardware_interface::EffortJointInterface *hw, ros::NodeHandle &n)
        {
            // ********* 1. Get joint name / gain from the parameter server *********
            // 1.0 Target Position objective
            if (!n.getParam("target_obj", target_obj))
            {
                ROS_ERROR("Could not find control objective");
                return false;
            }

            if (!n.getParam("ctrl_type", ctrl_type))
            {
                if(ctrl_type != 1 && ctrl_type != 2)
                {
                    ROS_ERROR("Could not find controller mode");
                    return false;
                }
            }

            // 1.1 Joint Name
            if (!n.getParam("joints", joint_names_))
            {
                ROS_ERROR("Could not find joint name");
                return false;
            }
            n_joints_ = joint_names_.size();

            if (n_joints_ == 0)
            {
                ROS_ERROR("List of joint names is empty.");
                return false;
            }
            else
            {
                ROS_INFO("Found %d joint names", n_joints_);
                for (int i = 0; i < n_joints_; i++)
                {
                    ROS_INFO("%s", joint_names_[i].c_str());
                }
            }

            // 1.2 Gain
            // 1.2.1 Task-space Controller

            if (!n.getParam("/dualarm/impedance_control/desired/mass/m1", des_m1))
            {
                ROS_ERROR("Cannot find Right-arm desired mass");
                return false;
            }

            if (!n.getParam("/dualarm/impedance_control/desired/mass/m2", des_m2))
            {
                ROS_ERROR("Cannot find Left-arm desired mass");
                return false;
            }

            if (!n.getParam("/dualarm/impedance_control/gains/mbk/KpT", Kp_trans))
            {
                ROS_ERROR("Cannot find Left-arm pid/Kp Translation proportional gain");
                return false;
            }

            if (!n.getParam("/dualarm/impedance_control/gains/mbk/KpR", Kp_rot))
            {
                ROS_ERROR("Cannot find Left-arm pid/Kp Rotation proportional gain");
                return false;
            }

            if (!n.getParam("/dualarm/impedance_control/gains/mbk/KdT", Kd_trans))
            {
                ROS_ERROR("Cannot find Left-arm pid/Kd Translation derivative gain");
                return false;
            }

            if (!n.getParam("/dualarm/impedance_control/gains/mbk/KdR", Kd_rot))
            {
                ROS_ERROR("Cannot find Left-arm pid/Kd Rotation derivative gain");
                return false;
            }

            // 2. ********* urdf *********
            urdf::Model urdf;
            if (!urdf.initParam("robot_description"))
            {
                ROS_ERROR("Failed to parse urdf file");
                return false;
            }
            else
            {
                ROS_INFO("Found robot_description");
            }

            // 3. ********* Get the joint object to use in the realtime loop [Joint Handle, URDF] *********
            for (int i = 0; i < n_joints_; i++)
            {
                try
                {
                    joints_.push_back(hw->getHandle(joint_names_[i]));
                }
                catch (const hardware_interface::HardwareInterfaceException &e)
                {
                    ROS_ERROR_STREAM("Exception thrown: " << e.what());
                    return false;
                }

                urdf::JointConstSharedPtr joint_urdf = urdf.getJoint(joint_names_[i]);
                if (!joint_urdf)
                {
                    ROS_ERROR("Could not find joint '%s' in urdf", joint_names_[i].c_str());
                    return false;
                }
                joint_urdfs_.push_back(joint_urdf);
            }

            // 4. ********* KDL *********
            // 4.1 kdl parser
            if (!kdl_parser::treeFromUrdfModel(urdf, kdl_tree_))
            {
                ROS_ERROR("Failed to construct kdl tree");
                return false;
            }

            // 4.2 kdl chain
            std::string root_name, tip_name1, tip_name2;
            if (!n.getParam("root_link", root_name))
            {
                ROS_ERROR("Could not find root link name");
                return false;
            }
            if (!n.getParam("tip_link1", tip_name1))
            {
                ROS_ERROR("Could not find tip link name");
                return false;
            }
            if (!n.getParam("tip_link2", tip_name2))
            {
                ROS_ERROR("Could not find tip link name");
                return false;
            }

            if (!kdl_tree_.getChain(root_name, tip_name1, kdl_chain_))
            {
                ROS_ERROR_STREAM("Failed to get KDL chain from tree: ");
                ROS_ERROR_STREAM("  " << root_name << " --> " << tip_name1);
                ROS_ERROR_STREAM("  Chain has " << kdl_chain_.getNrOfJoints() << " joints");
                ROS_ERROR_STREAM("  Chain has " << kdl_chain_.getNrOfSegments() << " segments");
                ROS_ERROR_STREAM("  The segments are:");

                KDL::SegmentMap segment_map = kdl_tree_.getSegments();
                KDL::SegmentMap::iterator it;

                for (it = segment_map.begin(); it != segment_map.end(); it++)
                    ROS_ERROR_STREAM("    " << (*it).first);

                return false;
            }
            else
            {
                ROS_INFO_STREAM("Got kdl first chain");
                ROS_INFO_STREAM("  " << root_name << " --> " << tip_name1);
                ROS_INFO_STREAM("  Chain has " << kdl_chain_.getNrOfJoints() << " joints");
                ROS_INFO_STREAM("  Chain has " << kdl_chain_.getNrOfSegments() << " segments");
            }

            if(!kdl_tree_.getChain(root_name, tip_name2, kdl_chain2_))
            {
                ROS_ERROR_STREAM("Failed to get KDL chain from tree: ");
                ROS_ERROR_STREAM("  " << root_name << " --> " << tip_name2);
                ROS_ERROR_STREAM("  Chain has " << kdl_chain2_.getNrOfJoints() << " joints");
                ROS_ERROR_STREAM("  Chain has " << kdl_chain2_.getNrOfSegments() << " segments");
                ROS_ERROR_STREAM("  The segments are:");

                KDL::SegmentMap segment_map = kdl_tree_.getSegments();
                KDL::SegmentMap::iterator it;

                for (it = segment_map.begin(); it != segment_map.end(); it++)
                    ROS_ERROR_STREAM("    " << (*it).first);

                return false;
            }
            else
            {
                ROS_INFO_STREAM("Got kdl second chain");
                ROS_INFO_STREAM("  " << root_name << " --> " << tip_name2);
                ROS_INFO_STREAM("  Chain has " << kdl_chain2_.getNrOfJoints() << " joints");
                ROS_INFO_STREAM("  Chain has " << kdl_chain2_.getNrOfSegments() << " segments");
            }

            // 4.3 inverse dynamics solver 초기화
            g_kdl_ = KDL::Vector::Zero();
            g_kdl_(2) = -9.81; // 0: x-axis 1: y-axis 2: z-axis

            id_solver_.reset(new KDL::ChainDynParam(kdl_chain_, g_kdl_));
            jnt_to_jac_solver_.reset(new KDL::ChainJntToJacSolver(kdl_chain_));
            jnt_to_jacdot_solver.reset(new KDL::ChainJntToJacDotSolver(kdl_chain_));
            fk_pos_solver_.reset(new KDL::ChainFkSolverPos_recursive(kdl_chain_));
            J1_kdl_.resize(kdl_chain_.getNrOfJoints());
            J1dot.resize(kdl_chain_.getNrOfJoints());
            M_kdl_.resize(kdl_chain_.getNrOfJoints());
            C_kdl_.resize(kdl_chain_.getNrOfJoints());
            G_kdl_.resize(kdl_chain_.getNrOfJoints());

            id_solver1_.reset(new KDL::ChainDynParam(kdl_chain2_, g_kdl_));
            jnt_to_jac_solver1_.reset(new KDL::ChainJntToJacSolver(kdl_chain2_));
            jnt_to_jacdot_solver1.reset(new KDL::ChainJntToJacDotSolver(kdl_chain2_));
            fk_pos_solver1_.reset(new KDL::ChainFkSolverPos_recursive(kdl_chain2_));
            J2_kdl_.resize(kdl_chain2_.getNrOfJoints());
            J2dot.resize(kdl_chain2_.getNrOfJoints());
            M1_kdl_.resize(kdl_chain2_.getNrOfJoints());
            C1_kdl_.resize(kdl_chain2_.getNrOfJoints());
            G1_kdl_.resize(kdl_chain2_.getNrOfJoints());

            Jdot.setZero(12,n_joints_);

            q1_.resize(9);
            q1dot_.resize(9);
            q2_.resize(9);
            q2dot_.resize(9);

            // ********* 5. 각종 변수 초기화 *********
            // 5.1 KDL Vector 초기화 (사이즈 정의 및 값 0)
            x_cmd_.data = Eigen::VectorXd::Zero(12);
            ex_.setZero(12);
            ex_dot_.setZero(12);

            dtwist.setZero(12);
            dtwist2.setZero(12);
            ex2_.setZero(12);
            ex2_dot_.setZero(12);

            dxdot.setZero(12);
            dxdot2.setZero(12);

            dxddot.setZero(12);
            xdot_logging.setZero(12);

            KpTask.setZero(12);
            KdTask.setZero(12);
            KpNull.setZero(16);
            KdNull.setZero(16);

            emgsig.setZero(12);

            qd_.data = Eigen::VectorXd::Zero(n_joints_);
            qd_dot_.data = Eigen::VectorXd::Zero(n_joints_);
            qd_ddot_.data = Eigen::VectorXd::Zero(n_joints_);

            q_.data = Eigen::VectorXd::Zero(n_joints_);
            qdot_.data = Eigen::VectorXd::Zero(n_joints_);
            torque.setZero(n_joints_);
            ft_sensor.setZero(12);
            torque_ext.setZero(n_joints_);
            des_m.setZero(2);
            targetpos.setZero(12);
            xa.setZero(12);
            xa2.setZero(12);
            Kp_emg.setZero(12);
            Kd_emg.setZero(12);
            Mx.setZero(12);

            wpInv_lambda[0].setZero(3);
            wpInv_lambda[1].setZero(3);
            frictiontorque_.data = Eigen::VectorXd::Zero(n_joints_);

            // ********* 6. ROS 명령어 *********
            // 6.1 publisher
            state_pub_.reset(new realtime_tools::RealtimePublisher<dualarm_controller::TaskCurrentState>(n, "states", 10));
            state_pub_->msg_.header.stamp = ros::Time::now();
            //state_pub_->msg_.header.frame_id = "dualarm";
            state_pub_->msg_.header.seq = 0;
            for(int i=0; i<n_joints_; i++)
            {
                state_pub_->msg_.q.push_back(q_.data(i));
                state_pub_->msg_.qdot.push_back(qdot_.data(i));
                state_pub_->msg_.dq.push_back(qd_.data(i));
                state_pub_->msg_.dqdot.push_back(qd_dot_.data(i));
                state_pub_->msg_.torque.push_back(torque(i));
                state_pub_->msg_.torque_ext.push_back(torque_ext(i));

            }
            for(int j=0; j<2; j++) {
                state_pub_->msg_.InverseConditionNum.push_back(InverseConditionNumber[j]);
                state_pub_->msg_.SingleMM.push_back(SingleMM[j]);
            }
            state_pub_->msg_.x.resize(2);
            state_pub_->msg_.dx.resize(2);
            state_pub_->msg_.ex.resize(2);
            state_pub_->msg_.DAMM = DAMM;
            state_pub_->msg_.TODAMM = TODAMM2;
            state_pub_->msg_.lambda1.resize(3);
            state_pub_->msg_.lambda2.resize(3);
            state_pub_->msg_.Kp_R = Kp_rot;
            state_pub_->msg_.Kp_T = Kp_trans;
            state_pub_->msg_.Kp_emgg.resize(12);
            state_pub_->msg_.Kd_emgg.resize(12);
            state_pub_->msg_.emgsigg=emgsig_int;
            state_pub_->msg_.Mx_diag.resize(12);

            state_pub_->msg_.ftsensor.resize(12);

            state_pub_->msg_.xdot.resize(12);
            state_pub_->msg_.dxdot.resize(12);
            state_pub_->msg_.dxddot.resize(12);

            pub_buffer_.writeFromNonRT(std::vector<double>(n_joints_, 0.0));

            // 6.2 subscriber
            const auto joint_state_cb = utils::makeCallback<dualarm_controller::TaskDesiredState>([&](const auto& msg){

                ControlIndex1 = 5;
                ControlIndex2 = 1;
                ControlSubIndex = 1;
                JointState = ControlSubIndex;

            });

            sub_x_cmd_ = n.subscribe<dualarm_controller::TaskDesiredState>( "command", 5, joint_state_cb);
            sub_ft_sensor_R = n.subscribe<geometry_msgs::WrenchStamped>("/ft_sensor_topic_R", 5, &Impedance_Control::UpdateFTsensorR, this);
            sub_ft_sensor_L = n.subscribe<geometry_msgs::WrenchStamped>("/ft_sensor_topic_L", 5, &Impedance_Control::UpdateFTsensorL, this);
            vive_pose_sub_R  = n.subscribe<nav_msgs::Odometry>( "/LHR_18744CF1_odom", 5, &Impedance_Control::Update_vive_pose_R, this);
            vive_pose_sub_L  = n.subscribe<nav_msgs::Odometry>( "/LHR_FC007723_odom", 5, &Impedance_Control::Update_vive_pose_L, this);
            vive_pose_sub_R2  = n.subscribe<nav_msgs::Odometry>( "/LHR_539B1118_odom", 5, &Impedance_Control::Update_vive_pose_R2, this);
            vive_pose_sub_L2  = n.subscribe<nav_msgs::Odometry>( "/LHR_078BFD2F_odom", 5, &Impedance_Control::Update_vive_pose_L2, this);
            Myo_emg  = n.subscribe<ros_myo::EmgArray>( "/myo_raw/myo_emg", 6, &Impedance_Control::Update_myo_emg, this);

            return true;
        }

        void Update_myo_emg(const ros_myo::EmgArray::ConstPtr &msg)
        {

            emgsig.setConstant(msg->rms_data);
            emgsig_int=emgsig(0);

        }

        void Update_vive_pose_R(const nav_msgs::Odometry::ConstPtr &msg)
        {

            vive_dR_R.w()= msg->pose.pose.orientation.w;
            vive_dR_R.x()= msg->pose.pose.orientation.x;
            vive_dR_R.y()= msg->pose.pose.orientation.y;
            vive_dR_R.z()= msg->pose.pose.orientation.z;

            vive_dP_R(0)=msg->pose.pose.position.x;
            vive_dP_R(1)=msg->pose.pose.position.y;
            vive_dP_R(2)=msg->pose.pose.position.z;

            dtwist(0)=msg->twist.twist.angular.x;
            dtwist(1)=msg->twist.twist.angular.y;
            dtwist(2)=msg->twist.twist.angular.z;
            dtwist(3)=msg->twist.twist.linear.x;
            dtwist(4)=msg->twist.twist.linear.y;
            dtwist(5)=msg->twist.twist.linear.z;

        }

        void Update_vive_pose_L(const nav_msgs::Odometry::ConstPtr &msg)
        {

            vive_dR_L.w()= msg->pose.pose.orientation.w;
            vive_dR_L.x()= msg->pose.pose.orientation.x;
            vive_dR_L.y()= msg->pose.pose.orientation.y;
            vive_dR_L.z()= msg->pose.pose.orientation.z;

            vive_dP_L(0)=msg->pose.pose.position.x;
            vive_dP_L(1)=msg->pose.pose.position.y;
            vive_dP_L(2)=msg->pose.pose.position.z;

            dtwist(6)=msg->twist.twist.angular.x;
            dtwist(7)=msg->twist.twist.angular.y;
            dtwist(8)=msg->twist.twist.angular.z;
            dtwist(9)=msg->twist.twist.linear.x;
            dtwist(10)=msg->twist.twist.linear.y;
            dtwist(11)=msg->twist.twist.linear.z;

        }

        void Update_vive_pose_R2(const nav_msgs::Odometry::ConstPtr &msg)
        {

            vive_dP_R2(0)=msg->pose.pose.position.x;
            vive_dP_R2(1)=msg->pose.pose.position.y;
            vive_dP_R2(2)=msg->pose.pose.position.z;

            dtwist2(3)=msg->twist.twist.linear.x;
            dtwist2(4)=msg->twist.twist.linear.y;
            dtwist2(5)=msg->twist.twist.linear.z;

        }

        void Update_vive_pose_L2(const nav_msgs::Odometry::ConstPtr &msg)
        {

            vive_dP_L2(0)=msg->pose.pose.position.x;
            vive_dP_L2(1)=msg->pose.pose.position.y;
            vive_dP_L2(2)=msg->pose.pose.position.z;

            dtwist2(9)=msg->twist.twist.linear.x;
            dtwist2(10)=msg->twist.twist.linear.y;
            dtwist2(11)=msg->twist.twist.linear.z;

        }

        void UpdateFTsensorR(const geometry_msgs::WrenchStamped::ConstPtr &msg)
        {
            geometry_msgs::Wrench ft_measure = msg->wrench;
            ft_sensor(0) = ft_measure.torque.x;
            ft_sensor(1) = ft_measure.torque.y;
            ft_sensor(2) = ft_measure.torque.z;
            ft_sensor(3) = ft_measure.force.x;
            ft_sensor(4) = ft_measure.force.y;
            ft_sensor(5) = ft_measure.force.z;
        }

        void UpdateFTsensorL(const geometry_msgs::WrenchStamped::ConstPtr &msg)
        {
            geometry_msgs::Wrench ft_measure = msg->wrench;
            ft_sensor(6) = ft_measure.torque.x;
            ft_sensor(7) = ft_measure.torque.y;
            ft_sensor(8) = ft_measure.torque.z;
            ft_sensor(9) = ft_measure.force.x;
            ft_sensor(10) = ft_measure.force.y;
            ft_sensor(11) = ft_measure.force.z;
        }

        void starting(const ros::Time &time) override {
            t = 0.0;
            InitTime = 2.0;

            cManipulator = std::make_shared<SerialManipulator>();

            Control = std::make_unique<HYUControl::Controller>(cManipulator);
            motion = std::make_unique<HYUControl::Motion>(cManipulator);

            cManipulator->UpdateManipulatorParam();

            des_m(0) = des_m1;
            des_m(1) = des_m2;

            KpTask.segment(0, 3).setConstant(Kp_rot);
            KpTask.segment(3, 3).setConstant(Kp_trans);
            KpTask.segment(6, 3).setConstant(Kp_rot);
            KpTask.segment(9, 3).setConstant(Kp_trans);

            KdTask.segment(0, 3).setConstant(Kd_rot);
            KdTask.segment(3, 3).setConstant(Kd_trans);
            KdTask.segment(6, 3).setConstant(Kd_rot);
            KdTask.segment(9, 3).setConstant(Kd_trans);

            KpNull.setConstant(16,0.1);
            KdNull.setConstant(16,0.4 );
            Control->SetImpedanceGain(KpTask, KdTask, KpNull, KdNull, des_m);
//
//            ControlIndex1 = CTRLMODE_IDY_JOINT;
//            ControlIndex2 = SYSTEM_BEGIN;
//            ControlSubIndex = MOVE_ZERO;
//            JointState = MOVE_ZERO;

            ControlIndex1 = 5;
            ControlIndex2 = 1;
            ControlSubIndex = 1;
            JointState = 1;

            ROS_INFO("Starting Impedance Controller");

        }

        void update(const ros::Time &time, const ros::Duration &period) override
        {
            std::vector<double> &commands = *pub_buffer_.readFromRT();
            // ********* 0. Get states from gazebo *********
            // 0.1 sampling time
            double dt = period.toSec();

            clock_gettime(CLOCK_MONOTONIC, &begin);

            // 0.2 joint state
            for (int i = 0; i < n_joints_; i++)
            {
                q_(i) = joints_[i].getPosition();
                qdot_(i) = joints_[i].getVelocity();
                //torque_(i) = joints_[i].getEffort();
            }

            //----------------------
            // dynamics calculation
            //----------------------
            cManipulator->pKin->PrepareJacobian(q_.data);
            cManipulator->pDyn->PrepareDynamics(q_.data, qdot_.data);
            cManipulator->pKin->GetForwardKinematics(ForwardPos, ForwardOri, NumChain);
            cManipulator->pKin->GetAngleAxis(ForwardAxis, ForwardAngle, NumChain);
            cManipulator->pKin->GetInverseConditionNumber(InverseConditionNumber);
            cManipulator->pKin->GetAnalyticJacobian(AJac);

            xdot_logging = AJac*qdot_.data;

            q1_.data = q_.data.head(9);
            q1dot_.q = q1_;
            q1dot_.qdot.data = qdot_.data.head(9);

            q2_.data.tail(7) = q_.data.tail(7);
            q2_.data.head(2) = q_.data.head(2);
            q2dot_.q = q2_;
            q2dot_.qdot.data.tail(7) = qdot_.data.tail(7);
            q2dot_.qdot.data.head(2) = qdot_.data.head(2);

            fk_pos_solver_->JntToCart(q1_, x_[0]);
            fk_pos_solver1_->JntToCart(q2_, x_[1]);

            jnt_to_jac_solver_->JntToJac(q1_, J1_kdl_);
            jnt_to_jac_solver1_->JntToJac(q2_, J2_kdl_);

            jnt_to_jacdot_solver->JntToJacDot(q1dot_, J1dot);
            jnt_to_jacdot_solver1->JntToJacDot(q2dot_, J2dot);

            cManipulator->pDyn->M_Matrix(M);
            KineticEnergy = qdot_.data.transpose()*M*qdot_.data;

            manipulability_data();

            ctrl_type = ControlIndex2;
            target_obj = ControlSubIndex;

            if( ControlIndex1 == CTRLMODE_IMPEDANCE_TASK )
            {
                if(ControlIndex2 == 3)
                {
                    cManipulator->pKin->GetForwardKinematicsWithRelative(xa);
                    cManipulator->pKin->GetRelativeJacobian(RelativeJac);
                    SingleMM[0] = cManipulator->pKin->GetManipulabilityMeasure(AJac.block(0,0,6,16));
                    SingleMM[1] = cManipulator->pKin->GetManipulabilityMeasure(RelativeJac);
                    AJac.block(6,0,6,16) = RelativeJac;
                    MM = cManipulator->pKin->GetManipulabilityMeasure(AJac);
                }
                else
                {
                    cManipulator->pKin->GetForwardKinematics(xa);
                    cManipulator->pKin2->GetForwardKinematics(xa2);

                    SingleMM[0] = cManipulator->pKin->GetManipulabilityMeasure(AJac.block(0,0,6,16));
                    SingleMM[1] = cManipulator->pKin->GetManipulabilityMeasure(AJac.block(6,0,6,16));
                    MM = cManipulator->pKin->GetManipulabilityMeasure();
                }
                motion->TaskMotion2(q_R,q_L, TargetPos_Linear_R, TargetPos_Linear_L, TargetPos_Linear_R2, TargetPos_Linear_L2,
                                    vive_dR_R, vive_dR_L,vive_dP_R,vive_dP_L, vive_dP_R2,vive_dP_L2,
                                    dtwist, dtwist2, dxdot,dxdot2, dxddot, targetpos, xa, qdot_.data, t, JointState, ControlSubIndex);
//                std::cout<<TargetPos_Linear_R2<<std::endl;

                //Control->TaskImpedanceController(q_.data, qdot_.data, dx, dxdot, dxddot, ft_sensor, torque,ControlIndex2);
//                Control->TaskImpedanceController2(q_.data, qdot_.data, dx, dxdot, dxddot, ft_sensor, torque, q_R,q_L, TargetPos_Linear_R, TargetPos_Linear_L, frictiontorque,ControlIndex2);

                Control->TaskImpedanceController3(q_.data, qdot_.data, dx, dxdot,dxdot2, dxddot, ft_sensor, torque, q_R,q_L, TargetPos_Linear_R, TargetPos_Linear_L, TargetPos_Linear_R2,TargetPos_Linear_L2,ControlIndex2,emgsig, Mx, Kd_emg, Kp_emg);

                Control->GetControllerStates(qd_.data, qd_dot_.data, ex_);
                Control->GetControllerStates2(qd_.data, qd_dot_.data, ex2_);
//                std::cout<<des_m<<std::endl;

                cManipulator->pKin->GetWDampedpInvLambda(wpInv_lambda);
            }
            else if( ControlIndex1 == CTRLMODE_IDY_JOINT )
            {
                MM = cManipulator->pKin->GetManipulabilityMeasure();
                qd_dot_.data.setZero(16);
                motion->JointMotion(qd_.data, qd_dot_.data, qd_ddot_.data, targetpos, q_.data, qdot_.data, t, JointState, ControlSubIndex);
                Control->InvDynController2(q_.data, qdot_.data, qd_.data, qd_dot_.data, qd_ddot_.data, torque, frictiontorque_.data, dt);
            }

            for (int i = 0; i < n_joints_; i++)
            {
                joints_[i].setCommand(torque(i));
            }

            clock_gettime(CLOCK_MONOTONIC, &end);
            // ********* 4. data 저장 *********
            publish_data();

            // ********* 5. state 출력 *********
            print_state();

            t = t + dt;
        }

        void stopping(const ros::Time &time) override
        {
            ROS_INFO("Stop Impedance Controller");
        }

        void manipulability_data()
        {
            const auto desired_manipulability =
                    manipulability_metrics::Ellipsoid{ { { (Eigen::Matrix<double, 6, 1>{} << 1, 0, 0, 0, 0, 0).finished(), 1.0 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 1, 0, 0, 0, 0).finished(), 1.0 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 0, 1, 0, 0, 0).finished(), 1.0 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 0, 0, 1, 0, 0).finished(), 1.0 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 0, 0, 0, 1, 0).finished(), 1.0 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 0, 0, 0, 0, 1).finished(), 1.0 } } };

            const auto desired_Singelmanipulability =
                    manipulability_metrics::Ellipsoid{ { { (Eigen::Matrix<double, 6, 1>{} << 1, 0, 0, 0, 0, 0).finished(), 0.2 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 1, 0, 0, 0, 0).finished(), 0.8 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 0, 1, 0, 0, 0).finished(), 0.8 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 0, 0, 1, 0, 0).finished(), 0.9 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 0, 0, 0, 1, 0).finished(), 1.7 },
                                                               { (Eigen::Matrix<double, 6, 1>{} << 0, 0, 0, 0, 0, 1).finished(), 1.7 } } };
            auto left_ellipsoid = manipulability_metrics::ellipsoidFromJacobian(J2_kdl_.data);
            auto right_ellipsoid = manipulability_metrics::ellipsoidFromJacobian(J1_kdl_.data);
            TOMM[0] = manipulability_metrics::inverseShapeDiscrepancy(desired_Singelmanipulability, J1_kdl_.data);
            TOMM[1] = manipulability_metrics::inverseShapeDiscrepancy(desired_Singelmanipulability, J2_kdl_.data);
            DAMM = std::max(manipulability_metrics::volumeIntersection(left_ellipsoid, J1_kdl_.data),
                            manipulability_metrics::volumeIntersection(right_ellipsoid, J2_kdl_.data));
            TODAMM2 = manipulability_metrics::dualInverseShapeDiscrepancy(desired_manipulability, J2_kdl_.data, J1_kdl_.data);
        }

        void publish_data()
        {
            static int loop_count_ = 0;
            if(loop_count_ > 2)
            {
                if(state_pub_->trylock())
                {
                    state_pub_->msg_.header.stamp = ros::Time::now();
                    state_pub_->msg_.header.seq++;
                    torque_ext = AJac.transpose()*ft_sensor;
                    for(size_t i=0; i<n_joints_; i++)
                    {
                        state_pub_->msg_.q[i] = q_.data(i);
                        state_pub_->msg_.qdot[i] = qdot_.data(i);
                        state_pub_->msg_.dq[i] = qd_.data(i);
                        state_pub_->msg_.dqdot[i] = qd_dot_.data(i);
                        state_pub_->msg_.torque[i] = torque(i);
                        state_pub_->msg_.torque_ext[i] = torque_ext(i);
                    }

                    Quaterniond tmp;
                    Vector3d RTmp1;
                    Vector3d RTmp_R, RTmp_L;
                    Quaterniond q_aR,q_aL;

                    q_aR= AngleAxisd(xa(0), Vector3d::UnitX())
                         * AngleAxisd(xa(1), Vector3d::UnitY())
                         * AngleAxisd(xa(2), Vector3d::UnitZ());
                    q_aL= AngleAxisd(xa(6), Vector3d::UnitX())
                          * AngleAxisd(xa(7), Vector3d::UnitY())
                          * AngleAxisd(xa(8), Vector3d::UnitZ());
                    if (q_R.coeffs().dot(q_aR.coeffs()) < 0.0)
                    {
                        q_R.coeffs() << -q_R.coeffs();
                    }
                    if (q_L.coeffs().dot(q_aL.coeffs()) < 0.0)
                    {
                        q_L.coeffs() << -q_L.coeffs();
                    }

                    RTmp_R = q_R.toRotationMatrix().eulerAngles(2,1,0);
                    RTmp_L = q_L.toRotationMatrix().eulerAngles(2,1,0);
                    for(int j=0; j<2; j++)
                    {
//                        tmp = dx[j].r;

                        Quaterniond q_d;

                        if(j==0) {
                            state_pub_->msg_.dx[j].orientation.x = RTmp_R(2);
                            state_pub_->msg_.dx[j].orientation.y = RTmp_R(1);
                            state_pub_->msg_.dx[j].orientation.z = RTmp_R(0);
                            state_pub_->msg_.dx[j].position.x = TargetPos_Linear_R(0);
                            state_pub_->msg_.dx[j].position.y = TargetPos_Linear_R(1);
                            state_pub_->msg_.dx[j].position.z = TargetPos_Linear_R(2);
//                            state_pub_->msg_.dx2[j].position.x = TargetPos_Linear_R2(0);
//                            state_pub_->msg_.dx2[j].position.y = TargetPos_Linear_R2(1);
//                            state_pub_->msg_.dx2[j].position.z = TargetPos_Linear_R2(2);
                            }
                        else{
                            state_pub_->msg_.dx[j].orientation.x = RTmp_L(2);
                            state_pub_->msg_.dx[j].orientation.y = RTmp_L(1);
                            state_pub_->msg_.dx[j].orientation.z = RTmp_L(0);
                            state_pub_->msg_.dx[j].position.x = TargetPos_Linear_L(0);
                            state_pub_->msg_.dx[j].position.y = TargetPos_Linear_L(1);
                            state_pub_->msg_.dx[j].position.z = TargetPos_Linear_L(2);
//                            state_pub_->msg_.dx2[j].position.x = TargetPos_Linear_R2(0);
//                            state_pub_->msg_.dx2[j].position.y = TargetPos_Linear_R2(1);
//                            state_pub_->msg_.dx2[j].position.z = TargetPos_Linear_R2(2);
                            }

                        state_pub_->msg_.x[j].orientation.x = xa(6*j);
                        state_pub_->msg_.x[j].orientation.y = xa(6*j+1);
                        state_pub_->msg_.x[j].orientation.z = xa(6*j+2);
                        state_pub_->msg_.x[j].position.x = xa(6*j+3);
                        state_pub_->msg_.x[j].position.y = xa(6*j+4);
                        state_pub_->msg_.x[j].position.z = xa(6*j+5);

                        state_pub_->msg_.ex[j].orientation.x = ex_(6*j);
                        state_pub_->msg_.ex[j].orientation.y = ex_(6*j+1);
                        state_pub_->msg_.ex[j].orientation.z = ex_(6*j+2);
                        state_pub_->msg_.ex[j].position.x = ex_(6*j+3);
                        state_pub_->msg_.ex[j].position.y = ex_(6*j+4);
                        state_pub_->msg_.ex[j].position.z = ex_(6*j+5);

                        state_pub_->msg_.InverseConditionNum[j] = InverseConditionNumber[j];
                        state_pub_->msg_.SingleMM[j] = SingleMM[j];
                    }

                    state_pub_->msg_.MM = MM;
                    state_pub_->msg_.DAMM = DAMM;
                    state_pub_->msg_.TODAMM = TODAMM2;
                    state_pub_->msg_.lambda1[0] = wpInv_lambda[0](0);
                    state_pub_->msg_.lambda1[1] = wpInv_lambda[0](1);
                    state_pub_->msg_.lambda1[2] = wpInv_lambda[0](2);
                    state_pub_->msg_.lambda2[0] = wpInv_lambda[1](0);
                    state_pub_->msg_.lambda2[1] = wpInv_lambda[1](1);
                    state_pub_->msg_.lambda2[2] = wpInv_lambda[1](2);

                    state_pub_->msg_.emgsigg = emgsig_int;

                    for(int k=0; k<12; k++)
                    {
                        state_pub_->msg_.ftsensor[k] = ft_sensor(k);
                        state_pub_->msg_.xdot[k] = xdot_logging(k);
                        state_pub_->msg_.dxdot[k] = dxdot(k);
                        state_pub_->msg_.dxddot[k] = dxddot(k);
                        state_pub_->msg_.Kp_emgg[k] = Kp_emg(k);
                        state_pub_->msg_.Kd_emgg[k] = Kd_emg(k);
                        state_pub_->msg_.Mx_diag[k] = Mx(k);

                    }
                    state_pub_->msg_.KE = KineticEnergy;

                    state_pub_->unlockAndPublish();
                }
                loop_count_=0;
            }
            loop_count_++;
        }

        void print_state()
        {
            static int count = 0;
            if (count > 499)
            {
                printf("*********************************************************\n\n");
                printf("*** Calcutaion Time (unit: sec)  ***\n");
                printf("t_cal = %0.9lf\n", static_cast<double>(end.tv_sec - begin.tv_sec) + static_cast<double>(end.tv_nsec - begin.tv_nsec) / 1000000000.0);
                printf("*** Simulation Time (unit: sec)  ***\n");
                printf("t = %0.3lf\n", t);
                printf("Index1:%d\n", ControlIndex1);
                printf("Index2:%d\n", ControlIndex2);
                printf("SubIndex:%d\n", ControlSubIndex);
                printf("MsgSubIndex:%d\n", JointState);
                printf("\n");

                printf("*** Command from Subscriber in Task Space (unit: m) ***\n");
                printf("*** States in Joint Space (unit: deg) ***\n");
                for(int i=0; i < n_joints_; i++)
                {
                    printf("[%s]:  \t", joint_names_[i].c_str());
                    printf("q: %0.2lf,\t", q_.data(i) * R2D);
                    printf("dq: %0.2lf,\t", qd_.data(i) * R2D);

                    printf("qdot: %0.2lf,\t", qdot_.data(i) * R2D);
                    printf("tau: %0.2f\n", torque(i));
                }

                printf("\nForward Kinematics:\n");
                for(int j=0; j<NumChain; j++)
                {
                    printf("no.%d, Actual: x:%0.3lf, y:%0.3lf, z:%0.3lf, u:%0.2lf, v:%0.2lf, w:%0.2lf\n", j,
                           ForwardPos[j](0), ForwardPos[j](1),ForwardPos[j](2),
                           ForwardOri[j](0), ForwardOri[j](1), ForwardOri[j](2));
                    double a, b, g;
                    x_[j].M.GetEulerZYX(a, b, g);

                    printf("no.%d, DH: x:%0.3lf, y:%0.3lf, z:%0.3lf, u:%0.2lf, v:%0.2lf, w:%0.2lf\n",
                           j, x_[j].p(0), x_[j].p(1),x_[j].p(2), g, b, a);
                    printf("no.%d, AngleAxis x: %0.2lf, y: %0.2lf, z: %0.2lf, Angle: %0.3lf\n\n",
                           j, ForwardAxis[j](0), ForwardAxis[j](1), ForwardAxis[j](2), ForwardAngle[j]);
                }
                Vector3d RTmp_R, RTmp_L;
                RTmp_R = q_R.toRotationMatrix().eulerAngles(2,1,0);
                RTmp_L = q_L.toRotationMatrix().eulerAngles(2,1,0);
                printf("Right Desired: x:%0.3lf, y:%0.3lf, z:%0.3lf, u:%0.2lf, v:%0.2lf, w:%0.2lf\n",
                       TargetPos_Linear_R(0), TargetPos_Linear_R(1), TargetPos_Linear_R(2), RTmp_R(2), RTmp_R(1), RTmp_R(0));
                printf("Left Desired: x:%0.3lf, y:%0.3lf, z:%0.3lf, u:%0.2lf, v:%0.2lf, w:%0.2lf\n",
                       TargetPos_Linear_L(0), TargetPos_Linear_L(1), TargetPos_Linear_L(2), RTmp_L(2), RTmp_L(1), RTmp_L(0));

//                printf("Right Desired: x:%0.3lf, y:%0.3lf, z:%0.3lf, Left Desired: x:%0.2lf, y:%0.2lf, z:%0.2lf\n",
//                       TargetPos_Linear_R2(0), TargetPos_Linear_R2(1), TargetPos_Linear_R2(2), TargetPos_Linear_L2(0), TargetPos_Linear_L2(1), TargetPos_Linear_L2(2));
//                printf("elbow_act: x:%0.3lf, y:%0.3lf, z:%0.3lf\n",
//                       xa2(3),xa2(4),xa2(5));
//                printf("elbow_val: x:%0.3lf, y:%0.3lf, z:%0.3lf\n",
//                       dtwist2(3),dtwist2(4),dtwist2(5));

                printf("Right e(u):%0.3lf, e(v):%0.3lf, e(w):%0.3lf, e(x):%0.3lf, e(y):%0.3lf, e(z):%0.3lf\n",
                       ex_(0)*RADtoDEG, ex_(1)*RADtoDEG, ex_(2)*RADtoDEG, ex_(3), ex_(4), ex_(5));
                printf("Left e(u):%0.3lf, e(v):%0.3lf, e(w):%0.3lf, e(x):%0.3lf, e(y):%0.3lf, e(z):%0.3lf\n\n",
                       ex_(6)*RADtoDEG, ex_(7)*RADtoDEG, ex_(8)*RADtoDEG, ex_(9), ex_(10), ex_(11));

                printf("FT Sensor(Right): torque_u:%0.3lf, torque_v:%0.3lf, torque_w:%0.3lf, force_x:%0.3lf, force_y:%0.3lf, force_z:%0.3lf\n",
                       ft_sensor(0), ft_sensor(1), ft_sensor(2),ft_sensor(3),ft_sensor(4),ft_sensor(5));
                printf("FT Sensor(Left): torque_u:%0.3lf, torque_v:%0.3lf, torque_w:%0.3lf, force_x:%0.3lf, force_y:%0.3lf, force_z:%0.3lf\n\n",
                       ft_sensor(6), ft_sensor(7), ft_sensor(8),ft_sensor(9),ft_sensor(10),ft_sensor(11));

//                printf("Inverse Condition Number: Right:%0.5lf, Left:%0.5lf \n", InverseConditionNumber[0], InverseConditionNumber[1]);
//                printf("SingleMM: Right:%0.5lf, Left:%0.5lf\n", SingleMM[0], SingleMM[1]);
//                printf("TOMM: Right:%0.5lf, Left:%0.5lf\n", TOMM[0], TOMM[1]);
                printf("MM: %0.5lf\n", MM);
                printf("Mx_R: x:%0.3lf, y:%0.3lf, z:%0.3lf\n", Mx(3,3),Mx(4,4),Mx(5,5));
                printf("Mx_L: x:%0.3lf, y:%0.3lf, z:%0.3lf\n", Mx(9,9),Mx(10,10),Mx(11,11));
                printf("Kp_emg: x:%0.3lf, y:%0.3lf, z:%0.3lf\n", Kp_emg(3),Kp_emg(4),Kp_emg(5));
                printf("Kd_emg: x:%0.3lf, y:%0.3lf, z:%0.3lf\n", Kd_emg(3),Kd_emg(4),Kd_emg(5));
                cout<<emgsig(0)<<endl;
                //                printf("DAMM: %0.5lf\n", DAMM);
//                printf("TODAMM: %0.5lf\n", TODAMM2);
                printf("*********************************************************\n");
                count = 0;

                //std::cout << "Weight Damped Jacobian Pseudoinverse:" << std::endl;
                //std::cout << WdampedpInvJac << "\n" << std::endl;
                //std::cout << "lambda:" << std::endl;
                //std::cout << wpInv_lambda[0] << std::endl;
                //std::cout << wpInv_lambda[1] << std::endl;

                //std::cout << "J1:" <<std::endl;
                //std::cout << J1_kdl_.data << "\n"<< std::endl;
                //std::cout << "J2:" <<std::endl;
                //std::cout << J2_kdl_.data << "\n"<< std::endl;
                //std::cout << "Analytic Jacobian:" <<std::endl;
                //std::cout << AJac << "\n"<< std::endl;
                //std::cout << "J1dot:" <<std::endl;
                //std::cout << J1dot.data << "\n"<< std::endl;
                //std::cout << "J2dot:" <<std::endl;
                //std::cout << J2dot.data << "\n"<< std::endl;

                //std::cout << "Body Jacobiandot:" <<std::endl;
                //std::cout << BodyJacDot << "\n"<< std::endl;
                /*
                Jdot.setZero();
                Jdot.block(0,0,3,9) = J1dot.data.block(3,0,3,9);
                Jdot.block(3,0,3,9) = J1dot.data.block(0,0,3,9);
                Jdot.block(6,0,3,2) = J2dot.data.block(3,0,3,2);
                Jdot.block(9,0,3,2) = J2dot.data.block(0,0,3,2);
                Jdot.block(6,9,3,7) = J2dot.data.block(3,2,3,7);
                Jdot.block(9,9,3,7) = J2dot.data.block(0,2,3,7);

                std::cout << "Jdot:" <<std::endl;
                std::cout << Jdot << "\n"<< std::endl;
                std::cout << "Analytic Jacobiandot:" <<std::endl;
                std::cout << AJacDot << "\n"<< std::endl;
                */

            }
            count++;
        }

    private:
        // others
        double t=0.0;
        int target_obj=0;
        int ctrl_type=0;
        double InitTime=0.0;

        unsigned char ControlIndex1;
        unsigned char ControlIndex2;
        unsigned char ControlSubIndex;
        unsigned char JointState;

        struct timespec begin, end;
        //Joint handles
        unsigned int n_joints_;
        std::vector<std::string> joint_names_;
        std::vector<hardware_interface::JointHandle> joints_;
        std::vector<urdf::JointConstSharedPtr> joint_urdfs_;

        // kdl
        KDL::Tree kdl_tree_;
        KDL::Chain kdl_chain_;
        KDL::Chain kdl_chain2_;

        KDL::JntSpaceInertiaMatrix M_kdl_, M1_kdl_;
        KDL::JntArray C_kdl_, C1_kdl_;
        KDL::JntArray G_kdl_, G1_kdl_;
        KDL::Vector g_kdl_;
        Eigen::VectorXd g_vec_collect;
        Eigen::MatrixXd g_mat_collect;
        Eigen::MatrixXd M_mat_collect;

        KDL::Jacobian J1_kdl_, J2_kdl_;
        KDL::Jacobian J1dot, J2dot;

        // kdl solver
        boost::scoped_ptr<KDL::ChainFkSolverPos_recursive> fk_pos_solver_, fk_pos_solver1_; //Solver to compute the forward kinematics (position)
        boost::scoped_ptr<KDL::ChainJntToJacSolver> jnt_to_jac_solver_, jnt_to_jac_solver1_; //Solver to compute the jacobian
        boost::scoped_ptr<KDL::ChainJntToJacDotSolver> jnt_to_jacdot_solver, jnt_to_jacdot_solver1;
        boost::scoped_ptr<KDL::ChainDynParam> id_solver_, id_solver1_;               // Solver To compute the inverse dynamics

        MatrixXd M;
        VectorXd G;
        VectorXd C;

        Vector3d ForwardPos[2];
        Vector3d ForwardPos2[2];
        Vector3d ForwardOri[2];
        int NumChain=2;
        Vector3d ForwardAxis[2];
        double ForwardAngle[2];
        double InverseConditionNumber[2];
        double SingleMM[2];
        double TOMM[2];
        double MM;
        double DAMM;
        double TODAMM2;
        double KineticEnergy;

        // kdl and Eigen Jacobian
        Eigen::MatrixXd pInvJac;
        Eigen::MatrixXd AJac, AJacDot, Jdot;
        Eigen::MatrixXd RelativeJac;
        Eigen::MatrixXd BodyJacDot;

        // Joint Space State
        KDL::JntArray qd_;
        KDL::JntArray qd_dot_;
        KDL::JntArray qd_ddot_;
        KDL::JntArray q_, q1_, q2_;
        KDL::JntArray qdot_;
        KDL::JntArrayVel q1dot_, q2dot_;
        KDL::JntArray frictiontorque_;
        VectorXd frictiontorque;

        Eigen::VectorXd xa, xa2;
        Eigen::VectorXd xdot_logging;
        Eigen::VectorXd q0dot;
        Eigen::VectorXd torque;
        Eigen::VectorXd ft_sensor;
        Eigen::VectorXd torque_ext;
        Eigen::VectorXd targetpos;
        Eigen::VectorXd wpInv_lambda[2];

        // Task Space State
        // ver. 01
        KDL::Frame xd_[2];
        KDL::Frame x_[2];
        KDL::Twist ex_temp_;

        // KDL::Twist xd_dot_, xd_ddot_;
        Cartesiand dx[2];
        Eigen::VectorXd ex_, ex2_;
        Eigen::VectorXd ex_dot_,ex2_dot_;
        Eigen::VectorXd dxdot, dxdot2;
        Eigen::VectorXd dxddot;
        Eigen::VectorXd dtwist, dtwist2;
        Eigen::Vector3d TargetPos_Linear_R, TargetPos_Linear_L, vive_dP_R, vive_dP_L,TargetPos_Linear_R2, TargetPos_Linear_L2, vive_dP_R2, vive_dP_L2;
        Eigen::VectorXd Mx;

        Eigen::Quaterniond q_R, q_L, vive_dR_R, vive_dR_L;

        // Input
        KDL::JntArray x_cmd_;

        // gains
        Eigen::VectorXd KpTask, KdTask, des_m;
        Eigen::VectorXd KpNull, KdNull;
        Eigen::VectorXd aKp_, aKi_, aKd_, aK_inf_;
        Eigen::VectorXd emgsig, Kp_emg, Kd_emg;
        double emgsig_int;
        double des_m1, des_m2;
        double Kp_trans, Kp_rot, Kd_trans, Kd_rot;

        // publisher
        realtime_tools::RealtimeBuffer<std::vector<double>> pub_buffer_;
        boost::scoped_ptr<realtime_tools::RealtimePublisher<dualarm_controller::TaskCurrentState>> state_pub_;

        // subscriber
        ros::Subscriber sub_x_cmd_;
        ros::Subscriber sub_ft_sensor_R, sub_ft_sensor_L;
        ros::Subscriber vive_pose_sub_R, vive_pose_sub_L;
        ros::Subscriber vive_pose_sub_R2, vive_pose_sub_L2;
        ros::Subscriber Myo_emg;

        std::shared_ptr<SerialManipulator> cManipulator;
        std::unique_ptr<HYUControl::Controller> Control;
        std::unique_ptr<HYUControl::Motion> motion;
    };
}

PLUGINLIB_EXPORT_CLASS(dualarm_controller::Impedance_Control,controller_interface::ControllerBase)


/*
 * Copyright 2020 California  Institute  of Technology (“Caltech”)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <x/vio/msckf_slam_update.h>
#include <x/vio/tools.h>
#include <x/ekf/state.h>
#include <boost/math/distributions.hpp>

using namespace x;
using namespace Eigen;

MsckfSlamUpdate::MsckfSlamUpdate(const x::TrackList& trks,
                                 const x::AttitudeList& quats,
                                 const x::TranslationList& pos,
                                 const Triangulation& triangulator,
                                 const MatrixXd& cov_s,
                                 const int n_poses_max,
                                 const double sigma_img)
{
  // Number of features
  const size_t n_trks = trks.size();

  // Number of feature observations
  size_t n_obs = 0;
  for(size_t i=0; i < n_trks; i++)
    n_obs += trks[i].size();

  // Initialize MSCKF Kalman update matrices
  const size_t rows0 = 2 * n_obs - n_trks * 3;
  const size_t cols = cov_s.cols();
  jac_ = MatrixXd::Zero(rows0, cols);
  cov_m_diag_ = VectorXd::Ones(rows0);
  res_ = MatrixXd::Zero(rows0, 1);

  // Initialize MSCKF-SLAM feature initialization matrices
  const size_t rows1 = n_trks * 3;
  init_mats_.H1 = MatrixXd::Zero(rows1, cols);
  init_mats_.H2 = MatrixXd::Zero(rows1, n_trks * 3);
  init_mats_.r1 = MatrixXd::Zero(rows1, 1);
  init_mats_.features = MatrixXd::Zero(rows1, 1);

  // For each track, compute residual, Jacobian and covariance block
  const double var_img = sigma_img * sigma_img;
  size_t row_h = 0, row1 = 0;
  for (size_t i = 0; i < n_trks; ++i) {
    processOneTrack(trks[i],
                    quats,
                    pos,
                    triangulator,
                    cov_s,
                    n_poses_max,
                    var_img,
                    i,
                    row_h,
                    row1);
  }
}

void MsckfSlamUpdate::processOneTrack(const x::Track& track,
                                      const x::AttitudeList& C_q_G,
                                      const x::TranslationList& G_p_C,
                                      const Triangulation& triangulator,
																      const MatrixXd& P,
                                      const int n_poses_max,
                                      const double var_img,
                                      const size_t& j,
                                      size_t& row_h,
                                      size_t& row1)
{
  const size_t track_size = track.size();
  unsigned int rows_track_j = track_size * 2;
  const size_t cols = P.cols();
  MatrixXd h_j(MatrixXd::Zero(rows_track_j, cols));
  MatrixXd Hf_j(MatrixXd::Zero(h_j.rows(), kJacCols));
  MatrixXd res_j(MatrixXd::Zero(rows_track_j, 1));

  // Feature triangulation
  Vector3d feature; // inverse-depth parameters in last observation frame
  triangulator.triangulateGN(track, feature);
  const double alpha = feature(0);
  const double beta  = feature(1);
  const double rho   = feature(2);

  // Anchor pose
  x::Quaternion Cn_q_G;
  Cn_q_G.x() = C_q_G.back().ax;
  Cn_q_G.y() = C_q_G.back().ay;
  Cn_q_G.z() = C_q_G.back().az;
  Cn_q_G.w() = C_q_G.back().aw;

  Vector3d G_p_Cn(G_p_C.back().tx, G_p_C.back().ty, G_p_C.back().tz);

  // Coordinate of feature in global frame
  Vector3d G_p_fj = 1 / (rho)*Cn_q_G.normalized().toRotationMatrix() * Vector3d(alpha, beta, 1) + G_p_Cn;

  x::Quatern attitude_to_quaternion;

  // LOOP OVER ALL FEATURE OBSERVATIONS
  for (size_t i = 0; i < track_size; ++i)
  {
    const unsigned int pos = C_q_G.size() - track_size + i;

    Quaterniond Ci_q_G_ = attitude_to_quaternion(C_q_G[pos]);
    Vector3d G_p_Ci_(G_p_C[pos].tx, G_p_C[pos].ty, G_p_C[pos].tz);

    // Feature position expressed in camera frame.
    Vector3d Ci_p_fj;
    Ci_p_fj << Ci_q_G_.normalized().toRotationMatrix().transpose() * (G_p_fj - G_p_Ci_);

    // eq. 20(a)
    Vector2d z;
    z(0) = track[i].getX();
    z(1) = track[i].getY();

    Vector2d z_hat(z);
    assert(Ci_p_fj(2));
    z_hat(0) = Ci_p_fj(0) / Ci_p_fj(2);
    z_hat(1) = Ci_p_fj(1) / Ci_p_fj(2);

    // eq. 20(b)
    res_j(i * 2, 0) = z(0) - z_hat(0);
    res_j(i * 2 + 1, 0) = z(1) - z_hat(1);

    //============================
    // Measurement Jacobian matrix
    //============================

    if (i == track_size - 1)  // Handle special case
    {
      // Inverse-depth feature coordinates jacobian
      MatrixXd mat(MatrixXd::Zero(2, 3));
      mat(0, 0) = 1.0;
      mat(1, 1) = 1.0;

      // Update stacked Jacobian matrices associated to the current feature
      unsigned int row = i * kVisJacRows;
      Hf_j.block<kVisJacRows, kJacCols>(row, 0) = mat;
    }
    else
    {
      // Set Jacobian of pose for i'th measurement of feature j (eq.22, 23)
      VisJacBlock J_i(VisJacBlock::Zero());
      // first row
      J_i(0, 0) = 1.0 / Ci_p_fj(2);
      J_i(0, 1) = 0.0;
      J_i(0, 2) = -Ci_p_fj(0) / std::pow((double)Ci_p_fj(2), 2);
      // second row
      J_i(1, 0) = 0.0;
      J_i(1, 1) = 1.0 / Ci_p_fj(2);
      J_i(1, 2) = -Ci_p_fj(1) / std::pow((double)Ci_p_fj(2), 2);

      // Attitude
      Vector3d skew_vector = Ci_q_G_.normalized().toRotationMatrix().transpose() * (G_p_fj - G_p_Ci_);
      VisJacBlock J_attitude = J_i * x::Skew(skew_vector(0), skew_vector(1), skew_vector(2)).matrix;

      // Position
      VisJacBlock J_position = -J_i * Ci_q_G_.normalized().toRotationMatrix().transpose();

      // Anchor attitude
      VisJacBlock J_anchor_att = -1 / rho * J_i * Ci_q_G_.normalized().toRotationMatrix().transpose() *
                                      Cn_q_G.normalized().toRotationMatrix() * x::Skew(alpha, beta, 1).matrix;

      // Anchor position
      VisJacBlock J_anchor_pos = -J_position;

      // Inverse-depth feature coordinates
      MatrixXd mat(MatrixXd::Identity(3, 3));
      mat(0, 2) = -alpha / rho;
      mat(1, 2) = -beta / rho;
      mat(2, 2) = -1 / rho;
      VisJacBlock Hf_j1 = 1 / rho * J_i * Ci_q_G_.normalized().toRotationMatrix().transpose() *
                               Cn_q_G.normalized().toRotationMatrix() * mat;

      // Update stacked Jacobian matrices associated to the current feature
      unsigned int row = i * kVisJacRows;
      Hf_j.block<kVisJacRows, kJacCols>(row, 0) = Hf_j1;

      unsigned int col = pos * kJacCols;
      h_j.block<kVisJacRows, kJacCols>(row, kSizeCoreErr + col) = J_position;

      col += n_poses_max * kJacCols;
      h_j.block<kVisJacRows, kJacCols>(row, kSizeCoreErr + col) = J_attitude;

      col = (C_q_G.size() - 1) * kJacCols;
      h_j.block<kVisJacRows, kJacCols>(row, kSizeCoreErr + col) = J_anchor_pos;

      col += n_poses_max * kJacCols;
      h_j.block<kVisJacRows, kJacCols>(row, kSizeCoreErr + col) = J_anchor_att;
    }
  }  // LOOP OVER ALL FEATURE OBSERVATIONS

  //========================================================================
  // Left nullspace projection
  //========================================================================
  // Nullspace computation
  MatrixXd q = Hf_j.householderQr().householderQ();
  MatrixXd A = x::MatrixBlock(q, 0, 3);

  // Projections
  MatrixXd res0_j = A.transpose() * res_j;
  MatrixXd h0_j = A.transpose() * h_j;

  // New noise measurement matrix
  VectorXd r0_j_diag = var_img * VectorXd::Ones(rows_track_j - 3);
  MatrixXd r0_j = r0_j_diag.asDiagonal();

  //========================================================================
  // Column space projection
  //========================================================================
  // Only needed for persistent feature init (Li, 2012)

	// Column space
	const MatrixXd U = x::MatrixBlock(q, 0, 0, q.rows(), 3);

	// Projections to be used in Core::CorrectAfterCoreCorrection
	MatrixXd H1j = U.transpose() * h_j;
	MatrixXd H2j = U.transpose() * Hf_j;
	MatrixXd r1j = U.transpose() * res_j;

  init_mats_.H1.block(row1, 0, 3, cols) = H1j;
	init_mats_.H2.block(row1, row1, 3, 3) = H2j;
	init_mats_.r1.block(row1, 0, 3, 1)	= r1j;
	init_mats_.features.block(row1, 0, 3, 1)= feature;
	row1 += 3;

  //==========================================================================
  // Outlier rejection
  //==========================================================================
  MatrixXd S_inv = (h0_j * P * h0_j.transpose() + r0_j).inverse();
  MatrixXd gamma = res0_j.transpose() * S_inv * res0_j;
  boost::math::chi_squared_distribution<> my_chisqr(2 * track_size - 3);  // 2*Mj-3 DoFs
  double chi = quantile(my_chisqr, 0.95);                            // 95-th percentile

  if (gamma(0, 0) < chi)  // Inlier
  {
#ifdef VERBOSE
    inliers_.push_back(G_p_fj);
#endif

    jac_.block(row_h,            // startRow
             0,                 // startCol
             rows_track_j - 3,  // numRows
             cols) = h0_j;    // numCols

    // Residual vector [feature j]
    res_.block(row_h, 0, rows_track_j - 3, 1) = res0_j;

    // Measurement covariance matrix [feature j]
    cov_m_diag_.segment(row_h, rows_track_j - 3) = r0_j_diag;

    row_h += rows_track_j - 3;
	}
  else  // outlier
  {
#ifdef VERBOSE
    outliers_.push_back(G_p_fj);
#endif
  }
}


#include <boost/graph/spectral_embedding_3d_layout.hpp>
#include "graph_types.hpp"
#include <boost/graph/python/point2d.hpp>
#include <boost/graph/python/point3d.hpp>
#include <boost/graph/python/graph.hpp>
#include <boost/python.hpp>

namespace boost { namespace graph { namespace python {

template <typename Graph>
void
spectral_embedding_layout
  (Graph& g,
   const vector_property_map<point3d, typename Graph::VertexIndexMap>* in_pos)
{
  using boost::python::object;

  typedef typename property_map<Graph, vertex_index_t>::const_type VertexIndexMap;
  typedef vector_property_map<point3d, VertexIndexMap> PositionMap;
  typedef vector_property_map<vertex_index_t, VertexIndexMap> IndexMap;

  PositionMap pos = *in_pos;

  boost::spectral_embedding_layout(g, pos);
}

void export_spectral_embedding_layout()
{
  using boost::python::arg;
  using boost::python::def;
  using boost::python::object;

  def("spectral_embedding_layout",
      &spectral_embedding_layout<Graph>,
      (arg("graph"),
       arg("position") =
         (vector_property_map<point2d, Graph::VertexIndexMap>*)0));

}

} } } // end namespace boost::graph::python


#include <cctbx/boost_python/flex_fwd.h>

#include <cctbx/sgtbx/sym_equiv_sites.h>
#include <boost/python/class.hpp>
#include <boost/python/args.hpp>
#include <boost/python/return_value_policy.hpp>
#include <boost/python/copy_const_reference.hpp>
#include <boost/python/return_internal_reference.hpp>

namespace cctbx { namespace sgtbx { namespace boost_python {

namespace {

  struct sym_equiv_sites_wrappers
  {
    typedef sym_equiv_sites<> w_t;

    static void
    wrap()
    {
      using namespace boost::python;
      typedef return_value_policy<copy_const_reference> ccr;
      typedef return_internal_reference<> rir;
      class_<w_t>("sym_equiv_sites", no_init)
        .def(init<site_symmetry const&>((
          arg("site_symmetry"))))
        .def(init<uctbx::unit_cell const&,
                  sgtbx::space_group const&,
                  fractional<> const&,
                  sgtbx::site_symmetry_ops const&> ((
          arg("unit_cell"), arg("space_group"), arg("original_site"),
          arg("site_symmetry_ops"))))
        .def(init<wyckoff::mapping const&>((
          arg("wyckoff_mapping"))))
        .def(init<sgtbx::space_group const&,
                  fractional<> const&,
                  optional<uctbx::unit_cell const&> >((
          arg("space_group"), arg("original_site"), arg("unit_cell"))))
        .def(init<uctbx::unit_cell const&,
                  sgtbx::space_group const&,
                  fractional<> const&,
                  rt_mx const&,
                  optional<std::size_t> >((
          arg("unit_cell"), arg("space_group"), arg("original_site"),
          arg("special_op"), arg("multiplicity"))))
        .def(init<uctbx::unit_cell const&,
                  sgtbx::space_group const&,
                  fractional<> const&,
                  optional<double, double> >((
          arg("unit_cell"), arg("space_group"), arg("original_site"),
          arg("minimum_distance"), arg("tolerance"))))
        .def("unit_cell", &w_t::unit_cell, rir())
        .def("space_group", &w_t::space_group, rir())
        .def("original_site", &w_t::original_site, ccr())
        .def("special_op", &w_t::special_op, ccr())
        .def("max_accepted_tolerance", &w_t::max_accepted_tolerance)
        .def("coordinates", &w_t::coordinates)
        .def("sym_op", &w_t::sym_op, (arg("i_coor")))
        .def("sym_op_indices", &w_t::sym_op_indices)
        .def("is_special_position", &w_t::is_special_position)
      ;
    }
  };

  struct min_sym_equiv_distance_info_wrappers
  {
    typedef min_sym_equiv_distance_info<> w_t;

    static void
    wrap()
    {
      using namespace boost::python;
      typedef return_value_policy<copy_const_reference> ccr;
      class_<w_t>("min_sym_equiv_distance_info", no_init)
        .def(init<sym_equiv_sites<> const&,
                  fractional<> const&,
                  optional<af::tiny<bool, 3> const&> >(
          (arg("reference_sites"),
           arg("other"),
           arg("principal_continuous_allowed_origin_shift_flags")
             =no_continuous_allowed_shifts)))
        .def(init<sym_equiv_sites<> const&,
                  af::const_ref<scitbx::vec3<double> > const&,
                  optional<af::tiny<bool, 3> const&> >(
          (arg("reference_sites"),
           arg("others"),
           arg("principal_continuous_allowed_origin_shift_flags")
             =no_continuous_allowed_shifts)))
        .def("i_other", &w_t::i_other)
        .def("sym_op", &w_t::sym_op, ccr())
        .def("continuous_shifts", &w_t::continuous_shifts, ccr())
        .def("diff", &w_t::diff, ccr())
        .def("dist", &w_t::dist)
        .def("apply", &w_t::apply, (arg("sites_frac")))
      ;
    }
  };

} // namespace <anoymous>

  void wrap_sym_equiv_sites()
  {
    sym_equiv_sites_wrappers::wrap();
    min_sym_equiv_distance_info_wrappers::wrap();
  }

}}} // namespace cctbx::sgtbx::boost_python


#include <cctbx/boost_python/flex_fwd.h>

#include <cctbx/maptbx/statistics.h>
#include <boost/python/class.hpp>
#include <boost/python/def.hpp>
#include <boost/python/args.hpp>

namespace cctbx { namespace maptbx { namespace boost_python {

namespace {

  struct statistics_wrappers
  {
    typedef statistics<> w_t;

    static void
    wrap()
    {
      using namespace boost::python;
      class_<w_t>("statistics", no_init)
        .def(init<af::const_ref<float, af::flex_grid<> > const&>())
        .def(init<af::const_ref<double, af::flex_grid<> > const&>())
        .def("min", &w_t::min)
        .def("max", &w_t::max)
        .def("mean", &w_t::mean)
        .def("mean_sq", &w_t::mean_sq)
        .def("sigma", &w_t::sigma)
      ;
    }
  };

  struct more_statistics_wrappers
  {
    typedef more_statistics<> w_t;

    static void
    wrap()
    {
      using namespace boost::python;
      class_<w_t, bases<statistics<> > >("more_statistics", no_init)
        .def(init<af::const_ref<float, af::flex_grid<> > const&>())
        .def(init<af::const_ref<double, af::flex_grid<> > const&>())
        .def("skewness", &w_t::skewness)
        .def("kurtosis", &w_t::kurtosis)
      ;
    }
  };

  struct even_more_statistics_wrappers
  {
    static void
    wrap()
    {
      typedef mem_iteration<> w_t;
      using namespace boost::python;
      class_<w_t>("mem_iteration", no_init)
        .def(init<af::ref<double, af::c_grid<3> > const&,
                  af::ref<double, af::c_grid<3> > const&,
                  af::ref<double, af::c_grid<3> >,
                  double,
                  af::tiny<int, 3> const&,
                  double,
                  double,
                  bool >())
        .def("tp", &w_t::tp)
        .def("scale", &w_t::scale)
        .def("z", &w_t::z)
        .def("hw", &w_t::hw)
        .def("hn", &w_t::hn)
      ;
    }
  };

} // namespace <anoymous>

  void wrap_statistics()
  {
    statistics_wrappers::wrap();
    more_statistics_wrappers::wrap();
    even_more_statistics_wrappers::wrap();
  }

}}} // namespace cctbx::maptbx::boost_python


/*
Copyright (c) 2020 Naomasa Matsubayashi

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

#include <cstdint>
#include <cmath>
#include <array>
#include <iostream>
#include <fstream>
#include <algorithm>
#include <iterator>
#include <charconv>
#include <random>
#include <boost/container/flat_map.hpp>
#include <boost/program_options.hpp>
#include <omp.h>
#include <nlohmann/json.hpp>
#include "ifm/exp_match.h"
int main( int argc, char *argv[] ) {
  boost::program_options::options_description options("オプション");
  options.add_options()
    ("help,h",    "ヘルプを表示")
    ("decay,d", boost::program_options::value<bool>()->default_value( true ), "decay")
    ("alpha,a", boost::program_options::value<float>()->default_value( 1 ), "a")
    ("beta,b", boost::program_options::value<float>()->default_value( 1 ), "b")
    ("length,l", boost::program_options::value<float>()->default_value( 1 ), "l");
  boost::program_options::variables_map params;
  boost::program_options::store( boost::program_options::parse_command_line( argc, argv, options ), params );
  boost::program_options::notify( params );
  if( params.count("help") ) {
    std::cout << options << std::endl;
    return 0;
  }
  float x0, x1, x2, y0, y1, y2;
  float a = params["alpha"].as<float>();
  float b = params["beta"].as<float>();
  if( params["decay"].as<bool>() ) {
    const auto [d1,d2] = ifm::approxymate_decay( a, b, params["length"].as<float>() );
    x0 = 0;
    x1 = d1;
    x2 = d2;
    y0 = ifm::get_exp_envelope( a, b, x0 );
    y1 = ifm::get_exp_envelope( a, b, x1 );
    y2 = ifm::get_exp_envelope( a, b, x2 );
    std::cout << d1 << " " << d2 << std::endl;
  }
  else {
    x0 = 0;
    x1 = ifm::approxymate_attack( a, b, params["length"].as<float>() );
    x2 = params["length"].as<float>();
    y0 = ifm::get_exp_envelope( a, b, x0 );
    y1 = ifm::get_exp_envelope( a, b, x1 );
    y2 = ifm::get_exp_envelope( a, b, x2 );
    std::cout << a << std::endl;
  }
  const auto [tangent0,shift0] = ifm::get_linear_interpolation( x0, y0, x1, y1 );
  const auto [tangent1,shift1] = ifm::get_linear_interpolation( x1, y1, x2, y2 );
  float diff = 0;
  for( size_t i = 0; i != 1000; ++i ) {
    float x = i * params["length"].as<float>() / 1000.f;
    float approx_y;
    if( x < x1 ) approx_y = x * tangent0 + shift0;
    else if( x < x2 ) approx_y = x * tangent1 + shift1;
    else approx_y = y2;
    float y = ifm::get_exp_envelope( a, b, x );
    diff += approx_y - y;
  }
  std::cout << "max(" << tangent0 << "x+" << shift0 << "," << tangent1 << "x+" << shift1 << "," << y2 << ")" << std::endl;
  std::cout << "diff: " << diff/1000.f << std::endl;
}


/*
 * word2vec.cpp
 *
 */
#include <fstream>
#include <random>
#include <set>
#include <string>
#include <utility>

#include <boost/assert.hpp>

#include "core/utils.h"
#include "core/functions.h"
#include "core/updaters.h"
#include "layers/contlayer.h"
#include "layers/fclayer.h"
#include "layers/smaxlayer.h"

#include "word2vec.h"

using namespace std;
using namespace boost;
using namespace yann;
using namespace yann::word2vec;

namespace yann::word2vec {

class DataSource_SkipGram : public Trainer::DataSource
{
  typedef vector<pair<MatrixSize, MatrixSize>> TestsVector;

public:
  DataSource_SkipGram(
      const Text & text,
      const MatrixSize & window_size,
      const MatrixSize & batch_size
  ) :
    _text(text),
    _window_size(window_size),
    _batch_size(batch_size),
    _cur_pos(0)
  {
    YANN_CHECK_GT(_batch_size, 0);
    YANN_CHECK_GT(_text.get_dictionary_size(), 0);

    _tests = get_tests(_text, _window_size);
    resize_batch(_inputs_batch, _batch_size, _text.get_dictionary_size());
    resize_batch(_outputs_batch, _batch_size, _text.get_dictionary_size());
  }
  virtual ~DataSource_SkipGram()
  {
  }

  // Trainer::DataSource overwrites
  virtual std::string get_info() const
  {
    ostringstream oss;
    oss << "yann::DataSource_SkipGram("
        << "dictionary_size=" << _text.get_dictionary_size()
        << ", batch_size=" << _batch_size
        << ", window_size=" << _window_size
        << ", tests_num=" << _tests.size()
        << ")";
    return oss.str();
  }
  virtual MatrixSize get_batch_size() const
  {
    return _batch_size;
  }

  virtual MatrixSize get_num_batches() const
  {
    YANN_CHECK(_batch_size > 0);
    return _tests.size() / _batch_size;
  }

  virtual void start_epoch()
  {
    _cur_pos = 0;

    // todo: support sequential mode w/o shuffling too
    auto rng = default_random_engine { };
    shuffle(begin(_tests), end(_tests), rng);
  }

  virtual optional<Batch> get_next_batch()
  {
    YANN_CHECK_GE(_cur_pos, 0);
    auto res = create_batch(_cur_pos);
    if(res <= 0) {
      return boost::none;
    }
    _cur_pos += res;
    return Batch(_inputs_batch, _outputs_batch, get_batch_size()); // each row in the batch is a test
  }
  virtual void end_epoch()
  {
    // do nothing
  }

public:
  double test_all(const Network & nn)
  {
    size_t success_count = 0, total_count = 0;

    VectorBatch outputs;
    outputs.resizeLike(_outputs_batch);
    for(MatrixSize start_pos = 0; ; ) {
      auto res = create_batch(start_pos);
      if(res <= 0) {
        break;
      }
      start_pos += res;

      nn.calculate(_inputs_batch, outputs);
      for(MatrixSize jj = 0; jj < yann::get_batch_size(outputs); ++jj) {
        const auto & actual = get_batch(outputs, jj);
        const auto & expected = get_batch(_outputs_batch, jj);
        MatrixSize aa = 0, ee = 0;
        actual.array().maxCoeff(&aa);
        expected.array().maxCoeff(&ee);
        if(aa == ee) {
          ++success_count;
        }
        ++total_count;
      }
    }

    return total_count > 0 ? success_count / (double)total_count : 0.0;
  }

private:
  MatrixSize create_batch(const MatrixSize & start_pos)
  {
    YANN_CHECK_GE(start_pos, 0);

    _inputs_batch.setZero();
    _outputs_batch.setZero();
    for(MatrixSize ii = 0, pos = start_pos; ii < _batch_size; ++ii, ++pos) {
      if(pos >= (MatrixSize)_tests.size()) {
        return 0;
      }
      const auto & cur =  _tests[pos];
      _inputs_batch.insert(ii, cur.first) = 1.0; // RowMajor
      _outputs_batch(ii, cur.second) = 1.0; // RowMajor
    }
    return _batch_size;
  }

  static TestsVector get_tests(const Text & text, const MatrixSize & window_size)
  {
    TestsVector tests;
    for(MatrixSize ii = 0, size = text.sentences().size(); ii < size; ++ii) {
      const auto & sentence = text.sentences()[ii];
      // sentence[jj] is the "input" word
      for(MatrixSize jj = 0, sentence_size = sentence.size(); jj < sentence_size; ++jj) {
        // sentence[kk] is the "output" word
        for(MatrixSize kk = jj - window_size; kk <= jj + window_size; ++kk) {
          // check if we are in-range and that we are not pointing to the word itself
          if(kk < 0 || kk >= sentence_size || sentence[jj] == sentence[kk]) {
            continue;
          }
          tests.push_back(make_pair(sentence[jj], sentence[kk]));
        }
      }
    }
    return tests;
  }

private:
  const Text & _text;
  const MatrixSize _window_size;
  const MatrixSize _batch_size;

  TestsVector _tests;
  size_t      _cur_pos;

  SparseVectorBatch _inputs_batch;
  VectorBatch _outputs_batch;
}; // class DataSource_SkipGram

class DataSource_CBOW : public Trainer::DataSource
{
  typedef set<MatrixSize> BagOfWords;
  typedef vector<pair<MatrixSize, BagOfWords>> TestsVector;

public:
  DataSource_CBOW(
      const Text & text,
      const MatrixSize & window_size,
      const MatrixSize & batch_size
  ) :
    _text(text),
    _window_size(window_size),
    _batch_size(batch_size),
    _cur_pos(0)
  {
    YANN_CHECK_GT(_batch_size, 0);
    YANN_CHECK_GT(_text.get_dictionary_size(), 0);

    _tests = get_tests(_text, _window_size);
    resize_batch(_inputs_batch, _batch_size, _text.get_dictionary_size());
    resize_batch(_outputs_batch, _batch_size, _text.get_dictionary_size());
  }
  virtual ~DataSource_CBOW()
  {
  }

  // Trainer::DataSource overwrites
  virtual std::string get_info() const
  {
    ostringstream oss;
    oss << "yann::DataSource_CBOW("
        << "dictionary_size=" << _text.get_dictionary_size()
        << ", batch_size=" << _batch_size
        << ", window_size=" << _window_size
        << ", tests_num=" << _tests.size()
        << ")";
    return oss.str();
  }
  virtual MatrixSize get_batch_size() const
  {
    return _batch_size;
  }

  virtual MatrixSize get_num_batches() const
  {
    YANN_CHECK(_batch_size > 0);
    return _tests.size() / _batch_size;
  }
  virtual void start_epoch()
  {
    _cur_pos = 0;

    // todo: support sequential mode w/o shuffling too
    auto rng = default_random_engine { };
    shuffle(begin(_tests), end(_tests), rng);
  }

  virtual optional<Batch> get_next_batch()
  {
    YANN_CHECK_GE(_cur_pos, 0);
    auto res = create_batch(_cur_pos);
    if(res <= 0) {
      return boost::none;
    }
    _cur_pos += res;
    return Batch(_inputs_batch, _outputs_batch, get_batch_size()); // each row in the batch is a test
  }

  virtual void end_epoch()
  {
    // do nothing
  }

public:
  double test_all(const Network & nn)
  {
    size_t success_count = 0, total_count = 0;

    VectorBatch outputs;
    outputs.resizeLike(_outputs_batch);
    for(MatrixSize start_pos = 0; ; ) {
      auto res = create_batch(start_pos);
      if(res <= 0) {
        break;
      }
      start_pos += res;

      nn.calculate(_inputs_batch, outputs);
      for(MatrixSize jj = 0; jj < yann::get_batch_size(outputs); ++jj) {
        const auto & actual = get_batch(outputs, jj);
        const auto & expected = get_batch(_outputs_batch, jj);
        MatrixSize aa = 0, ee = 0;
        actual.array().maxCoeff(&aa);
        expected.array().maxCoeff(&ee);
        if(aa == ee) {
          ++success_count;
        }
        ++total_count;
      }
    }

    return total_count > 0 ? success_count / (double)total_count : 0.0;
  }

private:
  MatrixSize create_batch(const MatrixSize & start_pos)
  {
    YANN_CHECK_GE(start_pos, 0);

    _inputs_batch.setZero();
    _outputs_batch.setZero();
    for(MatrixSize ii = 0, pos = start_pos; ii < _batch_size; ++ii, ++pos) {
      if(pos >= (MatrixSize)_tests.size()) {
        return 0;
      }
      const auto & cur =  _tests[pos];
      _inputs_batch.insert(ii, cur.first) = 1.0; // RowMajor
      for(const auto & val : cur.second) {
        _outputs_batch(ii, val) = 1.0; // RowMajor
      }
    }
    return _batch_size;
  }

  static TestsVector get_tests(const Text & text, const MatrixSize & window_size)
  {
    TestsVector tests;
    for(MatrixSize ii = 0, size = text.sentences().size(); ii < size; ++ii) {
      const auto & sentence = text.sentences()[ii];
      // sentence[jj] is the "current" word
      for(MatrixSize jj = 0, sentence_size = sentence.size(); jj < sentence_size; ++jj) {
        // sentence[kk] is the "bag" word
        BagOfWords bow;
        for(MatrixSize kk = jj - window_size; kk <= jj + window_size; ++kk) {
          // check if we are in-range and that we are not pointing to the word itself
          if(kk < 0 || kk >= sentence_size || sentence[jj] == sentence[kk]) {
            continue;
          }
          bow.insert(sentence[kk]);
        }
        tests.push_back(make_pair(sentence[jj], bow));
      }
    }
    return tests;
  }

private:
  const Text & _text;
  const MatrixSize _window_size;
  const MatrixSize _batch_size;

  TestsVector _tests;
  size_t      _cur_pos;

  SparseVectorBatch _inputs_batch;
  VectorBatch _outputs_batch;
}; // class DataSource_CBOW

}; // namespace yann::word2vec

////////////////////////////////////////////////////////////////////////////////////////////////
//
// Word2Vec::TrainingParams implementation
//
yann::word2vec::Word2Vec::TrainingParams::TrainingParams():
    _window_size(5),
    _dimensions(100),
    _updater(make_unique<Updater_GradientDescent>(0.9)),
    _training_sampling_rate(0.0),
    _training_batch_size(100),
    _epochs(10),
    _epochs_callback(nullptr),
    _batch_callback(nullptr)
{
}

////////////////////////////////////////////////////////////////////////////////////////////////
//
// Word2Vec implementation
//

// the format is (d:<dictionary>,v:<vectors>)
ostream& std::operator<<(ostream & os, const Word2Vec & w2v)
{
  os << "(d:" << w2v._dictionary << ",v:" << w2v._vectors << ")";
  return os;
}

// the format is (d:<dictionary>,v:<vectors>)
istream& std::operator>>(istream & is, Word2Vec & w2v)
{
  char ch;
  if(is >> ch && ch != '(') {
    is.putback(ch);
    is.setstate(std::ios_base::failbit);
    return is;
  }
  if(is >> ch && ch != 'd') {
    is.putback(ch);
    is.setstate(std::ios_base::failbit);
    return is;
  }
  if(is >> ch && ch != ':') {
    is.putback(ch);
    is.setstate(std::ios_base::failbit);
    return is;
  }
  is >> w2v._dictionary;
  if(is.fail()) {
    return is;
  }
  if(is >> ch && ch != ',') {
    is.putback(ch);
    is.setstate(std::ios_base::failbit);
    return is;
  }
  if(is >> ch && ch != 'v') {
    is.putback(ch);
    is.setstate(std::ios_base::failbit);
    return is;
  }
  if(is >> ch && ch != ':') {
    is.putback(ch);
    is.setstate(std::ios_base::failbit);
    return is;
  }
  is >> w2v._vectors;
  if(is.fail()) {
    return is;
  }
  if(is >> ch && ch != ')') {
    is.putback(ch);
    is.setstate(std::ios_base::failbit);
    return is;
  }

  return is;
}

// cos alpha = A * B / (|A| * |B|)
Value yann::word2vec::Word2Vec::cosine(const RefConstVector & v1, const RefConstVector & v2)
{
  return (v1.array() * v2.array()).sum() / (v1.norm() * v2.norm());
}

yann::word2vec::Word2Vec::Word2Vec()
{
}

yann::word2vec::Word2Vec::~Word2Vec()
{
}

bool yann::word2vec::Word2Vec::is_equal(const Word2Vec& other, double tolerance) const
{
  return _dictionary == other._dictionary && _vectors.isApprox(other._vectors, tolerance);
}

optional<MatrixSize> yann::word2vec::Word2Vec::find_word(const std::string & word) const
{
  return _dictionary.find_word(word);
}
optional<const std::string &> yann::word2vec::Word2Vec::find_word(const MatrixSize & word_num) const
{
  return _dictionary.find_word(word_num);
}

optional<RefConstVector> yann::word2vec::Word2Vec::find_vector(const MatrixSize & word_num) const
{
  YANN_CHECK_LT(word_num, get_vectors_size());
  RefConstVector vv = _vectors.row(word_num);
  return vv;
}

optional<RefConstVector> yann::word2vec::Word2Vec::find_vector(const std::string & word) const
{
  auto word_num = find_word(word);
  if(!word_num) {
    return boost::none;
  }
  return find_vector(*word_num);
}

Value yann::word2vec::Word2Vec::distance(const MatrixSize & word_num1, const MatrixSize & word_num2) const
{
  auto vv1 = find_vector(word_num1);
  auto vv2 = find_vector(word_num2);
  return (vv1 && vv2) ? cosine(*vv1, *vv2) : 0;
}

Value yann::word2vec::Word2Vec::distance(const std::string word1, const std::string & word2) const
{
  auto vv1 = find_vector(word1);
  auto vv2 = find_vector(word2);
  return (vv1 && vv2) ? cosine(*vv1, *vv2) : 0;
}

vector<pair<const string &, Value>> yann::word2vec::Word2Vec::find_closest(const string & word, const size_t & num) const
{
  YANN_CHECK_GT(num, 0);

  auto word_num = find_word(word);
  if(!word_num) {
    return vector<pair<const string &, Value>>(); // not found
  }
  auto vv = find_vector(*word_num);
  YANN_CHECK(vv);

  // compare by the vector distance
  auto cmp = [](const auto & aa, const auto & bb) { return aa.second < bb.second; };
  set<pair<MatrixSize, Value>, decltype(cmp)> ordered(cmp);
  for(MatrixSize ii = 0; ii < _vectors.rows(); ++ii) {
    if(ii == *word_num) continue; // ignore itself

    auto distance = cosine(*vv, _vectors.row(ii));
    ordered.insert(make_pair(ii, distance));
    if(ordered.size() > num) {
      ordered.erase(ordered.begin());
    }
  }

  vector<pair<const string &, Value>> res;
  res.reserve(ordered.size());
  for(auto it = ordered.rbegin(); it != ordered.rend(); ++it) {
    // can't use make_pair() because of "const string &"
    res.push_back(pair<const string &, Value>(_dictionary.get_word(it->first), it->second));
  }
  return res;
}

vector<pair<const string &, Value>> yann::word2vec::Word2Vec::find_farthest(const string & word, const size_t & num) const
{
  YANN_CHECK_GT(num, 0);

  auto word_num = find_word(word);
  if(!word_num) {
    return vector<pair<const string &, Value>>(); // not found
  }
  auto vv = find_vector(*word_num);
  YANN_CHECK(vv);

  // compare by the vector distance
  auto cmp = [](const auto & aa, const auto & bb) { return aa.second > bb.second; };
  set<pair<MatrixSize, Value>, decltype(cmp)> ordered(cmp);
  for(MatrixSize ii = 0; ii < _vectors.rows(); ++ii) {
    if(ii == *word_num) continue; // ignore itself

    auto distance = cosine(*vv, _vectors.row(ii));
    ordered.insert(make_pair(ii, distance));
    if(ordered.size() > num) {
      ordered.erase(ordered.begin());
    }
  }

  vector<pair<const string &, Value>> res;
  res.reserve(ordered.size());
  for(auto it = ordered.rbegin(); it != ordered.rend(); ++it) {
    // can't use make_pair() because of "const string &"
    res.push_back(pair<const string &, Value>(_dictionary.get_word(it->first), it->second));
  }
  return res;
}

void yann::word2vec::Word2Vec::save(const string & filename) const
{
  ofstream ofs(filename, ofstream::out | ofstream::trunc);
  if(!ofs || ofs.fail()) {
    throw runtime_error("can't open file " + filename);
  }
  ofs << (*this);
  if(!ofs || ofs.fail()) {
    throw runtime_error("can't write to file " + filename);
  }
  ofs.close();
}

void yann::word2vec::Word2Vec::load(const std::string & filename)
{
  ifstream ifs(filename, ifstream::in);
  if(!ifs || ifs.fail()) {
    throw runtime_error("can't open file " + filename);
  }
  ifs >> (*this);
  if(!ifs || ifs.fail()) {
    throw runtime_error("can't read file " + filename);
  }
  ifs.close();
}

template<typename DataSourceType>
std::unique_ptr<Word2Vec> yann::word2vec::Word2Vec::train(
    const Text & text,
    const TrainingParams & params,
    DataSourceType & data_source)
{
  auto w2v = make_unique<Word2Vec>();
  YANN_CHECK(w2v);

  // create network
  auto nn = make_unique<Network>();
  YANN_CHECK(nn);
  nn->set_cost_function(make_unique<CrossEntropyCost>());

  auto fc_layer1 = make_unique<FullyConnectedLayer>(
      text.get_dictionary_size(), params._dimensions);
  fc_layer1->set_fixed_bias(0.0);
  fc_layer1->set_activation_function(make_unique<IdentityFunction>());
  if(0 < params._training_sampling_rate && params._training_sampling_rate < 1.0) {
    fc_layer1->set_sampling_rate(params._training_sampling_rate);
  }
  nn->append_layer(std::move(fc_layer1));

  auto fc_layer2 = make_unique<FullyConnectedLayer>(
      params._dimensions, text.get_dictionary_size());
  fc_layer2->set_fixed_bias(0.0);
  fc_layer2->set_activation_function(make_unique<FastSigmoidFunction>());
  if(0 < params._training_sampling_rate && params._training_sampling_rate < 1.0) {
    fc_layer2->set_sampling_rate(params._training_sampling_rate);
  }
  nn->append_layer(std::move(fc_layer2));
  nn->init(Layer::InitMode_Random, params._training_init_context);

  // train the network
  Trainer trainer(params._updater);
  trainer.set_batch_progress_callback(params._batch_callback);
  trainer.set_epochs_progress_callback(params._epochs_callback);
  trainer.train(*nn, data_source, params._training_batch_size,  params._epochs);

  // extract the weights
  auto fc_layer = dynamic_cast<const FullyConnectedLayer*>(nn->get_layer(0));
  YANN_CHECK(fc_layer);
  w2v->_vectors = fc_layer->get_weights();
  w2v->_dictionary = text.dictionary();

  // done
  return w2v;
}

unique_ptr<Word2Vec> yann::word2vec::Word2Vec::train_skip_gram(const Text & text, const TrainingParams & params)
{
  DataSource_SkipGram data_source(
        text,
        params._window_size,
        params._training_batch_size);
  return train(text, params, data_source);
}

std::unique_ptr<Word2Vec> yann::word2vec::Word2Vec::train_cbow(const Text & text, const TrainingParams & params)
{
  DataSource_CBOW data_source(
        text,
        params._window_size,
        params._training_batch_size);
  return train(text, params, data_source);
}