#include "mpc_particle_model.h"

using CppAD::AD;

// 设置静态变量值
double Mpc_particle_model::ref_v = 2.0;
double Mpc_particle_model::max_vel = 3.0;
double Mpc_particle_model::dt = 0.1;

class FG_eval_model
{
public:
  // Fitted polynomial coefficients
  Eigen::VectorXd coeffs;
  double m_ref_v, m_dt;
  double m_v, m_vth;
  double m_global_target_x, m_global_target_y;
  FG_eval_model(Eigen::VectorXd coeffs,double v,double vth,double ref_v,double dt, 
                double global_target_x, double global_target_y)
  {
    this->coeffs = coeffs;
    m_v = v;
    m_vth = vth;
    m_ref_v = ref_v;
    m_dt = dt;
    m_global_target_x = global_target_x;
    m_global_target_y = global_target_y;
  }

  typedef CPPAD_TESTVECTOR(AD<double>) ADvector;
  void operator()(ADvector &fg, const ADvector &vars)
  {
    fg[0] = 0;

    // Weights for how "important" each cost is - can be tuned
    const int y_cost_weight = 5000;
    const int th_cost_weight = 1000;
    const int v_cost_weight = 20;
    const int vth_cost_weight = 100;
    const int a_cost_weight = 1;
    const int ath_cost_weight = 5;

    const int target_cost_weight = 20;

    // // std::cout << "-------------"<< std::endl;
    // // 增加对接近目标时的速度限制, 接近1m内才生效
    // AD<double> x0 = vars[Mpc_particle_model::x_start];
    // AD<double> y0 = vars[Mpc_particle_model::y_start];
    // AD<double> th0 = vars[Mpc_particle_model::theta_start + Mpc_particle_model::N - 1];
    // AD<double> v0 = vars[Mpc_particle_model::v_start + Mpc_particle_model::N - 1];
    // AD<double> the_last_predict_x = x0 + v0 * CppAD::cos(th0) * m_dt * (Mpc_particle_model::N - 1);
    // AD<double> the_last_predict_y = y0 + v0 * CppAD::sin(th0) * m_dt * (Mpc_particle_model::N - 1);
    // AD<double> dist = CppAD::sqrt(CppAD::pow(the_last_predict_x - m_global_target_x, 2) + CppAD::pow(the_last_predict_y - m_global_target_y, 2));
    // AD<double> inv_sigmoid = 1 + CppAD::exp(10 * dist - 10);
    // AD<double> v_target = m_ref_v * (1 - 1 / inv_sigmoid);
    // std::cout << "dist, sigmoid: " << dist << ", " << (1 - 1 / inv_sigmoid) * m_ref_v << std::endl;
    // // fg[0] += target_cost_weight * dist / inv_sigmoid ;

    // Cost for CTE, psi error and velocity
    for (int t = 0; t < Mpc_particle_model::N; t++)
    {
      AD<double> xt = vars[ Mpc_particle_model::x_start + t];
      AD<double> fx = coeffs[0] + coeffs[1] * xt + coeffs[2] * CppAD::pow(xt, 2) + coeffs[3] * CppAD::pow(xt, 3);
      fg[0] += y_cost_weight * CppAD::pow(vars[Mpc_particle_model::y_start + t] - fx, 2);

      AD<double> fth = CppAD::atan(coeffs[1] + 2 * coeffs[2] * xt + 3 * coeffs[3] * CppAD::pow(xt, 2));
      fg[0] += th_cost_weight * CppAD::pow(vars[Mpc_particle_model::theta_start + t] - fth, 2);
      fg[0] += v_cost_weight * CppAD::pow(vars[Mpc_particle_model::v_start + t] - m_ref_v, 2);
      fg[0] += vth_cost_weight * CppAD::pow(vars[Mpc_particle_model::vth_start + t], 2);
      // std::cout<<"x, y, theta, v_ref: "<< xt << ", "<<fx<<", " << fth << ", " << m_ref_v <<std::endl;
    }

    // Costs for steering (delta) and acceleration (a)
    for (int t = 0; t < Mpc_particle_model::N - 1; t++)
    {
      fg[0] += a_cost_weight * CppAD::pow(vars[Mpc_particle_model::acc_start + t], 2);
      fg[0] += ath_cost_weight * CppAD::pow(vars[Mpc_particle_model::acc_vth_start + t], 2);
    }

    /////////////////////
    fg[1 + Mpc_particle_model::x_start] = vars[Mpc_particle_model::x_start];
    fg[1 + Mpc_particle_model::y_start] = vars[Mpc_particle_model::y_start];
    fg[1 + Mpc_particle_model::theta_start] = vars[Mpc_particle_model::theta_start];
    fg[1 + Mpc_particle_model::v_start] = vars[Mpc_particle_model::v_start];
    fg[1 + Mpc_particle_model::vth_start] = vars[Mpc_particle_model::vth_start];

    // The rest of the constraints
    for (int t = 1; t < Mpc_particle_model::N; t++)
    {
      // State at time t + 1
      AD<double> x1 = vars[Mpc_particle_model::x_start + t];
      AD<double> y1 = vars[Mpc_particle_model::y_start + t];
      AD<double> th1 = vars[Mpc_particle_model::theta_start + t];
      AD<double> v1 = vars[Mpc_particle_model::v_start + t];
      AD<double> vth1 = vars[Mpc_particle_model::vth_start + t];

      // State at time t
      AD<double> x0 = vars[Mpc_particle_model::x_start + t - 1];
      AD<double> y0 = vars[Mpc_particle_model::y_start + t - 1];
      AD<double> th0 = vars[Mpc_particle_model::theta_start + t - 1];
      AD<double> v0 = vars[Mpc_particle_model::v_start + t - 1];
      AD<double> vth0 = vars[Mpc_particle_model::vth_start + t - 1];

      // Actuator constraints at time t only
      AD<double> acc0 = vars[Mpc_particle_model::acc_start + t - 1];
      AD<double> acc_vth0 = vars[Mpc_particle_model::acc_vth_start + t - 1];

      // Setting up the rest of the model constraints
      fg[1 + Mpc_particle_model::x_start + t] = x1 - (x0 + (v0 + 0.5 * acc0 * m_dt) * CppAD::cos(th0 + 0.5 * vth0 * m_dt) * m_dt);
      fg[1 + Mpc_particle_model::y_start + t] = y1 - (y0 + (v0 + 0.5 * acc0 * m_dt) * CppAD::sin(th0 + 0.5 * vth0 * m_dt) * m_dt);
      // fg[1 + Mpc_particle_model::x_start + t] = x1 - (x0 + v0 * CppAD::cos(th0) * m_dt);
      // fg[1 + Mpc_particle_model::y_start + t] = y1 - (y0 + v0 * CppAD::sin(th0) * m_dt);
      fg[1 + Mpc_particle_model::theta_start + t] = th1 - (th0 + vth0 * m_dt);
      fg[1 + Mpc_particle_model::v_start + t] = v1 - (v0 + acc0 * m_dt);
      fg[1 + Mpc_particle_model::vth_start + t] = vth1 - (vth0 + acc_vth0 * m_dt);
    }
  }
};

Mpc_particle_model::Mpc_particle_model()
{
  // ref_v = 2.0;
  last_velocity = 0.0;
  last_velocity_theta = 0.0;
  m_last_pose.setIdentity();
  m_global_target.setIdentity();
  mb_has_global_target = false;
}

Mpc_particle_model::~Mpc_particle_model()
{
}

void Mpc_particle_model::set_global_target(tf::Pose target){
  m_global_target = target;
  mb_has_global_target = true;
}

// 求y
double Mpc_particle_model::polyeval(Eigen::VectorXd coeffs, double x)
{
  double result = 0.0;
  for (int i = 0; i < coeffs.size(); i++)
  {
    result += coeffs[i] * pow(x, i);
  }
  return result;
}

// 求参数
Eigen::VectorXd Mpc_particle_model::polyfit(Eigen::VectorXd xvals, Eigen::VectorXd yvals, int order)
{
  assert(xvals.size() == yvals.size());
  assert(order >= 1 && order <= xvals.size() - 1);
  Eigen::MatrixXd A(xvals.size(), order + 1);

  for (int i = 0; i < xvals.size(); i++)
  {
    A(i, 0) = 1.0;
  }

  for (int j = 0; j < xvals.size(); j++)
  {
    for (int i = 0; i < order; i++)
    {
      A(j, i + 1) = A(j, i) * xvals(j);
    }
  }

  auto Q = A.householderQr();
  auto result = Q.solve(yvals);
  return result;
}

// 输入局部路径，当前位姿，速度与角速度
// 输出用于下发的速度与角速度
bool Mpc_particle_model::Solve(double ref_velocity, std::vector<tf::Pose> ref_poses, tf::Pose curr_pose, double &v, double &vth, std::vector<tf::Pose> &pred_poses)
{
  ref_v = ref_velocity;
  // 转换坐标后的局部路径与当前位姿
  Eigen::VectorXd ptsx_car(ref_poses.size());
  Eigen::VectorXd ptsy_car(ref_poses.size());
  double t_global_target_trans_x, t_global_target_trans_y;
  // 反向旋转
  for (int i = 0; i < ref_poses.size(); i++)
  {
    double t_transformed_x = ref_poses[i].getOrigin().getX() - curr_pose.getOrigin().getX();
    double t_transformed_y = ref_poses[i].getOrigin().getY() - curr_pose.getOrigin().getY();
    ptsx_car[i] = t_transformed_x * curr_pose.getBasis()[0][0] + t_transformed_y * curr_pose.getBasis()[1][0];
    ptsy_car[i] = t_transformed_x * curr_pose.getBasis()[0][1] + t_transformed_y * curr_pose.getBasis()[1][1];
    // ptsx_car[i] = ref_poses[i].getOrigin().getX();
    // ptsy_car[i] = ref_poses[i].getOrigin().getY();
    // std::cout<<"x,y: "<<ptsx_car[i]<<", "<<ptsy_car[i]<<std::endl;
  }
  double trans_x = m_global_target.getOrigin().getX() - curr_pose.getOrigin().getX();
  double trans_y = m_global_target.getOrigin().getY() - curr_pose.getOrigin().getY();
  t_global_target_trans_x = trans_x * curr_pose.getBasis()[0][0] + trans_y * curr_pose.getBasis()[1][0];
  t_global_target_trans_y = trans_x * curr_pose.getBasis()[0][1] + trans_y * curr_pose.getBasis()[1][1];

  //拟合求斜率与y方向误差
  if(ptsx_car.size() < 3 || ptsy_car.size() < 3){
    std::cout<<"点数过少，无法拟合"<<std::endl;
    return false;
  }
  auto coeffs = polyfit(ptsx_car, ptsy_car, 3);
  // std::cout<<"coeffs: "<<coeffs<<std::endl;
  // if(std::isnan(coeffs[0])||std::isnan(coeffs[1])||std::isnan(coeffs[2])||std::isnan(coeffs[3]))
  //   return;

  bool ok = true;
  typedef CPPAD_TESTVECTOR(double) Dvector;

  size_t n_vars = N * 5 + (N - 1) * 2;

  size_t n_constraints = N * 5;

  Dvector vars(n_vars);
  for (int i = 0; i < n_vars; i++)
  {
    vars[i] = 0.0;
  }

  // 参数约束
  Dvector vars_lowerbound(n_vars);
  Dvector vars_upperbound(n_vars);
  // x, y constrains
  for (int i = 0; i < theta_start; i++)
  {
    vars_lowerbound[i] = -1.0e2;
    vars_upperbound[i] = 1.0e2;
  }
  // theta constrains
  for (int i = theta_start; i < v_start; i++)
  {
    vars_lowerbound[i] = -M_PI;
    vars_upperbound[i] = M_PI;
  }
  // velocity constrains
  for (int i = v_start; i < vth_start; i++)
  {
    vars_lowerbound[i] = -max_vel;
    vars_upperbound[i] = max_vel;
  }
  // velocity theta constrains
  for (int i = vth_start; i < acc_start; i++)
  {
    vars_lowerbound[i] = -M_PI / 16.0;
    vars_upperbound[i] = M_PI / 16.0;
  }
  // acceleration of velocity constrains
  for (int i = acc_start; i < acc_vth_start; i++)
  {
    vars_lowerbound[i] = -2.0;
    vars_upperbound[i] = 2.0;
  }
  // acceleration of velocity theta
  for (int i = acc_vth_start; i < n_vars; i++)
  {
    vars_lowerbound[i] = -M_PI / 16.0;
    vars_upperbound[i] = M_PI / 16.0;
  }

  // 函数约束
  Dvector constraints_lowerbound(n_constraints);
  Dvector constraints_upperbound(n_constraints);
  for (int i = 0; i < n_constraints; i++)
  {
    constraints_lowerbound[i] = 0;
    constraints_upperbound[i] = 0;
  }
  // v
  for (int i = v_start; i < v_start + 1; i++)
  {
    constraints_lowerbound[i] = v-0.1;
    constraints_upperbound[i] = v+0.1;
  }
  // vth
  for (int i = vth_start; i < vth_start + 1; i++)
  {
    constraints_lowerbound[i] = vth;
    constraints_upperbound[i] = vth;
  }

  double theta_diff = curr_pose.getRotation().getAngle() - m_last_pose.getRotation().getAngle();
  // Predict state after latency
  double pred_v = (curr_pose.getOrigin() - m_last_pose.getOrigin()).length() / dt;
  double pred_theta = theta_diff / dt;

  // object that computes objective and constraints
  FG_eval_model fg_eval(coeffs, pred_v, pred_theta, ref_v, dt, t_global_target_trans_x, t_global_target_trans_y);

  std::string options;
  // Uncomment this if you'd like more print information
  options += "Integer print_level  0\n";
  // NOTE: Setting sparse to true allows the solver to take advantage
  // of sparse routines, this makes the computation MUCH FASTER. If you
  // can uncomment 1 of these and see if it makes a difference or not but
  // if you uncomment both the computation time should go up in orders of
  // magnitude.
  // options += "String  sb           yes\n";
  // maximum number of iterations
  // options += "Integer max_iter     80\n";
  // approximate accuracy in first order necessary conditions;
  // see Mathematical Programming, Volume 106, Number 1,
  // Pages 25-57, Equation (6)
  // options += "Numeric tol          1e-4\n";
  options += "Sparse  true        forward\n";
  options += "Sparse  true        reverse\n";
  // NOTE: Currently the solver has a maximum time limit of 0.5 seconds.
  // Change this as you see fit.
  options += "Numeric max_cpu_time          0.1\n";

  // place to return solution
  CppAD::ipopt::solve_result<Dvector> solution;
  // solve the problem
  // std::cout<<vars.size()<<", "<<vars_lowerbound.size()<<", "<<vars_upperbound.size()<<", "<<constraints_lowerbound.size()<<", "<<constraints_upperbound.size()<<", "<<std::endl;
  CppAD::ipopt::solve<Dvector, FG_eval_model>(
      options, vars, vars_lowerbound, vars_upperbound, constraints_lowerbound,
      constraints_upperbound, fg_eval, solution);

  // Check some of the solution values
  ok &= solution.status == CppAD::ipopt::solve_result<Dvector>::success;

  // Cost
  auto cost = solution.obj_value;
  // std::cout << "model   status " << ok << ", Cost " << cost << std::endl;

  // 根据优化结果获取需下发 v vth
  std::vector<double> solved;
  if (ok)
  {
    // // 先存速度、角速度，后存预测路径
    // solved.push_back(solution.x[acc_start]);
    // solved.push_back(solution.x[acc_omega_start]);
    // for (int i = 0; i < N; ++i)
    // {
    //   solved.push_back(solution.x[x_start + i]);
    //   solved.push_back(solution.x[y_start + i]);
    //   std::cout<<"path: "<<solution.x[x_start + i]<<", "<<solution.x[y_start + i]<<std::endl;
    // }
    // // std::cout<<"return solved"<<std::endl;

    //更新 v vth
    v = solution.x[v_start];// + fabs(solution.x[acc_start]) * dt;
    vth = solution.x[vth_start];// + fabs(solution.x[acc_omega_start]) * dt;
    pred_poses.clear();
    for (size_t i = 0; i < N; i++)
    {
      double real_x = solution.x[x_start + i]*curr_pose.getBasis()[0][0] + solution.x[y_start + i]*curr_pose.getBasis()[0][1];
      double real_y = solution.x[x_start + i]*curr_pose.getBasis()[1][0] + solution.x[y_start + i]*curr_pose.getBasis()[1][1];
      // pred_poses.push_back(tf::Pose(tf::Quaternion(), tf::Vector3(solution.x[x_start + i], solution.x[y_start + i], 0.0)));      
      pred_poses.push_back(tf::Pose(
        tf::Quaternion(0, 0, sin(solution.x[theta_start]/2.0), cos(solution.x[theta_start]/2.0)),
        tf::Vector3(real_x, real_y, 0.0)
      ));
      // std::cout<<"path: "<<solution.x[x_start + i]<<", "<<solution.x[y_start + i]<<std::endl;
    }
    
    std::cout << std::setprecision(3) << "model   v, vth, acc, acc_vth: " << v << ", " << vth << ", " << solution.x[acc_start] << ", " << solution.x[acc_vth_start]<< std::endl;
    // std::cout << " ref_v : " << ref_v << std::endl;
    m_last_pose = curr_pose;
    last_velocity = v;
    last_velocity_theta = vth;
    return true;
  }
  else
  {
    m_last_pose.setIdentity();
    return false;
  }
}
