//State Vector- variable acceleration model
	// X(0) = quad_position.point.x;
    // X(1) = quad_velocity.point.x;
    // X(2) = quad_acc.point.x;
    // X(3) = quad_position.point.y;
    // X(4) = quad_velocity.point.y;
    // X(5) = quad_acc.point.y;
    // X(6) = quad_position.point.z;
    // X(7) = quad_velocity.point.z;
    // X(8) = quad_acc.point.z;
	// X(9)  = leader.point.x;
    // X(10) = leader.point.y;
    // X(11) = leader.point.z;
	// X(12) = neighbor.point.x;
	// X(13) = neighbor.point.y;
	// X(14) = neighbor.point.z;
///measeurement Vector
    //1.
	// Y(0) = quad.imu.acc.x
	// Y(1) = 0
	// Y(2) = 0
	// Y(3) = quad.imu.acc.y
	// Y(4) = 0
	// Y(5) = 0
	// Y(6) = quad.imu.acc.z
	// Y(7) = 0
	// Y(8) = 0
    //2.
    // Y(0) = quad.imu.vel.x
    // Y(1) = 0
    // Y(2) = 0
    // Y(3) = quad.imu.vel.y
    // Y(4) = 0
    // Y(5) = 0
    // Y(6) = quad.imu.vel.z
    // Y(7) = 0
    // Y(8) = 0
    //3.
	// Y() = leader_quad.point.x
	// Y() = leader_quad.point.y
	// Y() = leader_quad.point.z
    //4.
	// Y() = neighbor_quad.point.x
	// Y() = neighbor_quad.point.y
	// Y() = neighbor_quad.point.z

#include <Eigen/Dense>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <bits/stdc++.h>
#include <ros/ros.h>
#include <math.h>
#include <stdio.h>
#include <dwm1001/anchor.h>
#include <sensor_msgs/Imu.h>
#include "std_msgs/MultiArrayLayout.h"
#include "std_msgs/MultiArrayDimension.h"
#include "std_msgs/Float32MultiArray.h"
#include "geometry_msgs/PoseStamped.h"
#include "geometry_msgs/Vector3Stamped.h"
#include "geometry_msgs/TwistStamped.h"

// Declarations

double r_pred_l,r_meas_l;
Eigen::MatrixXd H_l(1,15);
double R_l;
double m_R_scale_l,sigma_r_l;
Eigen::VectorXd K_l(15);
Eigen::VectorXd X_p_l(15);
Eigen::MatrixXd P_p_l(15,15);
double error_l;
double precisionRangeErrEst_l;
//
double r_pred_n,r_meas_n;
Eigen::MatrixXd H_n(1,15);
double R_n;
double m_R_scale_n,sigma_r_n;
Eigen::VectorXd K_n(15);
Eigen::VectorXd X_p_n(15);
Eigen::MatrixXd P_p_n(15,15);
double error_n;
double precisionRangeErrEst_n;
//
geometry_msgs::Vector3Stamped pose;
geometry_msgs::Vector3Stamped vel;
geometry_msgs::Vector3Stamped acc_;

Eigen::Quaternionf q;
Eigen::MatrixXf R_mat(3,3);
Eigen::MatrixXf imuacc(3,1);
Eigen::MatrixXf acc(3,1);
double r_acc;
Eigen::MatrixXd r_meas_acc(3,1);
Eigen::MatrixXd r_pred_acc(3,1);
Eigen::MatrixXd H_acc(3,15);
Eigen::MatrixXd R_acc(3,3);
Eigen::MatrixXd K_acc(15,3);
Eigen::MatrixXd P_p_acc(15,15);
Eigen::VectorXd X_p_acc(15);
//
double r_vel;
Eigen::MatrixXd r_meas_vel(3,1);
Eigen::MatrixXd r_pred_vel(3,1);
Eigen::MatrixXd H_vel(3,15);
Eigen::MatrixXd R_vel(3,3);
Eigen::MatrixXd K_vel(15,3);
Eigen::MatrixXd P_p_vel(15,15);
Eigen::VectorXd X_p_vel(15);
//
double r_pos_l;
Eigen::MatrixXd pos_meas_l(3,1);
Eigen::MatrixXd pos_pred_l(3,1);
Eigen::MatrixXd H_pos_l(3,15);
Eigen::MatrixXd R_pos_l(3,3);
Eigen::MatrixXd K_pos_l(15,3);
Eigen::MatrixXd P_l_pos(15,15);
Eigen::VectorXd X_l_pos(15);
//
double r_pos_n;
Eigen::MatrixXd pos_meas_n(3,1);
Eigen::MatrixXd pos_pred_n(3,1);
Eigen::MatrixXd H_pos_n(3,15);
Eigen::MatrixXd R_pos_n(3,3);
Eigen::MatrixXd K_pos_n(15,3);
Eigen::MatrixXd P_n_pos(15,15);
Eigen::VectorXd X_n_pos(15);
//
Eigen::VectorXd X(15);
Eigen::VectorXd X_e(15);
Eigen::MatrixXd F(15,15);
Eigen::MatrixXd block_F(3,3);
Eigen::MatrixXd Q(15,15);
Eigen::MatrixXd P(15,15);
Eigen::MatrixXd M(15,15);
Eigen::MatrixXd block_Q(3,3);
Eigen::MatrixXd Q_lead(3,3);
Eigen::MatrixXd Q_neigh(3,3);
double tao_bias,m_y_damping_factor;
double T,error;
double m_last_range_time;
std_msgs::Float32MultiArray output;
double error_threshold,precisionRangeMm;
double T_sq,m_tao_bias_sqrt,T_cub,m_z_damping_factor,m_Q_scale;
double x,y,z,x_l,y_l,z_l,x_n,y_n,z_n;
double q_l,q_n;

void setState(Eigen::VectorXd Y)
{
    X = Y;
}
void setCovariance(Eigen::MatrixXd S)
{
    P = S;
}

void prediction_step()
{
    //ROS_WARN("Acceleration:## %f, ## %f, ## %f", ax,ay,az);

    T = ros::Time::now().toSec() - m_last_range_time;
    if(T>1){
        T = 1;
    } else if(T<0){
        T = 0.01;
    }

    T_sq = std::pow(T,2);
    T_cub = std::pow(T,3);
    // F is a 9x9 State Transition Matrix
    F = Eigen::MatrixXd::Zero(15,15);
    block_F << 1, T, -T_sq/2.0,
               0, 1, -T,
               0, 0, 1;
    F.block<3,3>(0,0) = block_F;
    F.block<3,3>(3,3) = block_F;
    F.block<3,3>(6,6) = block_F;
    F.block<6,6>(9,9) = Eigen::MatrixXd::Identity(6,6);

        int static count =0 ;
    if(count == 0)
       {ROS_WARN("prediction_step");
        count++;}
        //std::cout << F << '\n'<<'\n'; }
    // X is the predicted state vector and the predicted covariance matrix
    X_e = F*X ;
    // Q is the acceleration model
    tao_bias = m_tao_bias_sqrt * m_tao_bias_sqrt;
    Q = Eigen::MatrixXd::Zero(15,15);

    block_Q << (T_cub*T_sq)/20.0,   (T_sq*T_sq)/8.0 ,   -T_cub/6,
        	   (T_sq*T_sq)/8.0 ,    (T_cub)/3,          -T_sq/2,
         	    T_cub/6.0,		     -T_sq/2 ,               T  ;
    block_Q *= tao_bias;

    Q.block<3,3>(0,0) = block_Q;
    Q.block<3,3>(3,3) = block_Q ;
    Q.block<3,3>(6,6) = block_Q ;
    Q.block<3,3>(9,9) = Q_lead;
    Q.block<3,3>(12,12) = Q_neigh;
    Q *= m_Q_scale;
    //ROS_WARN("covariance:## %f, ## %f, ## %f",(T_cub*tao_acc/3.0)+(T_cub*T_sq)*tao_bias/20.0 , T*tao_acc+(T_cub*tao_bias/3), T*tao_bias);

    // if(count%50 == 1)
    // std::cout << X << '\n'<<'\n';
    // M is the predicted covariance matrix
	M = F*P*F.transpose() + Q;

    // time update
    m_last_range_time = ros::Time::now().toSec();
    //count++;
    // setState(X_e);
    // setCovariance(M);
        pose.header.stamp = ros::Time::now();
        pose.vector.x=X(0);
        pose.vector.y=X(3);
        pose.vector.z=X(6);

        vel.header.stamp = ros::Time::now();
        vel.vector.x=X(1);
        vel.vector.y=X(4);
        vel.vector.z=X(7);

        acc_.header.stamp = ros::Time::now();
        acc_.vector.x=X(2);
        acc_.vector.y=X(5);
        acc_.vector.z=X(8);
    //    fused.publish(output);

    if(error < error_threshold){
        //ROS_WARN("\n sucess too large: %f", error);
        setState(X_e);
        setCovariance(M);
        return ;
    } else {

        ROS_WARN("\n Estimate too large: %f", error);
        return ;
    }
    count++;
}

void correction_step_imu(const sensor_msgs::Imu::ConstPtr& msg)
{
        int static count =0 ;
    if(count == 0)
       {ROS_WARN("imu_step");
        count++; }

    //conerts the acceleration from body frame to earth frame(NED)
    q = Eigen::Quaternionf(msg->orientation.w, msg->orientation.x, msg->orientation.y, msg->orientation.z);
    R_mat= q.toRotationMatrix();

    imuacc << msg->linear_acceleration.x,msg->linear_acceleration.y,msg->linear_acceleration.z;
    acc= R_mat*imuacc;
    acc(0,0)=acc(0,0);
    acc(1,0)=acc(1,0);
    acc(2,0)=(acc(2,0)-9.8);

    r_pred_acc << X(2),X(5),X(8);

    r_meas_acc << acc(0,0),acc(1,0),acc(2,0);
    // K is the Kalman Gain
    R_acc = Eigen::MatrixXd::Identity(3,3) * (r_acc*r_acc);
    R_acc(1,1) = R_acc(1,1) * m_y_damping_factor;
    R_acc(2,2) = R_acc(2,2) * m_z_damping_factor;
    //ROS_WARN("R_matrix:## %f",R);
    R_acc = (H_acc*P*H_acc.transpose()) + R_acc ;
    K_acc = P*H_acc.transpose() * R_acc.inverse();
    // Update P for the a posteriori covariance matrix
    P_p_acc = ( Eigen::MatrixXd::Identity(15,15) - K_acc * H_acc ) * P;
    // Update the state
    X_p_acc = X + K_acc * (r_meas_acc - r_pred_acc);
    // decide to take the range info or not.

    if(error < error_threshold){
        setState(X_p_acc);
        setCovariance(P_p_acc);
        return ;
    } else {

        ROS_WARN("\n Estimate too large: %f", error);
        return ;
    }
}

void correction_step_vel(const geometry_msgs::TwistStamped::ConstPtr& msg)
{
        int static count =0 ;
    if(count == 0)
       {ROS_WARN("vel_step");
        count++; }

    r_pred_vel << X(1),X(4),X(7);

    r_meas_vel << msg->twist.linear.x,msg->twist.linear.y,msg->twist.linear.z;
    // K is the Kalman Gain
    R_vel = Eigen::MatrixXd::Identity(3,3) * (r_vel*r_vel);
    R_vel(1,1) = R_vel(1,1) * m_y_damping_factor;
    R_vel(2,2) = R_vel(2,2) * m_z_damping_factor;
    //ROS_WARN("R_matrix:## %f",R);
    R_vel = (H_vel * P * H_vel.transpose()) + R_vel ;
    K_vel = P * H_vel.transpose() * R_vel.inverse();
    // Update P for the a posteriori covariance matrix
    P_p_vel = ( Eigen::MatrixXd::Identity(15,15) - K_vel * H_vel ) * P;
    // Update the state
    X_p_vel = X + K_vel * (r_meas_vel - r_pred_vel);
    // decide to take the range info or not.

    if(error < error_threshold){
        //ROS_WARN("\n sucess velocity: %f", );

        setState(X_p_vel);
        setCovariance(P_p_vel);
        return ;
    } else {

        ROS_WARN("\n Estimate too large: %f", error);
        return ;
    }
}

void correction_step_leader(const dwm1001::anchor::ConstPtr& msg)
{
        int static count =0 ;
    if(count == 0)
       {ROS_WARN("leader_r_step");
        count++; }

	r_pred_l = std::sqrt( std::pow(X(0) -  X(9), 2) +
                          std::pow(X(3) - X(10), 2) +
                          std::pow(X(6) - X(11), 2) ) + 1e-5;

	r_meas_l = msg->range;
    // H is the linearized measurement matrix
    H_l << (X(0) -  X(9))/r_pred_l, 0, 0,
           (X(3) - X(10))/r_pred_l, 0, 0,
           (X(6) - X(11))/r_pred_l, 0, 0,
           (X(9) - X(0))/r_pred_l, (X(10) - X(3))/r_pred_l, (X(11) - X(6))/r_pred_l,
         	0,	0,	0;

    // K is the Kalman Gain
    sigma_r_l = double(precisionRangeErrEst_l) / 1000.0;
    R_l = std::pow(sigma_r_l,2) * m_R_scale_l;
    //ROS_WARN("R_matrix:## %f",R);
    K_l = P*H_l.transpose() / ( (H_l*P*H_l.transpose())(0,0) + R_l );
    // Update P for the a posteriori covariance matrix
    P_p_l = ( Eigen::MatrixXd::Identity(15,15) - K_l*H_l ) * P;
    // Return the measurement innovation
    error_l = std::fabs(r_meas_l - r_pred_l);
    // Update the state
    X_p_l = X + K_l * (r_meas_l - r_pred_l);
    // decide to take the range info or not.
    if(error_l < error_threshold){
        //ROS_WARN("\n sucess too large: %f", error);
        setState(X_p_l);
        setCovariance(P_p_l);

        return ;
    } else {

        ROS_WARN("Anchor id , Update too large: %f -- %f", r_pred_l,r_meas_l);
        std::cout << X << "\n"<<'\n';
        return ;
    }

}

void correction_step_neigh(const dwm1001::anchor::ConstPtr& msg)
{
        int static count =0 ;
    if(count == 0)
       {ROS_WARN("neghbor_r_step");
        count++; }

	r_pred_n = std::sqrt( std::pow(X(0) - X(12), 2) +
                          std::pow(X(3) - X(13), 2) +
                          std::pow(X(6) - X(14), 2) ) + 1e-5;

	r_meas_n = msg->range;
    // H is the linearized measurement matrix
    H_n << (X(0) - X(12))/r_pred_n, 0, 0,
           (X(3) - X(13))/r_pred_n, 0, 0,
           (X(6) - X(14))/r_pred_n, 0, 0,
         	0,	0,	0,
           (X(12) - X(0))/r_pred_n, (X(13) - X(3))/r_pred_n, (X(14) - X(6))/r_pred_n ;

    // K is the Kalman Gain
    sigma_r_n = double(precisionRangeErrEst_n) / 1000.0;
    R_n = std::pow(sigma_r_n,2) * m_R_scale_n;
    //ROS_WARN("R_matrix:## %f",R);
    K_n = P*H_n.transpose() / ( (H_n*P*H_n.transpose())(0,0) + R_n );
    // Update P for the a posteriori covariance matrix
    P_p_n = ( Eigen::MatrixXd::Identity(15,15) - K_n*H_n ) * P;
    // Return the measurement innovation
    error_n = std::fabs(r_meas_n - r_pred_n);
    // Update the state
    X_p_n = X + K_n * (r_meas_n - r_pred_n);
    // decide to take the range info or not.
    if(error < error_threshold){
        //ROS_WARN("\n sucess too large: %f", error);
        setState(X_p_n);
        setCovariance(P_p_n);

        return ;
    } else {

        ROS_WARN("Anchor id , Update too large: %f --", error_n);
        std::cout << msg->device_id << "\n";
        return ;
    }

}

void position_lead_cb(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
        int static count =0 ;
    if(count == 0)
       {ROS_WARN("leader_pos_step");
        count++; }
    pos_meas_l << msg->pose.position.x,msg->pose.position.y,msg->pose.position.z;

    pos_pred_l << X(9),X(10),X(11);
    // K is the Kalman Gain
    R_pos_l = Eigen::MatrixXd::Identity(3,3) * (r_pos_l*r_pos_l);
    //ROS_WARN("R_matrix:## %f",R);
    R_pos_l = (H_pos_l * P * H_pos_l.transpose()) + R_pos_l;
    K_pos_l = P * H_pos_l.transpose() * R_pos_l.inverse();
    // Update P for the a posteriori covariance matrix
    P_l_pos = ( Eigen::MatrixXd::Identity(15,15) - K_pos_l * H_pos_l ) * P;
    // Update the state
    X_l_pos = X + K_pos_l * (pos_meas_l - pos_pred_l);
    // decide to take the range info or not.
    if(error < error_threshold){
        //ROS_WARN("\n sucess too large: %f", error);
        setState(X_l_pos);
        setCovariance(P_l_pos);
        return ;
    } else {

        ROS_WARN("\n Estimate too large: %f", error);
        return ;
    }
}

void position_neigh_cb(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
        int static count =0 ;
    if(count == 0)
       {ROS_WARN("neighbor_pos_step");
        count++; }
    pos_meas_n << msg->pose.position.x,msg->pose.position.y,msg->pose.position.z;

    pos_pred_n << X(12),X(13),X(14);

    // K is the Kalman Gain
    R_pos_n = Eigen::MatrixXd::Identity(3,3) * (r_pos_n*r_pos_n);
    //ROS_WARN("R_matrix:## %f",R);
    R_pos_n = (H_pos_n * P * H_pos_n.transpose()) + R_pos_n;
    K_pos_n = P * H_pos_n.transpose() * R_pos_n.inverse();
    // Update P for the a posteriori covariance matrix
    P_n_pos = ( Eigen::MatrixXd::Identity(15,15) - K_pos_n * H_pos_n ) * P;
    // Update the state
    X_n_pos = X + K_pos_n * (pos_meas_n - pos_pred_n);
    // decide to take the range info or not.
    if(error < error_threshold){
        //ROS_WARN("\n sucess too large: %f", error);
        setState(X_n_pos);
        setCovariance(P_n_pos);
        return ;
    } else {

        ROS_WARN("\n Estimate too large: %f", error);
        return ;
    }
}

void param(ros::NodeHandle& nh)
{
    nh.getParam("KalmanFilter_swarm/start_x", x);
    nh.getParam("KalmanFilter_swarm/start_y", y);
    nh.getParam("KalmanFilter_swarm/start_z", z);
    nh.getParam("KalmanFilter_swarm/lead_x", x_l);
    nh.getParam("KalmanFilter_swarm/lead_y", y_l);
    nh.getParam("KalmanFilter_swarm/lead_z", z_l);
    nh.getParam("KalmanFilter_swarm/neighbor_x", x_n);
    nh.getParam("KalmanFilter_swarm/neighbor_y", y_n);
    nh.getParam("KalmanFilter_swarm/neighbor_z", z_n);
    nh.getParam("KalmanFilter_swarm/leader_covariance", q_l);
    nh.getParam("KalmanFilter_swarm/neighbor_covariance", q_n);
    nh.getParam("KalmanFilter_swarm/m_tao_bias_sqrt", m_tao_bias_sqrt );
    nh.getParam("KalmanFilter_swarm/m_y_damping_factor", m_y_damping_factor);
    nh.getParam("KalmanFilter_swarm/m_z_damping_factor", m_z_damping_factor);
    nh.getParam("KalmanFilter_swarm/m_Q_scale", m_Q_scale);
    nh.getParam("KalmanFilter_swarm/error_threshold", error_threshold);
    nh.getParam("KalmanFilter_swarm/m_R_scale_leader", m_R_scale_l);
    nh.getParam("KalmanFilter_swarm/precisionRangeErrEst_leader", precisionRangeErrEst_l);
    nh.getParam("KalmanFilter_swarm/m_R_scale_neighbor", m_R_scale_n);
    nh.getParam("KalmanFilter_swarm/precisionRangeErrEst_neighbor", precisionRangeErrEst_n);
    nh.getParam("KalmanFilter_swarm/r_acc", r_acc);
    nh.getParam("KalmanFilter_swarm/r_vel", r_vel);
    nh.getParam("KalmanFilter_swarm/r_pos_l", r_pos_l);
    nh.getParam("KalmanFilter_swarm/r_pos_n", r_pos_n);

    ROS_WARN("%f, %f, %f ",x,y,z);
}

void Initialize(ros::NodeHandle& nh)
{

    param(nh);

    Eigen::MatrixXd nine_cov = Eigen::MatrixXd::Identity(15,15);
    nine_cov(0,0) = 0.001;
    nine_cov(3,3) = 0.001;
    nine_cov(6,6) = 0.001;
    // set cov of vel
    nine_cov(1,1) = 0.01;
    nine_cov(4,4) = 0.01;
    nine_cov(7,7) = 0.01;
    // set cov of acc_bias
    nine_cov(2,2) = 0.01;
    nine_cov(5,5) = 0.01;
    nine_cov(8,8) = 0.01;

    nine_cov(9,9)   = 0.001;
    nine_cov(10,10) = 0.001;
    nine_cov(11,11) = 0.001;

    nine_cov(12,12) = 0.001;
    nine_cov(13,13) = 0.001;
    nine_cov(14,14) = 0.001;

    setCovariance(nine_cov);
    X<<x,0,0,y,0,0,z,0,0,x_l,y_l,z_l,x_n,y_n,z_n;

    Q_lead <<  q_l,  0,  0,
                0, q_l, 0,
                0,  0,  q_l ;
    Q_neigh << q_n,  0,  0,
                0, q_n, 0,
                0,  0,  q_n ;
        // H is the linearized measurement matrix
    H_acc << 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0;
        // H is the linearized measurement matrix
    H_vel << 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0;
        // H is the linearized measurement matrix
    H_pos_l << 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0,
               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0;
        // H is the linearized measurement matrix
    H_pos_n << 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0,
               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1;

    m_last_range_time = ros::Time::now().toSec();

}

int main(int argc, char** argv){

    ros::init(argc,argv,"KalmanFilter_swarm");
    ros::NodeHandle nh;
    //Initializinging the parameters
    Initialize(nh);
    //param(nh);
    //Subscriber and Publisher for the data. remapped in the launch file to the topic required
    ros::Subscriber quad_imu = nh.subscribe("imu",20,correction_step_imu);
    ros::Subscriber quad_vel = nh.subscribe("velocity",20,correction_step_vel);
    ros::Subscriber anchor_1 = nh.subscribe("anchor_lead", 20,correction_step_leader);
    ros::Subscriber anchor_2 = nh.subscribe("anchor_neigh",20,correction_step_neigh);
    ros::Subscriber anchor_3 = nh.subscribe("position_leader",20,position_lead_cb);
    ros::Subscriber anchor_4 = nh.subscribe("position_neighour",20,position_neigh_cb);

    //ros::Publisher fused = nh.advertise<std_msgs::Float32MultiArray>("Filtered_data", 100);
    ros::Publisher fused_pose = nh.advertise<geometry_msgs::Vector3Stamped>("Filtered_pose", 10);
    ros::Publisher fused_vel = nh.advertise<geometry_msgs::Vector3Stamped>("Filtered_velocity", 10);
    ros::Publisher fused_acc = nh.advertise<geometry_msgs::Vector3Stamped>("Filtered_acc", 10);

    ros::Rate loop_rate(50);
    while(ros::ok()){
    //	output.data.clear();
	//	for (int i = 0; i < 15; i++)
	//		output.data.push_back(X(i));
        prediction_step();

        fused_pose.publish(pose);
        fused_vel.publish(vel);
        fused_acc.publish(acc_);

        ros::spinOnce();
        loop_rate.sleep();
    }
    return 0;
}


// load up a track and see if the closed-loop control even works in theory

#include <Eigen/Dense>

#include "drive/config.h"
#include "drive/controller.h"
#include "drive/trajtrack.h"

const float M_K1 = 2.58;  // DC motor response constants (measured)
const float M_K2 = 0.093;
const float M_K3 = 0.218;
const float M_OFFSET = 0.103;  // minimum control input (dead zone)

// return new velocity after applying throttle
float motor(float u, float v, float dt) {
  if (u >= 0) {
    float dv = u*M_K1 - u*M_K2*v - M_K3*v;
    return v + dv*dt;
  } else {
    float dv = u*M_K2*v - M_K3*v;
    return v + dv*dt;
  }
}

int main() {
  DriverConfig config;
  DriveController control;

  float x = 8, y = -7, theta = 3.14;
  float v = 0, w = 0;
  float s = 0;

  Eigen::Vector3f gyro(0, 0, 0), accel(0, 0, 0);

  const float dt = 1.0 / 30;

  for (int i = 0; i < 900; i++) {
    float throttle, steering;

    control.UpdateState(config, accel, gyro, s, dt);
    #if 0
    // broken now by a dependency on coneslam. oops.
    control.UpdateLocation(config, x, y, theta);
    if (!control.GetControl(config, 0, 0, &throttle, &steering, dt, true, i)) {
      break;
    }

    printf("%d xy %f %f theta %f v %f w %f input %f %f control ", i, x, y, theta, v, w, throttle, steering);
    control.Dump();
    printf("\n");
    #endif

    if (v > 0) {
      w = steering*v;
      gyro[2] = w;
    }
    v = motor(throttle, v, dt);
    theta += dt*w;
    float S = sin(theta), C = cos(theta);
    x += dt*v*C;
    y += dt*v*S;
    s += dt*v;
  }
}


#pragma once
#include <vector>
#include <memory>
#include <Eigen/Sparse>
#include "../DerivedPtrHolder.hh"

namespace kt84 {

template <int N>
struct LaplaceDirect_VertexTraits {
    struct Data {
        bool is_fixed;
        Eigen::Matrix<double, N, 1> value;
        Eigen::Matrix<double, N, 1> laplacian;

        Data()
            : is_fixed()
            , value    (Eigen::Matrix<double, N, 1>::Zero())
            , laplacian(Eigen::Matrix<double, N, 1>::Zero())
        {}
    } laplaceDirect;
};

struct LaplaceDirect_HalfedgeTraits {
    double laplaceDirect_weight;
    LaplaceDirect_HalfedgeTraits()
        : laplaceDirect_weight(1)
    {}
};

template <class TMeshBase, class TMesh, int N>
struct LaplaceDirect : public DerivedPtrHolder<TMesh, LaplaceDirect<TMeshBase, TMesh, N>> {
    typedef Eigen::Matrix<double,  N, 1> Value ;
    typedef Eigen::Matrix<double, -1, N> Vector;
    typedef Eigen::SparseMatrix <double> Matrix;

    double laplaceDirect_constraintWeight;
    Matrix L;
    std::shared_ptr<Eigen::SimplicialCholesky<Matrix>> solver;

    LaplaceDirect()
        : laplaceDirect_constraintWeight(1000.0)            // better to set large value as default? not really sure...
    {}

    void laplaceDirect_factorize() {
        TMesh* mesh = get_mesh();

        int nv = mesh->n_vertices();

        typedef Eigen::Triplet<double> Triplet;
        std::vector<Triplet> triplets;
        for (int i = 0; i < nv; ++i) {
            auto v = mesh->vertex_handle(i);

            triplets.push_back(Triplet(i, i, 1));

            double weight_sum = 0;
            for (auto h = mesh->voh_iter(v); h.is_valid(); ++h)
                weight_sum += mesh->data(*h).laplaceDirect_weight;

            for (auto h = mesh->voh_iter(v); h.is_valid(); ++h) {
                auto w = mesh->to_vertex_handle(*h);
                double weight = mesh->data(*h).laplaceDirect_weight / weight_sum;

                triplets.push_back(Triplet(i, w.idx(), -weight));
            }

            if (mesh->data(v).laplaceDirect.is_fixed)
                triplets.push_back(Triplet(i, i, laplaceDirect_constraintWeight));
        }

        L.resize(nv, nv);
        L.setFromTriplets(triplets.begin(), triplets.end());

        solver.reset(new Eigen::SimplicialCholesky<Matrix>(L.transpose() * L));
    }

    void laplaceDirect_solve() {
        TMesh* mesh = get_mesh();

        int nv = mesh->n_vertices();

        if (!solver || L.rows() != nv)
            laplaceDirect_factorize();

        // set right hand side
        Vector b(nv, N);
        for (int i = 0; i < nv; ++i) {
            auto& vdata = mesh->data(mesh->vertex_handle(i)).laplaceDirect;

            b.row(i) = vdata.laplacian.transpose();

            if (vdata.is_fixed)
                b.row(i) += laplaceDirect_constraintWeight * vdata.value.transpose();
        }

        // solve!
        Vector x = solver->solve(L.transpose() * b);

        // copy result to Data::value
        for (int i = 0; i < nv; ++i) {
            auto& vdata = mesh->data(mesh->vertex_handle(i)).laplaceDirect;

            if (!vdata.is_fixed)
                vdata.value = x.row(i).transpose();
        }
    }

    void laplaceDirect_set_laplacian_from_value() {
        TMesh* mesh = get_mesh();

        for (auto v : mesh->vertices()) {
            auto& vdata = mesh->data(v).laplaceDirect;

            vdata.laplacian = vdata.value;

            double weight_sum = 0;
            for (auto h = mesh->voh_iter(v); h.is_valid(); ++h)
                weight_sum += mesh->data(*h).laplaceDirect_weight;

            for (auto h = mesh->voh_iter(v); h.is_valid(); ++h) {
                auto w = mesh->to_vertex_handle(*h);
                auto& wdata = mesh->data(w).laplaceDirect;
                double weight = mesh->data(*h).laplaceDirect_weight / weight_sum;

                vdata.laplacian -= weight * wdata.value;
            }
        }
    }
private:
    TMesh* get_mesh() const { return DerivedPtrHolder<TMesh, LaplaceDirect<TMeshBase, TMesh, N>>::derived_ptr; }
};

}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_SQRT_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_SQRT_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!

    @ingroup group-arithmetic
    Function object implementing sqrt capabilities

    Computes the square root of its parameter. For integers it is the
    truncation of the real square root.

    @par semantic:
    For any given value @c x of type @c T:

    @code
    T r = sqrt(x);
    @endcode

    If full accuracy is not needed a sometimes faster less accurate version of the function
    can be reached using the fast_ decorator.

    @par Decorators

    std_, fast_ for floating entries

  **/
  const boost::dispatch::functor<tag::sqrt_> sqrt = {};
} }
#endif

#include <boost/simd/function/scalar/sqrt.hpp>
#include <boost/simd/function/simd/sqrt.hpp>

#endif


// ROS includes
#include "ros/ros.h"
#include "ros/assert.h"
#include "dynamic_reconfigure/server.h"
#include "create_driver/vicon_driver.h"
#include "geometry_msgs/Twist.h"
#include <sensor_msgs/Joy.h>
#include <geometry_msgs/Vector3.h>

// Library includes
#include <string>
#include <vector>
#include <map>
#include <Eigen/Core>
#include <Eigen/LU>
#include <algorithm>

#include <stdio.h>
//#include <tchar.h>
#include <iostream>
#include <fstream>
#include <iomanip>
#include <cmath>

#include "jingfu/control_formation.h"
#include <Eigen/Core>
#include <Eigen/Dense>
using Eigen::MatrixXd;
using namespace Eigen;

using namespace std;
using std::setw;
using std::setprecision;

int main(int argc, char **argv)
{
	// ROS Initalization
	ros::init(argc, argv, "control");

	ros::NodeHandle node;
	ros::NodeHandle private_node("~");

	// ROS Parameters

	// List of robot names
	vector<string> robot_names;
	XmlRpc::XmlRpcValue robot_list;
	private_node.getParam("robot_list", robot_list);
	ROS_ASSERT(robot_list.getType() == XmlRpc::XmlRpcValue::TypeArray);

	for (int i = 0; i < robot_list.size(); i++)
	{
		ROS_ASSERT(robot_list[i].getType() == XmlRpc::XmlRpcValue::TypeString);
		robot_names.push_back(static_cast<string>(robot_list[i]));
	}

	// Number of robots
	const int num_robots = robot_names.size();

	// Robot name
	string robot_name;
	private_node.param<std::string>("robotname", robot_name, "defaultname");

	int index = find(robot_names.begin(), robot_names.end(), robot_name) - robot_names.begin();

	// ROS Subscribers
	map<string, create_driver::ViconStream> vicon;
	vector<string>::iterator name_it;
	vector<ros::Subscriber> sub;
	for (name_it = robot_names.begin(); name_it != robot_names.end(); name_it++)
	{
		vicon.insert(pair<string, create_driver::ViconStream>(*name_it, create_driver::ViconStream()));
		sub.push_back(node.subscribe("/" + *name_it + "/tf", 10, &create_driver::ViconStream::callback, &vicon[*name_it]));
	}

	// ROS Publishers
	ros::Publisher vel = node.advertise<geometry_msgs::Twist>("cmd_vel", 10);

	// ROS loop
	ros::Rate loop_rate(250); // 250 Hz

	ros::spinOnce();
	/// initial poses from Vicon///
	pose2D q10; q10.x = vicon[robot_names[0]].x();q10.y = vicon[robot_names[0]].y();q10.h = vicon[robot_names[0]].theta();
	pose2D q20; q20.x = vicon[robot_names[1]].x();q20.y = vicon[robot_names[1]].y();q20.h = vicon[robot_names[1]].theta();
	pose2D q30; q30.x = vicon[robot_names[2]].x();q30.y = vicon[robot_names[2]].y();q30.h = vicon[robot_names[2]].theta();
	pose2D quad; quad.x = vicon[robot_names[3]].x();quad.y = vicon[robot_names[3]].y();quad.h = vicon[robot_names[3]].theta();

	cout << robot_names[0] <<endl;
	cout << robot_names[1] << endl;
	cout << robot_names[2] << endl;
	/// initlal offests ///
	double k = 0.6; // radius
	int mm = 3;

	pose2D q1off; q1off.x = k*cos(1*PI/mm); q1off.y = k*sin(1*PI/mm); q1off.h = 0;
	pose2D q2off; q2off.x = k*cos(3*PI/mm); q2off.y = k*sin(3*PI/mm); q2off.h = 0;
	pose2D q3off; q3off.x = k*cos(5*PI/mm); q3off.y = k*sin(5*PI/mm); q3off.h = 0;

//	pose2D q1off; q1off.x = 1; q1off.y = 0; q1off.h = 0;
//	pose2D q2off; q2off.x = 2; q2off.y = 0; q2off.h = 0;
//	pose2D q3off; q3off.x = 3; q3off.y = 0; q3off.h = 0;

	/// end of initial offsets ///

	double df12 = sqrt((q1off.x - q2off.x)*(q1off.x - q2off.x) + (q1off.y - q2off.y)*(q1off.y - q2off.y));
	double df13 = sqrt((q1off.x - q3off.x)*(q1off.x - q3off.x) + (q1off.y - q3off.y)*(q1off.y - q3off.y));
	double df23 = sqrt((q2off.x - q3off.x)*(q2off.x - q3off.x) + (q2off.y - q3off.y)*(q2off.y - q3off.y));

	pose2D q10_v; pose2D q20_v; pose2D q30_v;
	Vector2d p12; Vector2d p13; Vector2d p23;
	double d12=0; double d13=0; double d23=0;

	Wheel vd;
	vd.l_v = 0.8;
	vd.a_v = 0;

	/// initial control inputs ///
	Wheel MyVel[3];
	MyVel[0].l_v = 0;
	MyVel[0].a_v = 0;

	MyVel[1].l_v = 0;
	MyVel[1].a_v = 0;

	MyVel[2].l_v = 0;
	MyVel[2].a_v = 0;

	/////////////////////////////
	control_formation r1;
	//mobileRobotSystem TestSys1;
	control_formation r2;
	//mobileRobotSystem TestSys2;
	control_formation r3;
	//mobileRobotSystem TestSys3;

	MatrixXd A(3,3); // adjacency matrix
	A <<0,1,1,
	    1,0,1,
	    1,1,0;

	pose2D virtualGoal_1; pose2D virtualGoal_2; pose2D virtualGoal_3; // virtual goal
	point2D po1; point2D po2; point2D po3; // position of obstacle

	ofstream myfile_1("/home/heckert/utd-lars/src/jingfu/data/robot_1.txt",ios::out | ios::binary);
	ofstream myfile_2("/home/heckert/utd-lars/src/jingfu/data/robot_2.txt",ios::out | ios::binary);
	ofstream myfile_3("/home/heckert/utd-lars/src/jingfu/data/robot_3.txt",ios::out | ios::binary);
	ofstream myfile_4("/home/heckert/utd-lars/src/jingfu/data/quad.txt",ios::out | ios::binary);

	while (ros::ok())
	{
		// Retrieve Vicon Data
		ros::spinOnce();

		/* An instance of this code is run on each robot in the formation.
		 *
		 * You can access each robot's position and oreintation through the following statements:
		 * vicon[robot_names[i]].x()
		 * vicon[robot_names[i]].y()
		 * vicon[robot_names[i]].theta()
		 *
		 * The index of the robot running the instance of the code is contain in the "index"
		 * variable. For example, if you wanted to get the distance of this robot from all of its
		 * neighbors you could run the following:
		 * double dist[num_robots];
		 * for (int i = 0; i < num_robots; i++) {
		 * 		dist[i] = sqrt(pow(vicon[robot_names[i]].x() - vicon[robot_names[index]].x(),2)
		 * 				     + pow(vicon[robot_names[i]].y() - vicon[robot_names[index]].y(),2));
		 * }
		 */

		 /// virtual poses
		q10_v.x = q10.x + q1off.x; q10_v.y = q10.y + q1off.y; q10_v.h = q10.h + q1off.h;
		q20_v.x = q20.x + q2off.x; q20_v.y = q20.y + q2off.y; q20_v.h = q20.h + q2off.h;
		q30_v.x = q30.x + q3off.x; q30_v.y = q30.y + q3off.y; q30_v.h = q30.h + q3off.h;

		int kk = 2;
		///
		/// virtual goals ///
		virtualGoal_1.x = (q20_v.x + q30_v.x)/kk;
		virtualGoal_1.y = (q20_v.y + q30_v.y)/kk;
		virtualGoal_1.h = (q20_v.h + q30_v.h)/kk;

		virtualGoal_2.x = (q10_v.x + q30_v.x)/kk;
		virtualGoal_2.y = (q10_v.y + q30_v.y)/kk;
		virtualGoal_2.h = (q10_v.h + q30_v.h)/kk;

		virtualGoal_3.x = (q20_v.x + q10_v.x)/kk;
		virtualGoal_3.y = (q20_v.y + q10_v.y)/kk;
		virtualGoal_3.h = (q20_v.h + q10_v.h)/kk;

		//////////////////////////////////

		p12(0) = q10.x - q20.x; p12(1) = q10.y - q20.y;
		p13(0) = q10.x - q30.x; p13(1) = q10.y - q30.y;
		p23(0) = q20.x - q30.x; p23(1) = q20.y - q30.y;

		/// Distance between robots ///
		d12 = p12.norm();
		d13 = p13.norm();

		d23 = p23.norm();

		//cout << "e12 = " << d12 - df12 << endl;
		//cout << "e13 = " << d13 - df13 << endl;
		//cout << "e23 = " << d23 - df23 << endl;

 		double d1[2] = {d12,d13};
 		int nneighbor = std::min_element(d1,d1+2) - d1;

 		switch (nneighbor) {
 			case 0:
 				po1.x = q20.x;
				po1.y = q20.y;
 				break;
 			case 1:
 				po1.x = q30.x;
				po1.y = q30.y;
 				break;
 		}

 		double d2[2] = {d12,d23};
 		nneighbor = std::min_element(d2,d2+2) - d2;

 		switch (nneighbor) {
 			case 0:
 				po2.x = q10.x;
				po2.y = q10.y;
 				break;
 			case 1:
 				po2.x = q30.x;
				po2.y = q30.y;
 				break;
 		}

 		double d3[2] = {d13,d23};
 		nneighbor = std::min_element(d3,d3+2) - d3;

 		switch (nneighbor) {
 			case 0:
 				po3.x = q10.x;
				po3.y = q10.y;
 				break;
 			case 1:
 				po3.x = q20.x;
				po3.y = q20.y;
 				break;
 		}

		MyVel[0] = r1.switching_w_obstacle(q10_v, virtualGoal_1, q10, po1, vd);

		q10.x = vicon[robot_names[0]].x();
		q10.y = vicon[robot_names[0]].y();
		q10.h = vicon[robot_names[0]].theta();
		q10.h = r1.between2PI(q10.h);

		MyVel[1] = r2.switching_w_obstacle(q20_v, virtualGoal_2, q20, po2, vd);

		q20.x = vicon[robot_names[1]].x();
		q20.y = vicon[robot_names[1]].y();
		q20.h = vicon[robot_names[1]].theta();
		q20.h = r2.between2PI(q20.h);

		MyVel[2] = r3.switching_w_obstacle(q30_v, virtualGoal_3, q30, po3, vd);

		q30.x = vicon[robot_names[2]].x();
		q30.y = vicon[robot_names[2]].y();
		q30.h = vicon[robot_names[2]].theta();
		q30.h = r3.between2PI(q30.h);

		quad.x = vicon[robot_names[3]].x();
		quad.y = vicon[robot_names[3]].y();
		quad.h = vicon[robot_names[3]].theta();

		if ((MyVel[0].l_v-MyVel[1].l_v)<=0.008 && (MyVel[0].l_v-MyVel[2].l_v)<=0.008)
		{
			vd.l_v = vd.l_v - 0.01;
		}

		if (vd.l_v<=0)
		{
			vd.l_v = 0;
		}

		//cout<<"Current State x, y, theta"<<endl;
 		//cout<<q10.x<<setw(20)<<q10.y<<setw(20)<<q10.h<<endl;
		//cout<<q20.x<<setw(20)<<q20.y<<setw(20)<<q20.h<<endl;
		//cout<<q30.x<<setw(20)<<q30.y<<setw(20)<<q30.h<<endl;
 		//cout<<endl;
		/*
		cout<<"Current velocities" <<endl;
		cout<<MyVel[0].l_v<<setw(20)<<endl;
		cout<<MyVel[1].l_v<<setw(20)<<endl;
		cout<<MyVel[2].l_v<<setw(20)<<endl;
 		cout<<endl;
		*/

		myfile_1 <<q10.x<<setw(20)<<q10.y<<setw(20)<<q10.h<<setw(20)<<MyVel[0].l_v<<setw(20)<<MyVel[0].a_v<<endl;
		myfile_2 <<q20.x<<setw(20)<<q20.y<<setw(20)<<q20.h<<setw(20)<<MyVel[1].l_v<<setw(20)<<MyVel[1].a_v<<endl;
		myfile_3 <<q30.x<<setw(20)<<q30.y<<setw(20)<<q30.h<<setw(20)<<MyVel[2].l_v<<setw(20)<<MyVel[2].a_v<<endl;
		myfile_4 <<quad.x<<setw(20)<<quad.y<<setw(20)<<quad.h<<endl;

		double v = MyVel[index].l_v*300; // linear velocity output
		double w = MyVel[index].a_v; // angular velocity output

		// Send wheel velocities to driver
		geometry_msgs::Twist msg;
		msg.linear.x = v;
		msg.linear.y = 0;
		msg.linear.z = 0;
		msg.angular.x = 0;
		msg.angular.y = 0;
		msg.angular.z = w;
		vel.publish(msg);

		loop_rate.sleep();
	}
	myfile_1.close();
	myfile_2.close();
	myfile_3.close();
	myfile_4.close();
	return 0;
}


#ifndef MY_UTILS
#define MY_UTILS

#define MAX_COST 99999
#define TIME_INTERVAL 0.1
#define TRAJECTORY_TIME_INTERVAL 0.02
#define TRAJECTORY_LENGTH 50
#define TRAJECTORY_KEEP_MAX 25
#define COEFFNUMBER 6
#define HORIZONREF 5 // 5 seconds
#define MPSTOMPH 22.4
#define SENSORRANGE 50

#include <cfloat>
#include <Eigen/Dense>
#include <math.h>
#include <vector>

enum class State{KeepLane, PreLaneChangeLeft, PreLaneChangeRight, LaneChangeLeft, LaneChangeRight};
enum class MIOPoses{FrontLeft, FrontRight, Front, RearLeft, RearRight};

double polynominal(const Eigen::VectorXd& coeff, const double x){
    double result{0.0};
    for(int i = 0; i < coeff.size(); ++i){
        result += coeff(i) * pow(x, i);
    }
    return result;
}
#endif

int getLaneNumber(const double d){
    if(d >= 0 && d <= 4){
        return 0;
    }
    else if(d >4 && d <= 8){
        return 1;
    }
    else{
        return 2;
    }

}

double getCartesionVel(double x_vel, double y_vel){
    sqrt(pow(x_vel, 2) + pow(y_vel, 2));
}

// TODO: Backbone functions
std::vector<double> getTrajCoeff(Eigen::Vector3d init_cond, Eigen::Vector3d final_cond, double horizon_t){}

std::vector<State> getAvailableStates(State curr_state, int curr_d){
    int lane_num = getLaneNumber(curr_d);
    std::vector<State> nextStates;
    if(curr_state == State::KeepLane){
        nextStates.push_back(State::KeepLane);
        if(lane_num == 0){
            nextStates.push_back(State::PreLaneChangeRight);
        }
        else if(lane_num == 1){
            nextStates.push_back(State::PreLaneChangeLeft);
            nextStates.push_back(State::PreLaneChangeRight);
        }
        else{
            nextStates.push_back(State::PreLaneChangeLeft);
        }
    }
    else if(curr_state == State::PreLaneChangeLeft){
        nextStates.resize(3);
        nextStates[0] = State::PreLaneChangeLeft;
        nextStates[1] = State::KeepLane;
        nextStates[2] = State::LaneChangeLeft;
    }
    else if(curr_state == State::PreLaneChangeRight){
        nextStates.resize(3);
        nextStates[0] = State::PreLaneChangeRight;
        nextStates[1] = State::KeepLane;
        nextStates[2] = State::LaneChangeRight;
    }
    else{
        nextStates.resize(2);
        nextStates[0] = curr_state;
        nextStates[1] = State::KeepLane;
    }
    return nextStates;
}

std::vector<std::vector<double>> generateFinalContidionsKL(Eigen::Vector3d s_init, Eigen::Vector3d d_init, double horizon_t){}
std::vector<std::vector<double>> generateFinalContidionsPLC(Eigen::Vector3d s_init, Eigen::Vector3d d_init, double horizon_t){}
std::vector<std::vector<double>> generateFinalContidionsLC(Eigen::Vector3d s_init, Eigen::Vector3d d_init, double horizon_t, int target_lane){}

double calculateCost(const std::vector<double>& s_traj_coeff, const std::vector<double>& d_traj_coeff,
                     const Eigen::Vector3d& s_init, const Eigen::Vector3d& d_init, double horizon_t)
{

}

// cost functions


#ifndef ASLAM_BACKEND_OPTIMIZER_2_HPP
#define ASLAM_BACKEND_OPTIMIZER_2_HPP

#include <boost/shared_ptr.hpp>
//#include <boost/function.hpp>
#include <sm/assert_macros.hpp>
#include <Eigen/Core>
#include "Optimizer2Options.hpp"
#include "backend.hpp"
#include "OptimizationProblemBase.hpp"
#include "OptimizerCallbackManager.hpp"
#include <aslam/Exceptions.hpp>
#include <sm/timing/Timer.hpp>
#include <boost/thread.hpp>
#include <sparse_block_matrix/linear_solver.h>
#include <aslam/backend/util/CommonDefinitions.hpp>
#include <aslam/backend/TrustRegionPolicy.hpp>
#include <aslam/backend/LevenbergMarquardtTrustRegionPolicy.hpp>
#include <aslam/backend/GaussNewtonTrustRegionPolicy.hpp>
#include <aslam/backend/DogLegTrustRegionPolicy.hpp>
#include <aslam/backend/util/OptimizerProblemManagerBase.hpp>

namespace sm {

  class ConstPropertyTree;

}
namespace aslam {
  namespace backend {
    class LinearSystemSolver;

    /**
     * \class Optimizer2
     *
     * A sparse Gauss-Newton/LM optimizer.
     *
     * The notation in this file follows Harley and Zisserman, Appendix 6.
     *
     * Some Additions to the standard algorithm:
     * Choice of Lambda:
     * H.B. Nielsen. Damping Parameter in Marquardts Method. Technical Report
     * IMM-REP-1999-05, Technical University of Denmark, 1999.
     *
     * Jacobian Normalisation:
     * Jean-Claude Trigeassou, Thierry Poinot & Sandrine Moreau (2003):
     * A Methodology for Estimation of Physical Parameters, Systems Analysis Modelling Simulation, 43:7, 925-943
     *
     * Erik Etien, Damien Halbert, and Thierry Poinot
     * Improved JilesAtherton Model for Least Square Identification Using Sensitivity Function Normalization
     * IEEE TRANSACTIONS ON MAGNETICS, VOL. 44, NO. 7, JULY 2008
     *
     */
    class Optimizer2 : public OptimizerProblemManagerBase {
    public:
      typedef aslam::backend::Timer Timer;
      typedef sparse_block_matrix::SparseBlockMatrix<Eigen::MatrixXd> SparseBlockMatrix;
      typedef Optimizer2Options Options;
      struct Status : public OptimizerStatus {
        SolutionReturnValue srv;
       private:
        void resetImplementation() override;
      };

      SM_DEFINE_EXCEPTION(Exception, aslam::Exception);

      Optimizer2(const Options& options = Options());
      Optimizer2(const sm::ConstPropertyTree& config, boost::shared_ptr<LinearSystemSolver> linearSystemSolver, boost::shared_ptr<TrustRegionPolicy> trustRegionPolicy);
      ~Optimizer2() override;

      /// \brief initialize the linear solver specified in the optimizer options.
      void initializeLinearSolver();

      void initializeTrustRegionPolicy();

      /// \brief Run the optimization
      SolutionReturnValue optimize();

      /// \brief Return the status
      const Status& getStatus() const override { return _status; }

      /// \brief Const getter for the optimizer options.
      const Options& getOptions() const override { return _options; }

      /// \brief Get the optimizer options.
      Options& options() { return _options; }

      /// \brief Set the optimizer options.
      void setOptions(const Options& options) { _options = options; }

      /// \brief Set the optimizer options.
      void setOptions(const OptimizerOptionsBase& options) override { static_cast<OptimizerOptionsBase&>(_options) = options; }

      /// \brief return the reduced system dx
      const Eigen::VectorXd& dx() const;

      /// The value of the objective function.
      double J() const;

      /// \brief compute the full covariance matrix. This is expensive.
      void computeCovariances(SparseBlockMatrix& outP, double lambda);

      /// \brief compute only the diagonal covariance blocks.
      void computeDiagonalCovariances(SparseBlockMatrix& outP, double lambda);

      /// \brief compute only the covariance blocks associated with the block indices passed as an argument
      void computeCovarianceBlocks(const std::vector<std::pair<int, int> >& blockIndices, SparseBlockMatrix& outP, double lambda);

      void computeHessian(SparseBlockMatrix& outH, double lambda);

      /// \brief Evaluate the error at the current state.
      double evaluateError(bool useMEstimator);

      /// \brief Get dense design variable i.
      DesignVariable* designVariable(size_t i);

      /// \brief how many dense design variables are involved in the problem
      size_t numDesignVariables() const;

      /// \brief print the internal timing information.
      void printTiming() const;

      /// \brief Do a bunch of checks to see if the problem is well-defined. This includes checking that every error term is
      ///        hooked up to design variables and running finite differences on error terms where this is possible.
      void checkProblemSetup() override;

      /// \brief Build the Gauss-Newton matrices.
      void buildGnMatrices(bool useMEstimator);

      /// Returns the linear solver
      template <class L>
      L* getSolver(bool assertNonNull = true){
        L* solver = dynamic_cast<L*>(_solver.get());
        if(assertNonNull) {
          SM_ASSERT_TRUE(Exception, solver != NULL, "The solver is null");
        }
        return solver;
      }

      /// \brief expose callback registry
      callback::Registry & callback(){
        return _callbackManager;
      }

      const Matrix * getJacobian() const;

      const LinearSystemSolver * getBaseSolver() const;

    private:

      /// \brief Zero the Gauss-Newton matrices.
      void zeroMatrices();

      /// \brief Revert the last state update.
      void revertLastStateUpdate();

      /// \brief Apply a state update.
      double applyStateUpdate();

      /// \brief issue callback for given event
      template<typename Event>
      void issueCallback();

      void optimizeImplementation() override;

      void initializeImplementation() override;

      /// \brief The dense update vector.
      Eigen::VectorXd _dx;

      /// \brief The previous value of the cost function.
      double _p_J;

      boost::shared_ptr<LinearSystemSolver> _solver;

      boost::shared_ptr<TrustRegionPolicy> _trustRegionPolicy;

      /// \brief the current set of options
      Options _options;

      /// \brief the current status
      Status _status;

      /// \brief A class that manages the optimizer callbacks
      callback::Manager _callbackManager;
    };

} // namespace backend
} // namespace aslam

#endif /* ASLAM_BACKEND_OPTIMIZER_HPP */


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_CONSTANT_SIGNMASK_HPP_INCLUDED
#define BOOST_SIMD_CONSTANT_SIGNMASK_HPP_INCLUDED

/*!
  @ingroup group-constant
  @defgroup constant-Signmask Signmask (function template)

  Generates a constant able to mask the sign bit of any value.

  @headerref{<boost/simd/constant/signmask.hpp>}

  @par Description

  1.  @code
      template<typename T> T Signmask();
      @endcode

  2.  @code
      template<typename T> T Signmask( boost::simd::as_<T> const& target );
      @endcode

  Generates a value of type @c T that evaluates to a value where the most significant bit is set
  to 1 (includign for unsigned type).

  @par Parameters

  | Name                | Description                                                         |
  |--------------------:|:--------------------------------------------------------------------|
  | **target**          | a [placeholder](@ref type-as) value encapsulating the constant type |

  @par Return Value
  A value of type @c T that evaluates to `bitwise_cast<T>(1 << sizeof(scalar_of_t<T>)*8-1)`

  @par Requirements
  - **T** models Value
**/

#include <boost/simd/constant/scalar/signmask.hpp>
#include <boost/simd/constant/simd/signmask.hpp>

#endif


// Distributed under the MIT License.
// See LICENSE.txt for details.

#pragma once

#include <array>
#include <boost/optional.hpp>
#include <cstddef>
#include <limits>

#include "DataStructures/Tensor/TypeAliases.hpp"
#include "Domain/OrientationMap.hpp"
#include "Utilities/TypeTraits.hpp"

namespace PUP {
class er;
}  // namespace PUP

namespace CoordinateMaps {

/*!
 * \ingroup CoordinateMapsGroup
 *
 * \brief Three dimensional map from the cube to a wedge.
 * \image html Shell.png "A shell can be constructed out of six wedges."
 *
 * \details The mapping that goes from a reference cube to a three-dimensional
 *  wedge centered on a coordinate axis covering a volume between an inner
 *  surface and outer surface. Each surface can be given a curvature
 *  between flat (a sphericity of 0) or spherical (a sphericity of 1).
 *
 *  The first two logical coordinates correspond to the two angular coordinates,
 *  and the third to the radial coordinate.
 *
 *  The Wedge3D map is constructed by linearly interpolating between a bulged
 *  face of radius `radius_of_inner_surface` to a bulged face of
 *  radius `radius_of_outer_surface`, where the radius of each bulged face
 *  is defined to be the radius of the sphere circumscribing the bulge.
 *
 *  We make a choice here as to whether we wish to use the logical coordinates
 *  parameterizing these surface as they are, in which case we have the
 *  equidistant choice of coordinates, or whether to apply a tangent map to them
 *  which leads us to the equiangular choice of coordinates. In terms of the
 *  logical coordinates, the equiangular coordinates are:
 *
 *  \f[\textrm{equiangular xi} : \Xi(\xi) = \textrm{tan}(\xi\pi/4)\f]
 *
 *  \f[\textrm{equiangular eta}  : \mathrm{H}(\eta) = \textrm{tan}(\eta\pi/4)\f]
 *
 *  With derivatives:
 *
 *  \f[\Xi'(\xi) = \frac{\pi}{4}(1+\Xi^2)\f]
 *
 *  \f[\mathrm{H}'(\eta) = \frac{\pi}{4}(1+\mathrm{H}^2)\f]
 *
 *  The equidistant coordinates are:
 *
 *  \f[ \textrm{equidistant xi}  : \Xi = \xi\f]
 *
 *  \f[ \textrm{equidistant eta}  : \mathrm{H} = \eta\f]
 *
 *  with derivatives:
 *
 *  <center>\f$\Xi'(\xi) = 1\f$, and \f$\mathrm{H}'(\eta) = 1\f$</center>
 *
 *  We also define the variable \f$\rho\f$, given by:
 *
 *  \f[\textrm{rho} : \rho = \sqrt{1+\Xi^2+\mathrm{H}^2}\f]
 *
 *  ### The Spherical Face Map
 *  The surface map for the spherical face of radius \f$R\f$ lying in the
 * \f$+z\f$
 *  direction in either choice of coordinates is then given by:
 *
 *  \f[\vec{\sigma}_{spherical}: \vec{\xi} \rightarrow \vec{x}(\vec{\xi})\f]
 *  Where
 *  \f[
 *  \vec{x}(\xi,\eta) =
 *  \begin{bmatrix}
 *  x(\xi,\eta)\\
 *  y(\xi,\eta)\\
 *  z(\xi,\eta)\\
 *  \end{bmatrix}  = \frac{R}{\rho}
 *  \begin{bmatrix}
 *  \Xi\\
 *  \mathrm{H}\\
 *  1\\
 *  \end{bmatrix}\f]
 *
 *  ### The Bulged Face Map
 *  The bulged surface is itself constructed by linearly interpolating between
 *  a cubical face and a spherical face. The surface map for the cubical face
 *  of side length \f$2L\f$ lying in the \f$+z\f$ direction is given by:
 *
 *  \f[\vec{\sigma}_{cubical}: \vec{\xi} \rightarrow \vec{x}(\vec{\xi})\f]
 *  Where
 *  \f[
 *  \vec{x}(\xi,\eta) =
 *  \begin{bmatrix}
 *  x(\xi,\eta)\\
 *  y(\xi,\eta)\\
 *  L\\
 *  \end{bmatrix}  = L
 *  \begin{bmatrix}
 *  \Xi\\
 *  \mathrm{H}\\
 *  1\\
 *  \end{bmatrix}\f]
 *
 *  To construct the bulged map we interpolate between this cubical face map
 *  and a spherical face map of radius \f$R\f$, with the
 *  interpolation parameter being \f$s\f$. The surface map for the bulged face
 *  lying in the \f$+z\f$ direction is then given by:
 *
 *  \f[\vec{\sigma}_{bulged}(\xi,\eta) = {(1-s)L + \frac{sR}{\rho}}
 *  \begin{bmatrix}
 *  \Xi\\
 *  \mathrm{H}\\
 *  1\\
 *  \end{bmatrix}\f]
 *
 *  We constrain L by demanding that the spherical face circumscribe the cube.
 *  With this condition, we have \f$L = R/\sqrt3\f$.
 *  \note This differs from the choice in SpEC where it is demanded that the
 *  surfaces touch at the center, which leads to \f$L = R\f$.
 *
 *  ### The Full Volume Map
 *  The final map for the wedge which lies along the \f$+z\f$ is obtained by
 *  interpolating between the two surfaces with the
 *  interpolation parameter being the logical coordinate \f$\zeta\f$. This
 *  results in:
 *
 *  \f[\vec{x}(\xi,\eta,\zeta) =
 *  \frac{1}{2}\left\{(1-\zeta)\Big[(1-s_{inner})\frac{R_{inner}}{\sqrt 3}
 *   + s_{inner}\frac{R_{inner}}{\rho}\Big] +
 *  (1+\zeta)\Big[(1-s_{outer})\frac{R_{outer}}{\sqrt 3} +s_{outer}
 *  \frac{R_{outer}}{\rho}\Big] \right\}\begin{bmatrix}
 *  \Xi\\
 *  \mathrm{H}\\
 *  1\\
 *  \end{bmatrix}\f]
 *
 *  We will define the variables \f$F(\zeta)\f$ and \f$S(\zeta)\f$, the frustum
 * and sphere factors: \f[F(\zeta) = F_0 + F_1\zeta\f] \f[S(\zeta) = S_0 +
 * S_1\zeta\f]
 *  Where \f{align*}F_0 &= \frac{1}{2} \big\{ (1-s_{outer})R_{outer} +
 * (1-s_{inner})R_{inner}\big\}\\
 *  F_1 &= \partial_{\zeta} F = \frac{1}{2} \big\{ (1-s_{outer})R_{outer} -
 * (1-s_{inner})R_{inner}\big\}\\
 *  S_0 &= \frac{1}{2} \big\{ s_{outer}R_{outer} + s_{inner}R_{inner}\big\}\\
 *  S_1 &= \partial_{\zeta} S = \frac{1}{2} \big\{ s_{outer}R_{outer} -
 * s_{inner}R_{inner}\big\}\f}
 *
 *  The map can then be rewritten as:
 * \f[\vec{x}(\xi,\eta,\zeta) = \left\{\frac{F(\zeta)}{\sqrt 3} +
 * \frac{S(\zeta)}{\rho}\right\}\begin{bmatrix}
 *  \Xi\\
 *  \mathrm{H}\\
 *  1\\
 *  \end{bmatrix}\f]
 *
 *  We provide some common derivatives:
 *  \f[\partial_{\xi}z = \frac{-S(\zeta)\Xi\Xi'}{\rho^3}\f]
 *  \f[\partial_{\eta}z = \frac{-S(\zeta)\mathrm{H}\mathrm{H}'}{\rho^3}\f]
 * \f[\partial_{\zeta}z = \frac{F'}{\sqrt 3} + \frac{S'}{\rho}\f]
 *  The Jacobian then is: \f[J =
 *  \begin{bmatrix}
 *  \Xi'z + \Xi\partial_{\xi}z & \Xi\partial_{\eta}z & \Xi\partial_{\zeta}z \\
 *  \mathrm{H}\partial_{\xi}z & \mathrm{H}'z +
 *  \mathrm{H}\partial_{\eta}z & \mathrm{H}\partial_{\zeta}z\\
 *   \partial_{\xi}z&\partial_{\eta}z &\partial_{\zeta}z \\
 *  \end{bmatrix}
 *  \f]
 *
 *  A common factor that shows up in the inverse jacobian is:
 *  \f[ T:= \frac{S(\zeta)}{(\partial_{\zeta}z)\rho^3}\f]
 *
 *  The inverse Jacobian then is: \f[J^{-1} =
 *  \frac{1}{z}\begin{bmatrix}
 *  \Xi'^{-1} & 0 & -\Xi\Xi'^{-1}\\
 *  0 & \mathrm{H}'^{-1} & -\mathrm{H}\mathrm{H}'^{-1}\\
 *  T\Xi &
 *  T\mathrm{H} &
 *  T + F(\partial_{\zeta}z)^{-1}/\sqrt 3\\
 *  \end{bmatrix}
 *  \f]
 *
 *  ### Changing the radial distribution of the gridpoints
 *  By default, Wedge3D linearly distributes its gridpoints in the radial
 *  direction. An exponential distribution of gridpoints can be obtained by
 *  linearly interpolating in the logarithm of the radius, in order to obtain
 *  a relatively higher resolution at smaller radii. Since this is a radial
 *  rescaling of Wedge3D, this option is only supported for fully spherical
 *  wedges with `sphericity_inner` = `sphericity_outer` = 1.
 *
 *  The linear interpolation done is:
 *  \f[
 *  \log r = \frac{1-\zeta}{2}\log R_{inner} +
 *  \frac{1+\zeta}{2}\log R_{outer}
 *  \f]
 *
 *  The map then is:
 *  \f[\vec{x}(\xi,\eta,\zeta) =
 *  \frac{\sqrt{R_{inner}^{1-\zeta}R_{outer}^{1+\zeta}}}{\rho}\begin{bmatrix}
 *  \Xi\\
 *  \mathrm{H}\\
 *  1\\
 *  \end{bmatrix}\f]
 *
 *  The jacobian simplifies similarly.
 *
 */
class Wedge3D {
 public:
  static constexpr size_t dim = 3;
  enum class WedgeHalves {
    /// Use the entire wedge
    Both,
    /// Use only the upper logical half
    UpperOnly,
    /// Use only the lower logical half
    LowerOnly
  };

  /*!
   * Constructs a 3D wedge.
   * \param radius_inner Distance from the origin to one of the
   * corners which lie on the inner surface.
   * \param radius_outer Distance from the origin to one of the
   * corners which lie on the outer surface.
   * \param orientation_of_wedge The orientation of the desired wedge relative
   * to the orientation of the default wedge which is a wedge that has its
   * curved surfaces pierced by the upper-z axis. The logical xi and eta
   * coordinates point in the cartesian x and y directions, respectively.
   * \param sphericity_inner Value between 0 and 1 which determines
   * whether the inner surface is flat (value of 0), spherical (value of 1) or
   * somewhere in between
   * \param sphericity_outer Value between 0 and 1 which determines
   * whether the outer surface is flat (value of 0), spherical (value of 1) or
   * somewhere in between
   * \param with_equiangular_map Determines whether to apply a tangent function
   * mapping to the logical coordinates (for `true`) or not (for `false`).
   * \param halves_to_use Determines whether to use the logical xi
   * coordinates in the [0,1] interval (value of `UpperOnly`) of the full wedge,
   * the coordinates in the [-1,0] interval (value of `LowerOnly`) of the
   * full wedge, or the full wedge entirely (value of `Both`). Half wedges are
   * currently only useful in constructing domains for binary systems.
   * \param with_logarithmic_map Determines whether to apply an exponential
   * function mapping to the "sphere factor", the effect of which is to
   * distribute the radial gridpoints logarithmically in physical space.
   */
  Wedge3D(double radius_inner, double radius_outer,
          OrientationMap<3> orientation_of_wedge, double sphericity_inner,
          double sphericity_outer, bool with_equiangular_map,
          WedgeHalves halves_to_use = WedgeHalves::Both,
          bool with_logarithmic_map = false) noexcept;

  Wedge3D() = default;
  ~Wedge3D() = default;
  Wedge3D(Wedge3D&&) = default;
  Wedge3D(const Wedge3D&) = default;
  Wedge3D& operator=(const Wedge3D&) = default;
  Wedge3D& operator=(Wedge3D&&) = default;

  template <typename T>
  std::array<tt::remove_cvref_wrap_t<T>, 3> operator()(
      const std::array<T, 3>& source_coords) const noexcept;

  /// For a \f$+z\f$-oriented `Wedge3D`, returns invalid if \f$z<=0\f$
  /// or if \f$(x,y,z)\f$ is on or outside the cone defined
  /// by \f$(x^2/z^2 + y^2/z^2+1)^{1/2} = -S/F\f$, where
  /// \f$S = \frac{1}{2}(s_1 r_1 - s_0 r_0)\f$ and
  /// \f$F = \frac{1}{2\sqrt{3}}((1-s_1) r_1 - (1-s_0) r_0)\f$.
  /// Here \f$s_0,s_1\f$ and \f$r_0,r_1\f$ are the specified sphericities
  /// and radii of the inner and outer \f$z\f$ surfaces.  The map is singular on
  /// the cone and on the xy plane.
  boost::optional<std::array<double, 3>> inverse(
      const std::array<double, 3>& target_coords) const noexcept;

  template <typename T>
  tnsr::Ij<tt::remove_cvref_wrap_t<T>, 3, Frame::NoFrame> jacobian(
      const std::array<T, 3>& source_coords) const noexcept;

  template <typename T>
  tnsr::Ij<tt::remove_cvref_wrap_t<T>, 3, Frame::NoFrame> inv_jacobian(
      const std::array<T, 3>& source_coords) const noexcept;

  // clang-tidy: google runtime references
  void pup(PUP::er& p) noexcept;  // NOLINT

 private:
  // factors out calculation of z needed for mapping and jacobian
  template <typename T>
  tt::remove_cvref_wrap_t<T> default_physical_z(const T& zeta,
                                                const T& one_over_rho) const
      noexcept;
  friend bool operator==(const Wedge3D& lhs, const Wedge3D& rhs) noexcept;

  double radius_inner_{std::numeric_limits<double>::signaling_NaN()};
  double radius_outer_{std::numeric_limits<double>::signaling_NaN()};
  OrientationMap<3> orientation_of_wedge_{};
  double sphericity_inner_{std::numeric_limits<double>::signaling_NaN()};
  double sphericity_outer_{std::numeric_limits<double>::signaling_NaN()};
  bool with_equiangular_map_ = false;
  WedgeHalves halves_to_use_ = WedgeHalves::Both;
  bool with_logarithmic_map_ = false;
  double scaled_frustum_zero_{std::numeric_limits<double>::signaling_NaN()};
  double sphere_zero_{std::numeric_limits<double>::signaling_NaN()};
  double scaled_frustum_rate_{std::numeric_limits<double>::signaling_NaN()};
  double sphere_rate_{std::numeric_limits<double>::signaling_NaN()};
};
bool operator!=(const Wedge3D& lhs, const Wedge3D& rhs) noexcept;
}  // namespace CoordinateMaps


/*
 * Copyright 2019 Xilinx, Inc.
 *
 * 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.
 */

/**
 *  @brief Kernel implementation file
 *
 *  $DateTime: 2019/04/09 12:00:00 $
 */

#ifndef _HLS_FD_SOLVER_H_
#define _HLS_FD_SOLVER_H_
#include "ap_int.h"
#include "xf_fintech/spmv.hpp"
#include "xf_fintech/dimv.hpp"
#include "xf_fintech/trsv.hpp"
#ifndef __SYNTHESIS__
#include <assert.h>
#include <boost/algorithm/string.hpp>
#include <cmath>
#include <fstream>
#include <iostream>
#include <vector>
#endif

#define PRAGMA_SUB(x) _Pragma(#x)
#define DO_PRAGMA(x) PRAGMA_SUB(x)

namespace xf {
namespace fintech {
namespace internal {

/// @brief Get V-inner index from an S ordered vector
inline unsigned int V2S(const unsigned int index, const unsigned int m1, const unsigned int m2) {
    unsigned int row = index % m2;
    unsigned int col = index / m2;
    return row * m1 + col;
}

/// @brief Get S-inner index from an V ordered vector
inline unsigned int S2V(const unsigned int index, const unsigned int m1, const unsigned int m2) {
    unsigned int row = index % m1;
    unsigned int col = index / m1;
    return row * m2 + col;
}

/// @brief Class to encapsulate the Finite Difference engine components
template <typename DT,
          unsigned int M_SIZE,
          unsigned int LOG2_M_SIZE,
          unsigned int A_SIZE,
          unsigned int MEM_WIDTH,
          unsigned int DIM2_SIZE1 = 3,
          unsigned int DIM2_SIZE2 = 5>
class Solver {
   public:
    /// @brief default constructor
    Solver() {
#pragma HLS inline
    }
    /// @brief Copy and reorder S ordered vector into V inner form
    /// @param[in]  v_in  Vector representing array of size m1 x m2 flattened in S
    /// inner form
    /// @param[in]  m1   Size of array [0..m1-1] in S direction
    /// @param[in]  m2   Size of array [0..m2-1] in V direction
    /// @param[out] v_out Vector representing array of size m1 x m2 flattened in V
    /// inner form
    void reorderS2V(const DT v_in[M_SIZE], DT v_out[M_SIZE], const unsigned int m1, const unsigned int m2) {
#pragma HLS ARRAY_PARTITION variable = v_out cyclic factor = MEM_WIDTH dim = 1

        for (unsigned int i = 0; i < M_SIZE; ++i) {
#pragma HLS UNROLL factor = MEM_WIDTH
            v_out[S2V(i, m1, m2)] = v_in[i];
        }
    }

    /// @brief Copy and reorder V ordered vector into S inner form
    /// @param[in]  v_in  Vector representing array of size m1 x m2 flattened in V
    /// inner form
    /// @param[in]  m1   Size of array [0..m1] in S direction
    /// @param[in]  m2   Size of array [0..m2] in V direction
    /// @param[out] v_out Vector representing array of size m1 x m2 flattened in S
    /// inner form
    void reorderV2S(const DT v_in[M_SIZE], DT v_out[M_SIZE], const unsigned int m1, const unsigned int m2) {
#pragma HLS ARRAY_PARTITION variable = v_out cyclic factor = MEM_WIDTH dim = 1

        for (unsigned int i = 0; i < M_SIZE; ++i) {
#pragma HLS UNROLL factor = MEM_WIDTH
            v_out[V2S(i, m1, m2)] = v_in[i];
        }
    }

    /// @brief Utility function to copy a vector
    void CopyVector(const DT v_in[M_SIZE], DT v_out[M_SIZE]) {
        for (unsigned int i = 0; i < M_SIZE; ++i) {
            v_out[i] = v_in[i];
        }
    }

    /// @brief Utility function to add two vectors
    void vectorAdd(const DT v_in0[M_SIZE], const DT v_in1[M_SIZE], DT v_out[M_SIZE]) {
#pragma HLS ARRAY_PARTITION variable = v_out cyclic factor = MEM_WIDTH dim = 1

        for (int i = 0; i < M_SIZE; ++i) {
#pragma HLS UNROLL factor = MEM_WIDTH
            v_out[i] = v_in0[i] + v_in1[i];
        }
    }

    /// @brief Utility function to subtract two vectors
    void vectorSub(const DT v_in0[M_SIZE], const DT v_in1[M_SIZE], DT v_out[M_SIZE]) {
#pragma HLS ARRAY_PARTITION variable = v_out cyclic factor = MEM_WIDTH dim = 1

        for (int i = 0; i < M_SIZE; ++i) {
#pragma HLS UNROLL factor = MEM_WIDTH
            v_out[i] = v_in0[i] - v_in1[i];
        }
    }

    /// @brief Wrapper to PCR tridiagonal solver
    /// @details Solves tridiagonal linear system
    /// @param[in]  X1    M x 3 array holding lower/main/upper diagonals of X1
    /// tridiagonal matrix
    /// @param[in]  rhs   Right hand side of linear system to be solved
    /// @param[in]  m1    Size of array [0..m1-1] in S direction
    /// @param[in]  m2    Size of array [0..m2-1] in V direction
    /// @param[out] v_out  Solution vector representing array of size m1 x m2
    /// flattened in V inner form
    void triDiagSovlerPCR(
        DT X1[M_SIZE][DIM2_SIZE1], DT rhs[M_SIZE], DT v_out[M_SIZE], const unsigned int m1, const unsigned int m2) {
        DT r[M_SIZE];
        DT a[M_SIZE];
        DT b[M_SIZE];
        DT c[M_SIZE];
#pragma HLS RESOURCE variable = r core = RAM_2P_BRAM
#pragma HLS RESOURCE variable = a core = RAM_2P_BRAM
#pragma HLS RESOURCE variable = b core = RAM_2P_BRAM
#pragma HLS RESOURCE variable = c core = RAM_2P_BRAM
        DO_PRAGMA(HLS array_partition variable = r cyclic factor = FD_NUM_PCR)
        DO_PRAGMA(HLS array_partition variable = a cyclic factor = FD_NUM_PCR)
        DO_PRAGMA(HLS array_partition variable = b cyclic factor = FD_NUM_PCR)
        DO_PRAGMA(HLS array_partition variable = c cyclic factor = FD_NUM_PCR)

        // Working copies of diagonals as solver will overwrite them
        for (int i = 0; i < M_SIZE; i++) {
            //#pragma HLS pipeline
            r[i] = rhs[i];
            a[i] = X1[i][0];
            b[i] = X1[i][1];
            c[i] = X1[i][2];
        }

        // Call the PCR tridiagonal solver
        xf::fintech::trsvCore<DT, M_SIZE, LOG2_M_SIZE, FD_NUM_PCR>(a, b, c, r);

        // Final division for result
        for (int k = 0; k < M_SIZE; k++) {
            // Note that v_out is written in V inner order using the S2V function
            v_out[S2V(k, m1, m2)] = r[k] / b[k];
        }
    }

    /// @brief Solve pentadiagonal form linear system
    /// @details This is a highly serial algorithm and is the bottleneck in this
    /// solver
    /// Unfortunately due to the formulation used by In 'T Hout & Foulon, the
    /// pentadiagonal array contains diagonals
    /// which are not fully populated.  This causes the common parallel
    /// pentadiagonal systems to fail due to
    /// divide-by-zero errors or similar.
    /// @param[in]  A          M x 5 array holding lower/lower/main/upper/upper
    /// diagonals of X1 tridiagonal matrix
    /// @param[in]  rhs        Right hand side of linear system to be solved
    /// @param[in]  precompute Flag to indicate the scaling factors should be
    /// computed
    /// @param[out] v_out       Solution vector representing array of size m1 x m2
    /// flattened in V inner form
    void pendaDiagSovler(DT A[M_SIZE][DIM2_SIZE2], DT rhs[M_SIZE], DT v_out[M_SIZE], bool precompute) {
        // These are preprocessed versions of the diagonals.
        // For a given array these are constant so they are calculated once and
        // cached (by marking as static) and using precompute variable.
        static DT d[M_SIZE];
        static DT a[M_SIZE];
        static DT e[M_SIZE];
        static DT c[M_SIZE];
        static DT f[M_SIZE];

        static DT xmult0[M_SIZE];
        static DT xmult1[M_SIZE];

        static DT xmult_r;

        static DT d_inv[M_SIZE];

        // Not static as this changes each time
        DT r[M_SIZE];

        // Working variable
        DT xmult = 0.0;

        // One time precomputation of the vectors
        if (precompute) {
            // Main diagonal
            for (unsigned int i = 0; i < M_SIZE; ++i) {
                d[i] = A[i][2];
            }

            // First upper/lower, dropping padded zeros
            for (unsigned int i = 0; i < M_SIZE - 1; ++i) {
                a[i] = A[i + 1][1];
                c[i] = A[i][3];
            }

            // Second upper/lower, dropping padded zeros
            for (unsigned int i = 0; i < M_SIZE - 2; ++i) {
                e[i] = A[i + 2][0];
                f[i] = A[i][4];
            }

            // Scaling factors
            for (unsigned int i = 1; i < (M_SIZE - 1); ++i) {
                xmult = a[i - 1] / d[i - 1];
                d[i] = d[i] - xmult * c[i - 1];
                c[i] = c[i] - xmult * f[i - 1];
                xmult0[i] = xmult;
                xmult = e[i - 1] / d[i - 1];
                a[i] = a[i] - xmult * c[i - 1];
                d[i + 1] = d[i + 1] - xmult * f[i - 1];
                xmult1[i] = xmult;
            }

            // d manipulation
            xmult_r = a[M_SIZE - 2] / d[M_SIZE - 2];
            d[M_SIZE - 1] = d[M_SIZE - 1] - xmult_r * c[M_SIZE - 2];

            // Invert d (one time operation and allows multiplies to be used later on)
            for (unsigned int i = 0; i < M_SIZE; ++i) {
                d_inv[i] = 1.0 / d[i];
            }
        }

        // Input vector changes each time
        for (unsigned int i = 0; i < M_SIZE; ++i) {
            r[i] = rhs[i];
        }

        // Apply the scaling factors
        for (unsigned int i = 1; i < (M_SIZE - 1); ++i) {
            r[i] = r[i] - xmult0[i] * r[i - 1];
            r[i + 1] = r[i + 1] - xmult1[i] * r[i - 1];
        }

        // Back solve
        v_out[M_SIZE - 1] = (r[M_SIZE - 1] - xmult_r * r[M_SIZE - 2]) * d_inv[M_SIZE - 1];
        v_out[M_SIZE - 2] = (r[M_SIZE - 2] - c[M_SIZE - 2] * v_out[M_SIZE - 1]) * d_inv[M_SIZE - 2];
        unsigned int i = M_SIZE - 3;
        do {
            v_out[i] = (r[i] - f[i] * v_out[i + 2] - c[i] * v_out[i + 1]) * d_inv[i];
        } while (i-- > 0);
    }
};

/// @brief Utility class to encapsulate the multiplier elements
template <typename DT,
          unsigned int MEM_WIDTH,
          unsigned int INDEX_WIDTH,
          unsigned int A_SIZE,
          unsigned int M_SIZE,
          unsigned int LOG2_M_SIZE,
          unsigned int DIM2_SIZE1 = 3,
          unsigned int DIM2_SIZE2 = 5>
class StreamWrapper {
   private:
    static const unsigned int M_SIZE_BLOCKS = M_SIZE / MEM_WIDTH;

   public:
    typedef xf::fintech::blas::WideType<DT, MEM_WIDTH> WideDataType;
    typedef hls::stream<WideDataType> WideStreamType;

   public:
    /// @brief default constructor
    StreamWrapper() {
#pragma HLS inline
    }
    /// @brief Computes multiplication of tridiagonal matrix by a vector
    /// @param[in]  A1         M x 3 array holding lower/main/upper diagonals of
    /// A1 tridiagonal matrix
    /// @param[in]  u          Vector to be multiplied
    /// @param[out] rhs1_tmp0 Multiplication result
    /// @param[out] u_out     Stream form of U to pass to pentadiagonal multiplier
    void streamDimv3(DT A1[M_SIZE][DIM2_SIZE1], DT u[M_SIZE], DT rhs1_tmp0[M_SIZE], WideStreamType& u_out) {
        for (unsigned int i = 0; i < M_SIZE_BLOCKS; ++i) {
#pragma HLS PIPELINE
            WideDataType val;
#pragma HLS ARRAY_PARTITION variable = val complete
            for (unsigned int j = 0; j < MEM_WIDTH; ++j) {
                val[j] = u[i * MEM_WIDTH + j];
            }
            u_out.write(val);
        }
        xf::fintech::blas::dimv<DT, M_SIZE, DIM2_SIZE1, MEM_WIDTH>(A1, u, M_SIZE, rhs1_tmp0);
    }

    /// @brief Computes multiplication of pentadiagonal matrix by a vector
    /// @param[in]  A2         M x 5 array holding lower/main/upper diagonals of
    /// A2 tridiagonal matrix
    /// @param[in]  u_in      Vector to be multiplied [stream format]
    /// @param[in]  m1         Size of array [0..m1-1] in S direction
    /// @param[in]  m2         Size of array [0..m2-1] in V direction
    /// @param[out] rhs2_tmp0 Multiplication result
    /// @param[out] u_out     Vector output [stream format]
    void streamDimv5(DT A2[M_SIZE][DIM2_SIZE2],
                     WideStreamType& u_in,
                     unsigned int m1,
                     unsigned int m2,
                     DT rhs2_tmp0[M_SIZE],
                     WideStreamType& u_out) {
        DT u_r0[M_SIZE]; // S-inner
        DT u_r1[M_SIZE]; // V-inner
#pragma HLS ARRAY_PARTITION variable = u_r0 cyclic factor = MEM_WIDTH dim = 1
#pragma HLS ARRAY_PARTITION variable = u_r1 cyclic factor = MEM_WIDTH dim = 1

        for (unsigned int i = 0; i < M_SIZE_BLOCKS; ++i) {
#pragma HLS PIPELINE
            WideDataType val;
#pragma HLS ARRAY_PARTITION variable = val complete
            val = u_in.read();
            u_out.write(val);
            for (unsigned int j = 0; j < MEM_WIDTH; ++j) {
                u_r0[i * MEM_WIDTH + j] = val[j];
            }
        }
        // S2V conversion
        Solver<DT, M_SIZE, LOG2_M_SIZE, A_SIZE, MEM_WIDTH> solver;
        solver.reorderS2V(u_r0, u_r1, m1, m2);
        xf::fintech::blas::dimv<DT, M_SIZE, DIM2_SIZE2, MEM_WIDTH>(A2, u_r1, M_SIZE, rhs2_tmp0);
    }

    /// @brief Computes multiplication of sparse matrix by vector plus a constant
    /// @param[in]  A          Sparse matrix value
    /// @param[in]  Ar         Sparse matrix row
    /// @param[in]  Ac         Sparse matrix column
    /// @param[in]  u_in      Vector to be multiplied [stream format]
    /// @param[in]  b          Vector to be added after sparse-mult stage
    /// @param[in]  Annz       Number of non-zeros in sparse matrix (how many
    /// elements of A/Ar/Ac are valid)
    /// @param[in]  M          Matrix M-size === (m1+1) x (m2+1)
    /// @param[out] y0         Result of mult-add in flattened S-inner form
    void streamSparseMultAdd(DT A[A_SIZE],
                             unsigned int Ar[A_SIZE],
                             unsigned int Ac[A_SIZE],
                             WideStreamType& u_in,
                             DT b[M_SIZE],
                             unsigned int Annz,
                             unsigned int M,
                             DT y0[M_SIZE]) {
#pragma HLS ARRAY_PARTITION variable = y0 cyclic factor = MEM_WIDTH dim = 1
        DT u[M_SIZE];
#pragma HLS ARRAY_PARTITION variable = u cyclic factor = MEM_WIDTH dim = 1
        DT y0_tmp1[M_SIZE];
#pragma HLS ARRAY_PARTITION variable = y0_tmp1 cyclic factor = MEM_WIDTH dim = 1

        for (unsigned int i = 0; i < M_SIZE_BLOCKS; ++i) {
#pragma HLS PIPELINE
            WideDataType val = u_in.read();
#pragma HLS ARRAY_PARTITION variable = val complete
            for (unsigned int j = 0; j < MEM_WIDTH; ++j) {
                u[i * MEM_WIDTH + j] = val[j];
            }
        }
        xf::fintech::blas::Spmv<DT, MEM_WIDTH, INDEX_WIDTH, M_SIZE, M_SIZE, A_SIZE> spmv;
        spmv.sparseMultAdd(A, Ar, Ac, u, b, y0_tmp1, Annz, M);
        unsigned int vec_blocks = M / MEM_WIDTH;
        for (unsigned int i = 0; i < vec_blocks; ++i) {
#pragma HLS PIPELINE
            for (unsigned int j = 0; j < MEM_WIDTH; ++j) {
                y0[i * MEM_WIDTH + j] = u[i * MEM_WIDTH + j] + y0_tmp1[i * MEM_WIDTH + j];
            }
        }
    }

    /// @brief Wrapper function to combine multipliers into a dataflow region and
    /// allow parallelization
    void parallelBlocks(DT A[A_SIZE],
                        unsigned int Ar[A_SIZE],
                        unsigned int Ac[A_SIZE],
                        DT u[M_SIZE],
                        DT b[M_SIZE],
                        unsigned int Annz,
                        unsigned int M,
                        DT A1[M_SIZE][DIM2_SIZE1],
                        DT A2[M_SIZE][DIM2_SIZE2],
                        unsigned int m1,
                        unsigned int m2,
                        DT y0[M_SIZE],
                        DT rhs1_tmp0[M_SIZE],
                        DT rhs2_tmp0[M_SIZE]) {
        WideStreamType u0;
        WideStreamType u1;
#pragma HLS DATAFLOW
        streamDimv3(A1, u, rhs1_tmp0, u0);
        streamDimv5(A2, u0, m1, m2, rhs2_tmp0, u1);
        streamSparseMultAdd(A, Ar, Ac, u1, b, Annz, M, y0);
    }
};

} // end of internal namespace block

/// @brief Top level callable function to perform the Douglas ADI method
/// @details This function creates the solver/stream wrapper objects and connects
/// them up
/// It also provides the extra connectivity for the non-streaming blocks
/// @param[in]  A          Sparse matrix value
/// @param[in]  Ar         Sparse matrix row
/// @param[in]  Ac         Sparse matrix column
/// @param[in]  Annz       Number of non-zeros in sparse matrix (how many
/// elements of A/Ar/Ac are valid)
/// @param[in]  A1         Tridiagonal matrix stored as three vectors
/// lower/main/upper
/// @param[in]  A2         Pentadiagonal matrix stored as five vectors
/// lower/lower/main/upper/upper
/// @param[in]  X1         Tridiagonal matrix stored as three vectors
/// lower/main/upper
/// @param[in]  X2         Pentadiagonal matrix stored as five vectors
/// lower/lower/main/upper/upper
/// @param[in]  b          Boundary condition vector
/// @param[in]  u0         Initial condition (payoff condition for a call
/// option)
/// @param[in]  M1         Size of array [0..M1] in S direction
/// @param[in]  M2         Size of array [0..M2] in V direction
/// @param[in]  N          Iteration count
/// @param[out] u          Calculated price grid
template <typename DT,
          unsigned int MEM_WIDTH,
          unsigned int INDEX_WIDTH,
          unsigned int A_SIZE,
          unsigned int M_SIZE,
          unsigned int LOG2_M_SIZE,
          unsigned int DIM2_SIZE1,
          unsigned int DIM2_SIZE2>
void FdDouglas(DT A[A_SIZE],
               unsigned int Ar[A_SIZE],
               unsigned int Ac[A_SIZE],
               unsigned int Annz,
               DT A1[M_SIZE][DIM2_SIZE1],
               DT A2[M_SIZE][DIM2_SIZE2],
               DT X1[M_SIZE][DIM2_SIZE1],
               DT X2[M_SIZE][DIM2_SIZE2],
               DT b[M_SIZE],
               DT u0[M_SIZE],
               unsigned int M1,
               unsigned int M2,
               unsigned int N,
               DT u[M_SIZE]) {
    DT y0[M_SIZE];
#pragma HLS ARRAY_PARTITION variable = y0 cyclic factor = MEM_WIDTH dim = 1
    DT rhs1[M_SIZE];
    DT y1[M_SIZE];
    DT rhs2[M_SIZE];
    DT y2[M_SIZE];

    DT rhs1_tmp0[M_SIZE];
    DT rhs2_tmp0[M_SIZE];

    internal::Solver<DT, M_SIZE, LOG2_M_SIZE, A_SIZE, MEM_WIDTH, DIM2_SIZE1, DIM2_SIZE2> solver;
    internal::StreamWrapper<DT, MEM_WIDTH, INDEX_WIDTH, A_SIZE, M_SIZE, LOG2_M_SIZE> stream_wrapper;

    solver.CopyVector(u0, u);

    // Perform the Douglas iteration
    // After N iterations the u vector will hold the output
    for (int i = 0; i < N; ++i) {
        // These operations are performed in parallel with a stream passing vector
        // data between them
        stream_wrapper.parallelBlocks(A, Ar, Ac, u, b, Annz, M_SIZE, A1, A2, M1, M2, y0, rhs1_tmp0, rhs2_tmp0);

        // These operations depending on the previous one completing and run
        // sequentially
        solver.vectorSub(y0, rhs1_tmp0, rhs1);
        solver.triDiagSovlerPCR(X1, rhs1, y1, M1, M2);
        solver.vectorSub(y1, rhs2_tmp0, rhs2);
        solver.pendaDiagSovler(X2, rhs2, y2, i == 0);

        // Vector u has to be reordered back to S-inner form
        solver.reorderV2S(y2, u, M1, M2);
    }
}
}
} // End of namespace xf::fintech

#endif