//
// Created by zx on 2019/9/3.
//

#include "TowBotMPC.h"
#include <cppad/cppad.hpp>
#include <cppad/ipopt/solve.hpp>
#include "Eigen/Core"
#include "Eigen/QR"

using CppAD::AD;

size_t K = 10;
size_t nx1 = 0;
size_t ny1 = nx1 + K;
size_t nth1 = ny1 + K;
size_t nv1 = nth1 + K;
size_t nvth1 = nv1 + K;
size_t nacc1 = nvth1 + K;
size_t naccth1 = nacc1 + K - 1;

size_t nx2 = naccth1 + K - 1;
size_t ny2 = nx2 + K;
size_t nth2 = ny2 + K;
size_t nv2 = nth2 + K;
size_t nvth2 = nv2 + K;
size_t nacc2 = nvth2 + K;
size_t naccth2 = nacc2 + K - 1;

size_t relative_pose_contrains = naccth2 + K - 1;

double Lx1 = -1.3;
double Ly1 = 0;
double Lx2 = 1.3;
double Ly2 = 0;

double ref_dist = 1.0;
// double min_dist = 1.5;
// double max_dist = 3.0;

class FG_eval_2
{
public:
    // Fitted polynomial coefficients
    Eigen::VectorXd coeffs, coeffs_b;
    Eigen::VectorXd m_data;

    FG_eval_2(Eigen::VectorXd coeffs, Eigen::VectorXd coeffs_b, Eigen::VectorXd P)
    {
        this->coeffs = coeffs;
        this->coeffs_b = coeffs_b;
        m_data = P;
    }

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

        // std::cout << "max size: " << naccth2 + K - 1 << std::endl;
        double v1 = m_data[0];
        double v2 = m_data[1];
        double vth1 = m_data[2];
        double vth2 = m_data[3];
        double ref_v = m_data[4];
        double dt = m_data[5];
        
        /// 两车初始位姿差   前车相对后车坐标系
        double dx = m_data[6]; 
        double dy = m_data[7];
        double dth = m_data[8]; ///两坐标系角度

        double target_dx = m_data[9];
        double target_dy = m_data[10];
        double slowdown_dist = m_data[11];

        const int brake_weight = 200;
        ///前车 cost weight
        const int y_cost_weight = 15000;
        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 = 10;

        fg[0] = 0;
        ///后车 cost weight
        const int back_y_cost_weight = 15000;
        const int back_th_cost_weight = 1000;
        const int back_v_cost_weight = 20;

        const int back_vth_cost_weight = 100;
        const int back_a_cost_weight = 1;
        const int back_ath_cost_weight = 10;

        // x, y, theta 路径约束
        fg[1 + nx1] = vars[nx1];// - vars[nv1] * dt;
        fg[1 + ny1] = vars[ny1];
        fg[1 + nth1] = vars[nth1];// - vars[nvth1] * dt;
        fg[1 + nv1] = vars[nv1];
        fg[1 + nvth1] = vars[nvth1];

        fg[1 + nx2 - 2*(K-1)] = vars[nx2];// - vars[nv2] * dt;
        fg[1 + ny2 - 2*(K-1)] = vars[ny2];
        fg[1 + nth2 - 2*(K-1)] = vars[nth2];// - vars[nvth2] * dt;
        fg[1 + nv2 - 2*(K-1)] = vars[nv2];
        fg[1 + nvth2 - 2*(K-1)] = vars[nvth2];

        AD<double> accumulate_fx = vars[nx1];
        AD<double> accumulate_fy = vars[ny1];
        AD<double> accumulate_bx = vars[nx2];
        AD<double> accumulate_by = vars[ny2];
        // The rest of the constraints
        for (int t = 1; t < K; t++)
        {
            /// 位置约束 1
            AD<double> fx1 = vars[nx1 + t];
            AD<double> fy1 = vars[ny1 + t];
            AD<double> fth1 = vars[nth1 + t];
            AD<double> fv1 = vars[nv1 + t];
            AD<double> fvth1 = vars[nvth1 + t];

            AD<double> fx0 = vars[nx1 + t - 1];
            AD<double> fy0 = vars[ny1 + t - 1];
            AD<double> fth0 = vars[nth1 + t - 1];
            AD<double> fv0 = vars[nv1 + t - 1];
            AD<double> fvth0 = vars[nvth1 + t - 1];

            AD<double> facc0 = vars[nacc1 + t - 1];
            AD<double> faccth0 = vars[naccth1 + t - 1];

            accumulate_fx += fv0 * CppAD::cos(fth0) * dt;
            accumulate_fy += fv0 * CppAD::sin(fth0) * dt;

            fg[1 + nx1 + t] = fx1 - (fx0 + fv0 * CppAD::cos(fth0) * dt);
            fg[1 + ny1 + t] = fy1 - (fy0 + fv0 * CppAD::sin(fth0) * dt);
            fg[1 + nth1 + t] = fth1 - (fth0 + fvth0 * dt);
            fg[1 + nv1 + t] = fv1 - (fv0 + facc0 * dt);
            fg[1 + nvth1 + t] = fvth1 - (fvth0 + faccth0 * dt);

            //位置约束2
            AD<double> bx1 = vars[nx2 + t];
            AD<double> by1 = vars[ny2 + t];
            AD<double> bth1 = vars[nth2 + t];
            AD<double> bv1 = vars[nv2 + t];
            AD<double> bvth1 = vars[nvth2 + t];

            AD<double> bx0 = vars[nx2 + t - 1];
            AD<double> by0 = vars[ny2 + t - 1];
            AD<double> bth0 = vars[nth2 + t - 1];
            AD<double> bv0 = vars[nv2 + t - 1];
            AD<double> bvth0 = vars[nvth2 + t - 1];

            AD<double> bacc0 = vars[nacc2 + t - 1];
            AD<double> baccth0 = vars[naccth2 + t - 1];

            accumulate_bx += bv0 * CppAD::cos(bth0) * dt;
            accumulate_by += bv0 * CppAD::sin(bth0) * dt;

            fg[1 + nx2 + t - 2*(K-1)] = bx1 - (bx0 + bv0 * CppAD::cos(bth0) * dt);
            fg[1 + ny2 + t - 2*(K-1)] = by1 - (by0 + bv0 * CppAD::sin(bth0) * dt);
            fg[1 + nth2 + t - 2*(K-1)] = bth1 - (bth0 + bvth0 * dt);
            fg[1 + nv2 + t - 2*(K-1)] = bv1 - (bv0 + bacc0 * dt);
            fg[1 + nvth2 + t - 2*(K-1)] = bvth1 - (bvth0 + baccth0 * dt);
        }

        // 筛选前进方向前车到目标偏移
        double current_dist = pow(target_dx, 2) + pow(target_dy, 2);
        double target_forward_weight = 1.0 / (1.0 + exp(-1000.0 * ref_v));
        double target_backward_weight = 1.0 / (1.0 + exp(1000.0 * ref_v));
        AD<double> final_dx = CppAD::pow(accumulate_fx-target_dx, 2) * target_forward_weight + CppAD::pow(accumulate_bx-target_dx, 2) * target_backward_weight;
        AD<double> final_dy = CppAD::pow(accumulate_fy-target_dy, 2) * target_forward_weight + CppAD::pow(accumulate_by-target_dy, 2) * target_backward_weight;
        // 预测末端与目标点偏移
        // AD<double> final_dist = CppAD::sqrt(final_dx + final_dy);

        // 减速模式权重,越接近目标,该权重越大,最终减速到达目标
        double slowdown_weight = 1.0 / (1.0 + exp(10.0 * sqrt(current_dist) - 5.0));
        // char disp[255];
        // memset(disp, 0, 255);
        // sprintf(disp, "to target: %.3f, pred dist: %.3f, slowdown weight: %.3f\n", current_dist, final_dx, slowdown_weight);
        // std::cout << disp << std::endl;

        for (int t = 0; t < K; t++)
        {
            // 前车坐标系下，y坐标差
            AD<double> xt = vars[nx1 + t];
            AD<double> fx = coeffs[0] + coeffs[1] * xt + coeffs[2] * pow(xt, 2) + coeffs[3] * pow(xt, 3);
            fg[0] += y_cost_weight * CppAD::pow(vars[ny1 + t] - fx, 2);

            // 前车坐标系下，theta cost & v cost
            AD<double> fth = CppAD::atan(coeffs[1] + 2 * coeffs[2] * xt + 3 * coeffs[3] * pow(xt, 2));
            fg[0] += th_cost_weight * CppAD::pow(vars[nth1 + t] - fth, 2);
            // 修改运动速度损失函数, 从ref_v 转到 接近终点时的最终预测点的dist误差与v误差
            fg[0] += (1-slowdown_weight)*v_cost_weight * CppAD::pow(vars[nv1 + t] - ref_v, 2) + slowdown_weight * brake_weight * (final_dx / K + final_dy / K + CppAD::pow(vars[nv1 + K - 1], 2));
            fg[0] += vth_cost_weight * CppAD::pow(vars[nvth1 + t], 2);

            // 在后车坐标系下，后车y坐标差
            AD<double> xt2 = vars[nx2 + t];
            AD<double> fx2 = coeffs_b[0] + coeffs_b[1] * xt2 + coeffs_b[2] * pow(xt2, 2) + coeffs_b[3] * pow(xt2, 3);
            fg[0] += back_y_cost_weight * CppAD::pow(vars[ny2 + t] - fx2, 2);

            // 在后车坐标系下，后车theta cost & v cost
            AD<double> fth2 = CppAD::atan(coeffs_b[1] + 2 * coeffs_b[2] * xt2 + 3 * coeffs_b[3] * pow(xt2, 2));
            fg[0] += back_th_cost_weight * CppAD::pow(vars[nth2 + t] - fth2, 2);
            // 修改运动速度损失函数
            fg[0] += (1-slowdown_weight)*back_v_cost_weight * CppAD::pow(vars[nv2 + t] - ref_v, 2) + slowdown_weight * brake_weight * (final_dx / K + final_dy / K + CppAD::pow(vars[nv2 + K - 1], 2));
            fg[0] += back_vth_cost_weight * CppAD::pow(vars[nvth2 + t], 2);
        }
        //加速度 cost
        // Costs for steering (delta) and acceleration (a)
        for (int t = 0; t < K - 1; t++)
        {
            fg[0] += a_cost_weight * CppAD::pow(vars[nacc1+ t], 2);
            fg[0] += ath_cost_weight * CppAD::pow(vars[naccth1 + t], 2);

            fg[0] += back_a_cost_weight * CppAD::pow(vars[nacc2 + t], 2);
            fg[0] += back_ath_cost_weight * CppAD::pow(vars[naccth2 + t], 2);
        }

        ///添加两车距离损失
        const int distance_weight = 800;
        const int delta_th_weight = 10;
        char buf[255]={0};
        AD<double> dtheta = vars[nth1] - vars[nth2] + dth;
        for (int t = 0; t < K; ++t)
        {
            AD<double> x1_in_front = vars[nx1 + t];
            AD<double> y1_in_front = vars[ny1 + t];
            ///将坐标旋转到后车坐标系下
            AD<double> x1_in_back = CppAD::cos(dtheta) * x1_in_front - CppAD::sin(dtheta) * y1_in_front + dx;
            AD<double> y1_in_back = CppAD::sin(dtheta) * x1_in_front + CppAD::cos(dtheta) * y1_in_front + dy;

            AD<double> distance_x = x1_in_back - vars[nx2 + t];
            AD<double> distance_y = y1_in_back - vars[ny2 + t];

            fg[0] += distance_weight * CppAD::pow(CppAD::sqrt(CppAD::pow(distance_x, 2) + CppAD::pow(distance_y, 2)) - ref_dist, 2);
            fg[0] += delta_th_weight * CppAD::pow(CppAD::atan(y1_in_back / x1_in_back), 2);
            //sprintf(buf+strlen(buf), " %.4f ", sqrt(pow(distance_x,2)+pow(distance_y,2)));
            // // 两车最短距离约束
            // fg[1 + relative_pose_contrains + t - 4*(K-1)] = CppAD::sqrt(CppAD::pow(distance_x, 2) + CppAD::pow(distance_y, 2));
        }
        //std::cout<<buf<<std::endl;
    }
};

#include <chrono>
bool TowBotMPC::Solve(Eigen::VectorXd state, double max_v, Eigen::VectorXd coeffs, Eigen::VectorXd coeffs_b, std::vector<double> &out)
{
    bool ok = true;
    typedef CPPAD_TESTVECTOR(double) Dvector;
    // State vector holds all current values neede for vars below

    // std::cout<<"state: "<<state<<std::endl;
    double v1 = state[0];
    double v2 = state[1];
    double vth1 = state[2];
    double vth2 = state[3];
    double ref_v = state[4];
    double dt = state[5];

    // double dx = state[6]; /// 两车初始位姿差   前车相对后车坐标系
    // double dy = state[7];
    // double dth = state[8]; ///两坐标系角度

    size_t n_vars = 2 * (K * 5 + (K - 1) * 2);
    size_t n_constraints = K * (5 * 2 + 0);

    Dvector vars(n_vars);
    for (int i = 0; i < n_vars; i++)
    {
        vars[i] = 0.0;
    }
    for (int i = 0; i < K; ++i)
    {
        vars[nv1 + i] = v1;
        vars[nv2 + i] = v2;

        vars[nvth1 + i] = vth1;
        vars[nvth2 + i] = vth2;
    }

    Dvector vars_lowerbound(n_vars);
    Dvector vars_upperbound(n_vars);

    for (int i = 0; i < n_vars; i++)
    {
        vars_lowerbound[i] = -1.0e19;
        vars_upperbound[i] = 1.0e19;
    }

    for (int i = 0; i < K; i++)
    {
        //// limit  v1 v2
        vars_lowerbound[nv1 + i] = -max_v;
        vars_upperbound[nv1 + i] = max_v;

        vars_lowerbound[nv2 + i] = -max_v;
        vars_upperbound[nv2 + i] = max_v;

        //// limit  vth1 2
        vars_lowerbound[nvth1 + i] = -0.25;
        vars_upperbound[nvth1 + i] = 0.25;

        vars_lowerbound[nvth2 + i] = -0.25;
        vars_upperbound[nvth2 + i] = 0.25;
    }
    for (int i = 0; i < K - 1; i++)
    {
        //// limit acc1 acc2
        vars_lowerbound[nacc1 + i] = -1.0;
        vars_upperbound[nacc1 + i] = 1.0;

        vars_lowerbound[nacc2 + i] = -1.0;
        vars_upperbound[nacc2 + i] = 1.0;

        //// limit accth1 accth2
        vars_lowerbound[naccth1 + i] = -0.05;//3度
        vars_upperbound[naccth1 + i] = 0.05;

        vars_lowerbound[naccth2 + i] = -0.05;//3度
        vars_upperbound[naccth2 + i] = 0.05;
    }
    

    ////   constraints
    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;
    }

    for (int i = 0; i < 1; i++)
    {
        ////acc v1 v2
        constraints_lowerbound[nv1 + i] = -3.0;
        constraints_upperbound[nv1 + i] = 3.0;
        constraints_lowerbound[nv2 + i - 2*(K-1)] = -3.0;
        constraints_upperbound[nv2 + i - 2*(K-1)] = 3.0;

        //acc th 1 2
        constraints_lowerbound[nvth1 + i] = -0.1;
        constraints_upperbound[nvth1 + i] = 0.1;
        constraints_lowerbound[nvth2 + i - 2*(K-1)] = -0.1;
        constraints_upperbound[nvth2 + i - 2*(K-1)] = 0.1;

        // constraints_lowerbound[relative_pose_contrains + i - 4*(K-1)] = min_dist;
        // constraints_upperbound[relative_pose_contrains + i - 4*(K-1)] = max_dist;
    }
    // // acc v1 v2
    // constraints_lowerbound[nv1]=v1-0.1;
    // constraints_upperbound[nv1]=v1+0.1;
    // constraints_lowerbound[nv2-2*(K-1)]=v2-0.1;
    // constraints_upperbound[nv2-2*(K-1)]=v2+0.1;

    // //acc th1 th2
    // constraints_lowerbound[nvth1]=vth1-0.017;
    // constraints_upperbound[nvth1]=vth1+0.017;
    // constraints_lowerbound[nvth2-2*(K-1)]=vth2-0.1;
    // constraints_upperbound[nvth2-2*(K-1)]=vth2+0.1;

    // object that computes objective and constraintsv
    FG_eval_2 fg_eval(coeffs, coeffs_b, state);

    std::string options;
    options += "Integer print_level  0\n";
    options += "Sparse  true        forward\n";
    options += "Sparse  true        reverse\n";
    options += "Numeric max_cpu_time          0.5\n";

    CppAD::ipopt::solve_result<Dvector> solution;

    CppAD::ipopt::solve<Dvector, FG_eval_2>(
        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;

    // Return the first actuator values, along with predicted x and y values to plot in the simulator.
    out.clear();
    m_trajectory_f.clear();
    m_trajectory_b.clear();
    // float speed1=solution.x[nv1];
    // if(fabs(speed1)<0.1)
    // {
    //     speed1=speed1>0?0.1:-0.1;
    // }

    // float speed2=solution.x[nv2];
    // if(fabs(speed2)<0.1)
    // {
    //     speed2=speed2>0?0.1:-0.1;
    // }
    out.push_back(solution.x[nv1]);
    out.push_back(solution.x[nvth1]);
    out.push_back(solution.x[nv2]);
    out.push_back(solution.x[nvth2]);
    for (int i = 0; i < K; ++i)
    {
        tf::Pose pose1, pose2;
        pose1.setOrigin(tf::Vector3(solution.x[nx1 + i], solution.x[ny1 + i], 0));
        pose2.setOrigin(tf::Vector3(solution.x[nx2 + i], solution.x[ny2 + i], 0));
        m_trajectory_f.push_back(pose1);
        m_trajectory_b.push_back(pose2);
    }
    static double velo = 0;
    static std::chrono::steady_clock::time_point last_time = std::chrono::steady_clock::now();
    if (velo != 0)
    {
        double interval = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - last_time).count() * 0.001;
        double acc = (solution.x[nv1] - velo) / dt;
        // printf("is ok? %d, cost: %.3f\n", ok, cost);
        printf("Cost: %.3f\tv1:%.3f\tvth1:%.3f\t||\tnv1:%.3f\tnvth1:%.3f\tnacc1:%.3f\tinterval:%.3f\n", solution.obj_value, v1, vth1, solution.x[nv1], solution.x[nvth1], acc, interval);
        
        // printf("Cost : %.3f\tref:%.3f\tcv : %.3f\tcvth : %.3f\t||\tv2:%.3f\tth2:%.3f\n",solution.obj_value,
        //     ref_v,v2,vth2,solution.x[nv2],solution.x[nvth2]);
    }
    velo = solution.x[nv1];
    last_time = std::chrono::steady_clock::now();
    return ok;
}

TowBotMPC::TowBotMPC()
{
    // ROS_WARN("构造");
    // std::cout << "构造" << std::endl;
}
TowBotMPC::~TowBotMPC()
{
    // ROS_WARN("析构");
    // std::cout << "析构" << std::endl;
}
