//
// Created by zx on 23-6-15.
//
#include "trajectory.h"
#include <gtsam/geometry/Pose2.h>
#include <gtsam/geometry/Pose3.h>
#include <gtsam/nonlinear/NonlinearFactorGraph.h>
#include <gtsam/nonlinear/Values.h>
#include <gtsam/inference/Symbol.h>
#include <gtsam/slam/PriorFactor.h>
#include <gtsam/slam/BetweenFactor.h>
#include <gtsam/nonlinear/GaussNewtonOptimizer.h>
#include <gtsam/nonlinear/Marginals.h>
#include <gtsam/nonlinear/ISAM2.h>
#include "PoseSpeedFactor.h"
/*
 * 根据参考点生成轨迹
 */
Trajectory update_traj(Trajectory traj,std::vector<gtsam::Vector2> speeds,int referenceId,double dt,double& loss){
  Trajectory ret=traj;
  if(referenceId>=traj.size() || referenceId<0)
  {
    printf(" reference < traj.size or referece <0 \n");
    return ret;
  }
  Pose2d refer=traj[referenceId];
  loss=0;
  for(int i=referenceId-1;i>=0;i--){
    double v=speeds[i][0];
    double w=speeds[i][1];
    double dyaw = w * dt;
    double theta=refer.theta()-dyaw;
    double dx = v * dt * cos(theta);
    double dy = v * dt * sin(theta);
    ret[i]=refer-Pose2d(dx,dy,dyaw);
    refer=ret[i];

    loss+=Pose2d::distance(ret[i],traj[i]);

  }

  refer=traj[referenceId];
  for(int i=referenceId+1;i<traj.size();++i){
    double v=speeds[i-1][0];
    double w=speeds[i-1][1];
    double dyaw = w * dt;
    double theta=refer.theta();
    double dx = v * dt * cos(theta);
    double dy = v * dt * sin(theta);
    ret[i]=refer+Pose2d(dx,dy,dyaw);

    refer=ret[i];
    loss+=Pose2d::distance(ret[i],traj[i]);

  }

  return ret;
}
bool SolveSpeed(Trajectory traj)
{

  if(traj.size()<2){
    printf("traj size must > 2\n");
    return false;
  }
  using gtsam::symbol_shorthand::V;
  //定义图
  gtsam::NonlinearFactorGraph graph;
  gtsam::ISAM2Params parameters;
  parameters.relinearizeThreshold = 0.1;
  parameters.relinearizeSkip = 1;
  gtsam::ISAM2* isam=new gtsam::ISAM2(parameters);

  //定义待求因子
  const int NUM=traj.size();
  double dt=0.1;
  //
  double sigma_x=0.005;
  double sigma_yaw=0.3*M_PI/180.0;

  gtsam::Values initials;
  for (int i=0;i<NUM-1;++i){
    //定义变量，计算初值
    Pose2d relative=Pose2d::relativePose(traj[i+1],traj[i]);
    double v=relative.x()/dt;
    double w=relative.theta()/dt;
    gtsam::Vector2 speed;
    speed<<v,w;
    initials.insert(V(i),speed);
    //初值分布
    Eigen::Matrix<double,2,1> sigma=Eigen::Matrix<double,2,1>::Zero();
    sigma(0,0)=sigma_x/dt;
    sigma(1,0)=sigma_yaw/dt;
    gtsam::noiseModel::Diagonal::shared_ptr priorModel = gtsam::noiseModel::Diagonal::Sigmas(sigma);
    graph.add(gtsam::PriorFactor<gtsam::Vector2>(V(i), speed, priorModel));
  }

  for(int i=0;i<NUM-2;++i){
    Eigen::Matrix<double,2,1> sigma_acc=Eigen::Matrix<double,2,1>::Zero();
    sigma_acc(0,0)=1*dt;
    sigma_acc(1,0)=10*M_PI/180.0*dt;
    gtsam::noiseModel::Diagonal::shared_ptr priorModel_t = gtsam::noiseModel::Diagonal::Sigmas(sigma_acc);
    //增加前后两点z值增量的概率分布，均值为0，方差为sigma_a的分布
    graph.add(gtsam::BetweenFactor<gtsam::Vector2>(V(i), V(i+1), gtsam::Vector2(0,0), priorModel_t));
  }
  isam->update(graph,initials);
  isam->update();
  gtsam::Values results=isam->calculateEstimate();

  std::vector<gtsam::Vector2> speeds;
  for(int i=0;i<NUM-1;++i){
    speeds.push_back(results.at<gtsam::Vector2>(V(i)));
  }

  double min_loss=1e16;
  int bestId=0;
  Trajectory best_traj=update_traj(traj,speeds,bestId,dt,min_loss);
  for(int i=1;i<traj.size();++i){
    double loss=0;
    Trajectory newTraj= update_traj(traj,speeds,i,dt,loss);
    if(loss<min_loss){
      min_loss=loss;
      best_traj=newTraj;
      bestId=i;
    }
  }

  printf(" best traj referId:%d P:%.5f %.5f %.5f , loss:%f\n",bestId,
         traj[bestId].x(),traj[bestId].y(),traj[bestId].theta(),min_loss);

  Pose2d diff=best_traj[0] - traj[0];
  printf("00 P:%.4f %.4f %.4f, OPT:%.4f %.4f %.4f, Diff:%.4f %.4f %.4f, Speed:%.4f %.4f\n",
         traj[0].x(),traj[0].y(),traj[0].theta(),best_traj[0].x(),best_traj[0].y(),best_traj[0].theta(),
         diff.x(),diff.y(),diff.theta(),results.at<gtsam::Vector2>(V(0))[0],results.at<gtsam::Vector2>(V(0))[1]*180/M_PI);

  for(int i=1;i<NUM;++i){
    diff=best_traj[i] - traj[i];
    if(i<NUM-1) {
      gtsam::Vector2 Speed=results.at<gtsam::Vector2>(V(i));
      printf("%02d P:%.4f %.4f %.4f, OPT:%.4f %.4f %.4f, Diff:%.4f %.4f %.4f, Speed:%.4f %.4f\n",i,
             traj[i].x(),traj[i].y(),traj[i].theta(),best_traj[i].x(),best_traj[i].y(),best_traj[i].theta(),
             diff.x(),diff.y(),diff.theta(),Speed[0],Speed[1]*180/M_PI);

    }else{
      printf("%02d P:%.4f %.4f %.4f, OPT:%.4f %.4f %.4f, Diff:%.4f %.4f %.4f\n",i,
             traj[i].x(),traj[i].y(),traj[i].theta(),best_traj[i].x(),best_traj[i].y(),best_traj[i].theta(),
             diff.x(),diff.y(),diff.theta());
    }
  }
  return true;

}

bool SolvePoses(Trajectory traj){
  if(traj.size()<2){
    printf("traj size must > 2\n");
    return false;
  }
  using gtsam::symbol_shorthand::P;
  //定义图
  gtsam::NonlinearFactorGraph graph;
  gtsam::ISAM2Params parameters;
  parameters.relinearizeThreshold = 0.1;
  parameters.relinearizeSkip = 1;
  gtsam::ISAM2* isam=new gtsam::ISAM2(parameters);

  //定义待求因子
  const int NUM=traj.size();
  double dt=0.1;
  //
  double sigma_x=0.005;
  double sigma_y=sigma_x;
  double sigma_yaw=0.3*M_PI/180.0;

  gtsam::Values initials;
  for (int i=0;i<NUM;++i){
    //定义变量，计算初值
    gtsam::Pose2 pose(traj[i].x(),traj[i].y(),traj[i].theta());
    initials.insert(P(i),pose);
    //初值分布
    Eigen::Matrix<double,3,1> sigma=Eigen::Matrix<double,3,1>::Zero();
    sigma(0,0)=sigma_x;
    sigma(1,0)=sigma_y;
    sigma(2,0)=sigma_yaw;
    gtsam::noiseModel::Diagonal::shared_ptr priorModel = gtsam::noiseModel::Diagonal::Sigmas(sigma);
    graph.add(gtsam::PriorFactor<gtsam::Pose2>(P(i), pose, priorModel));
    //增加速度约束及分布
    if(i>=1){
      Pose2d relative=Pose2d::relativePose(traj[i],traj[i-1]);
      double v=sqrt(relative.x()*relative.x()+relative.y()*relative.y())/dt;
      if(relative.x()<0) v=-v;
      double w=relative.theta()/dt;

      gtsam::Pose2 distance(v*dt*cos(traj[i-1].theta()),v*dt*sin(traj[i-1].theta()),w*dt);
      Eigen::Matrix<double,3,1> sigma_distance=Eigen::Matrix<double,3,1>::Zero();
      sigma_distance(0,0)=sqrt(2)*sigma_x;
      sigma_distance(1,0)=sqrt(2)*sigma_y;
      sigma_distance(2,0)=sqrt(2)*sigma_yaw;
      gtsam::noiseModel::Diagonal::shared_ptr priorModel_t = gtsam::noiseModel::Diagonal::Sigmas(sigma_distance);
      graph.add(gtsam::BetweenFactor<gtsam::Pose2>(P(i-1), P(i), distance, priorModel_t));
    }

  }


  isam->update(graph,initials);
  isam->update();
  gtsam::Values results=isam->calculateEstimate();

  for(int i=0;i<NUM;++i){

    gtsam::Pose2 optpose=results.at<gtsam::Pose2>(P(i));
    Pose2d last(optpose.x(),optpose.y(),optpose.theta());
    Pose2d diff=last-traj[i];

    if(i<NUM-1) {
      gtsam::Pose2 pose1=results.at<gtsam::Pose2>(P(i+1));
      Pose2d next(pose1.x(),pose1.y(),pose1.theta());

      Pose2d relative=Pose2d::relativePose(next,last);
      double v=sqrt(relative.x()*relative.x()+relative.y()*relative.y())/dt;
      if(relative.x()<0) v=-v;
      double w=relative.theta()/dt;

      printf("%02d P:%.4f %.4f %.4f,  OPT:%.4f %.4f %.4f,  Diff:%.4f %.4f %.4f,  Speed:%.4f %.4f\n",i,
             traj[i].x(),traj[i].y(),traj[i].theta(),optpose.x(),optpose.y(),optpose.theta(),
             diff.x(),diff.y(),diff.theta(),v,w*180.0/M_PI);

    }else{
      printf("%02d P:%.4f %.4f %.4f,  OPT:%.4f %.4f %.4f,  Diff:%.4f %.4f %.4f\n",i,
             traj[i].x(),traj[i].y(),traj[i].theta(),optpose.x(),optpose.y(),optpose.theta(),
             diff.x(),diff.y(),diff.theta());
    }
  }
  return true;
}

bool SolvePoseSpeed(Trajectory traj){
  if(traj.size()<2){
    printf("traj size must > 2\n");
    return false;
  }
  using gtsam::symbol_shorthand::P;
  using gtsam::symbol_shorthand::V;
  //定义图
  gtsam::NonlinearFactorGraph graph;
  gtsam::ISAM2Params parameters;
  parameters.relinearizeThreshold = 0.1;
  parameters.relinearizeSkip = 1;
  gtsam::ISAM2* isam=new gtsam::ISAM2(parameters);

  //定义待求因子
  const int NUM=traj.size();
  double dt=0.1;
  //
  double sigma_x=0.01;
  double sigma_y=sigma_x;
  double sigma_yaw=0.3*M_PI/180.0;

  gtsam::Values initials;
  for (int i=0;i<NUM;++i){
    //定义变量，计算初值
    gtsam::Pose2 pose(traj[i].x(),traj[i].y(),traj[i].theta());
    initials.insert(P(i),pose);
    //初值分布
    Eigen::Matrix<double,3,1> sigma=Eigen::Matrix<double,3,1>::Zero();
    sigma(0,0)=sigma_x;
    sigma(1,0)=sigma_y;
    sigma(2,0)=sigma_yaw;
    gtsam::noiseModel::Diagonal::shared_ptr priorModel = gtsam::noiseModel::Diagonal::Sigmas(sigma);
    graph.add(gtsam::PriorFactor<gtsam::Pose2>(P(i), pose, priorModel));

    if(i<NUM-1) {
      //增加速度分布
      Pose2d relative = Pose2d::relativePose(traj[i+1], traj[i]);
      double v = sqrt(relative.x() * relative.x() + relative.y() * relative.y()) / dt;
      if (relative.x() < 0) v = -v;
      double w = relative.theta() / dt;

      initials.insert(V(i), gtsam::Vector2(v, w));
      Eigen::Matrix<double, 2, 1> sigma_v = Eigen::Matrix<double, 2, 1>::Zero();
      sigma_v(0, 0) = 1.0;
      sigma_v(1, 0) = 10*M_PI/180.0;
      gtsam::noiseModel::Diagonal::shared_ptr priorModel_v = gtsam::noiseModel::Diagonal::Sigmas(sigma_v);
      graph.add(gtsam::PriorFactor<gtsam::Vector2>(V(i), Vector2(v,w), priorModel_v));
    }
  }

  for(int i=0;i<NUM-1;++i) {
    if(i<NUM-2) {
      Eigen::Matrix<double, 2, 1> sigma_a = Eigen::Matrix<double, 2, 1>::Zero();
      sigma_a(0, 0) = 1;
      sigma_a(1, 0) = 20 * M_PI / 180.0;
      gtsam::noiseModel::Diagonal::shared_ptr priorModel_a = gtsam::noiseModel::Diagonal::Sigmas(sigma_a);
      graph.add(gtsam::BetweenFactor<gtsam::Vector2>(V(i), V(i + 1), Vector2(0, 0), priorModel_a));
    }

    Eigen::Matrix<double, 3, 1> sigma_dt = Eigen::Matrix<double, 3, 1>::Ones()*0.005;
    gtsam::noiseModel::Diagonal::shared_ptr priorModel_dt = gtsam::noiseModel::Diagonal::Sigmas(sigma_dt);
    graph.add(PoseSpeedFactor(P(i), V(i),P(i+1), gtsam::Vector3(dt,dt,dt), priorModel_dt));
  }


  isam->update(graph,initials);
  isam->update();
  gtsam::Values results=isam->calculateEstimate();
  printf("  solve --completed--- \n");

  for(int i=0;i<NUM;++i){

    gtsam::Pose2 optpose=results.at<gtsam::Pose2>(P(i));
    Pose2d last(optpose.x(),optpose.y(),optpose.theta());
    Pose2d diff=last-traj[i];

    if(i<NUM-1) {
      gtsam::Vector2 speed=results.at<gtsam::Vector2>(V(i));
      printf("%02d P:%.4f %.4f %.4f,  OPT:%.4f %.4f %.4f,  Diff:%.4f %.4f %.4f,  Speed:%.4f %.4f\n",i,
             traj[i].x(),traj[i].y(),traj[i].theta(),optpose.x(),optpose.y(),optpose.theta(),
             diff.x(),diff.y(),diff.theta(),speed[0],speed[1]*180.0/M_PI);

    }else{
      printf("%02d P:%.4f %.4f %.4f,  OPT:%.4f %.4f %.4f,  Diff:%.4f %.4f %.4f\n",i,
             traj[i].x(),traj[i].y(),traj[i].theta(),optpose.x(),optpose.y(),optpose.theta(),
             diff.x(),diff.y(),diff.theta());
    }
  }
  return true;
}

int main(){
  Trajectory traj;
  traj.push_point(Pose2d(0.005,0,1.565));
  traj.push_point(Pose2d(0.001,0.01,1.581));
  traj.push_point(Pose2d(0.002,0.03,1.572));
  traj.push_point(Pose2d(0.010,0.09,1.573));
  traj.push_point(Pose2d(0.009,0.2,1.571));

  traj.push_point(Pose2d(0.0,0.305,1.571));
  traj.push_point(Pose2d(0.02,0.303,1.570));
  traj.push_point(Pose2d(0.019,0.303,1.571));


 /* traj.push_point(Pose2d(0,0.4,1.57));
  traj.push_point(Pose2d(0,0.45,1.57));
  traj.push_point(Pose2d(0,0.47,1.565));
  traj.push_point(Pose2d(0,0.5,1.571));

  traj.push_point(Pose2d(0,0.51,1.571));
  traj.push_point(Pose2d(0,0.508,1.572));
  traj.push_point(Pose2d(0,0.495,1.574));
  traj.push_point(Pose2d(0,0.499,1.565));
  traj.push_point(Pose2d(0,0.501,1.568));
  traj.push_point(Pose2d(0,0.50,1.57));

  traj.push_point(Pose2d(0.005,0.501,1.57));
  traj.push_point(Pose2d(0.004,0.499,1.57));
  traj.push_point(Pose2d(-0.002,0.498,1.57));
  traj.push_point(Pose2d(-0.002,0.502,1.57));
  traj.push_point(Pose2d(-0.001,0.501,1.57));
  traj.push_point(Pose2d(0,0.499,1.57));*/

   //SolvePoses(traj);
   printf("-------------------------------------------------\n");
  auto tp=std::chrono::steady_clock::now();

  //SolveSpeed(traj);
  SolvePoseSpeed(traj);
  auto now=std::chrono::steady_clock::now();
  auto duration = std::chrono::duration_cast<std::chrono::microseconds>(now - tp);
  double time = double(duration.count()) * std::chrono::microseconds::period::num /
      std::chrono::microseconds::period::den;
  printf("time :%.5f s\n",time);
  return 0;
}