#include "wheel_detector.h"
#include <ceres/ceres.h>
#include <ceres/cubic_interpolation.h>

struct singleWheelCostFunction {
    std::vector<float> line_;
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_;
    Eigen::Vector4f centroid;

    explicit singleWheelCostFunction(const std::vector<float> &line, pcl::PointCloud<pcl::PointXYZ>::Ptr cloud) {
        line_ = line;
        cloud_ = cloud;
        pcl::compute3DCentroid(*cloud, centroid);
    }

    template<typename JetT>
    bool operator()(const JetT *const x, JetT *residual) const {
        ceres::Grid1D<float> grid(line_.data(), 0, line_.size());
        ceres::CubicInterpolator<ceres::Grid1D<float>> cubic(grid);

        JetT delta_x = x[0] + JetT(centroid.x()  + 2);
        JetT theta = x[1];
        auto cos_t = ceres::cos(theta);
        auto sin_t = ceres::sin(theta);
        for (int i = 0; i < cloud_->size(); ++i) {
            JetT cx = JetT(cloud_->points[i].x - centroid.x()) * cos_t - JetT(cloud_->points[i].y - centroid.y()) * sin_t + delta_x;
            cubic.Evaluate(cx * 1000.,&residual[i]);
            residual[i]=JetT(1.0)-residual[i];
        }
        return true;
    }
};

struct doubleWheelCostFunction {
    std::vector<float> line_l_;
    std::vector<float> line_r_;
    pcl::PointCloud<pcl::PointXYZ>::Ptr m_left_cloud;
    pcl::PointCloud<pcl::PointXYZ>::Ptr m_right_cloud;


    explicit doubleWheelCostFunction(const std::vector<float> &linel,const std::vector<float> &liner, pcl::PointCloud<pcl::PointXYZ>::Ptr left_cloud, pcl::PointCloud<pcl::PointXYZ>::Ptr right_cloud) {
        line_l_ = linel;
        line_r_ = liner;
        m_left_cloud = left_cloud;
        m_right_cloud = right_cloud;
    }

    template<typename JetT>
    bool operator()(const JetT *const x, JetT *residual) const {
        ceres::Grid1D<float> gridl(line_l_.data(), 0, line_l_.size());
        ceres::CubicInterpolator<ceres::Grid1D<float>> cubicl(gridl);

        ceres::Grid1D<float> gridr(line_r_.data(), 0, line_r_.size());
        ceres::CubicInterpolator<ceres::Grid1D<float>> cubicr(gridr);


        JetT theta = x[0];
        JetT w = x[1];
        JetT tx = x[2];

        auto cos_t = ceres::cos(theta);
        auto sin_t = ceres::sin(theta);

        int resi_index = 0;
        for (auto &point: m_left_cloud->points) {
            JetT x=JetT(point.x);
            JetT y=JetT(point.y);
            JetT cx = x * cos_t - y * sin_t +tx + w*0.5;
            cubicl.Evaluate((cx +JetT(0.5)) * 1000.,&residual[resi_index]);
            residual[resi_index]=JetT(1.0)-residual[resi_index];
            resi_index++;
        }

        for (auto &point: m_right_cloud->points) {
            JetT x=JetT(point.x);
            JetT y=JetT(point.y);
            JetT cx = x * cos_t - y * sin_t + tx - w*0.5;
            cubicr.Evaluate((cx +JetT(2.0))* 1000.,&residual[resi_index]);
            residual[resi_index]=JetT(1.0)-residual[resi_index];
            resi_index++;
        }

        return true;
    }
};

struct WheelsTogetherCostFunction {
    std::vector<Eigen::Vector4f> m_vct_centroid;
    std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> m_vct_cloud;
    std::vector<ceres::CubicInterpolator<ceres::Grid1D<float>>*> m_vct_model;

    explicit WheelsTogetherCostFunction(std::vector<ceres::CubicInterpolator<ceres::Grid1D<float>>*> vct_model, std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> &vct_cloud, std::vector<Eigen::Vector4f> &vct_centroid) {
        m_vct_cloud = vct_cloud;
        m_vct_model = vct_model;
        m_vct_centroid = vct_centroid;
    }

    template<typename JetT>
    bool operator()(const JetT *const x, JetT *residual) const {
        int residual_index = 0;
        JetT delta_x[4] = {x[0] + JetT(m_vct_centroid[0].x() + 2), x[2] + JetT(m_vct_centroid[1].x() + 2), x[4] + JetT(m_vct_centroid[2].x() + 2), x[6] + JetT(m_vct_centroid[3].x() + 2)};
        JetT theta[4] = {x[1], x[3], x[5], x[7]};
        JetT cos_t[4] = {ceres::cos(theta[0]), ceres::cos(theta[1]), ceres::cos(theta[2]), ceres::cos(theta[3])};
        JetT sin_t[4] = {ceres::sin(theta[0]), ceres::cos(theta[1]), ceres::cos(theta[2]), ceres::cos(theta[3])};

        for (int device_index = 0; device_index < 4; device_index++) {
            for (auto &point: m_vct_cloud[device_index]->points) {
                JetT cx = JetT(point.x - m_vct_centroid[device_index].x()) * cos_t[device_index] - JetT(point.y - m_vct_centroid[device_index].y()) * sin_t[device_index] + delta_x[device_index];
                m_vct_model[device_index]->Evaluate(cx * 1000.,&residual[residual_index]);
                residual[residual_index]=JetT(1.0)-residual[residual_index];
                residual_index++;
            }
        }

        return true;
    }
};

bool WheelCeresDetector::detect(pcl::PointCloud<pcl::PointXYZ>::Ptr in_cloud_ptr,
                                pcl::PointCloud<pcl::PointXYZ>::Ptr out_cloud_ptr,
                                WheelCeresDetector::WheelDetectResult &result, bool left_model) {
    Eigen::Vector4f centroid;
    pcl::compute3DCentroid(*in_cloud_ptr, centroid);
    double t_vars[2] = {-centroid.x(), 0};
    LOG(INFO) << "ceres detect init info: " << t_vars[0] << " " << t_vars[1];
    auto model = left_model ? singleLeftWheelModel() : singleRightWheelModel();
    // 构建最小二乘问题
    ceres::Problem problem;

    problem.AddResidualBlock(new ceres::AutoDiffCostFunction<singleWheelCostFunction, ceres::DYNAMIC, 2>(
                                     (new singleWheelCostFunction(model, in_cloud_ptr)), in_cloud_ptr->size()),
                             nullptr, // 核函数，这里不使用，为空
                             t_vars   // 待估计参数
    );

//    problem.SetParameterLowerBound(t_vars, 1, (-20) * M_PI / 180.0f);
//    problem.SetParameterUpperBound(t_vars, 1, (20) * M_PI / 180.0f);

    // 配置求解器
    ceres::Solver::Options options;                             // 这里有很多配置项可以填
    options.linear_solver_type = ceres::DENSE_QR;               // 增量方程如何求解
    options.minimizer_progress_to_stdout = false;               // 输出到cout
    options.max_num_iterations = 100;

    ceres::Solver::Summary summary;                             // 优化信息
    ceres::Solve(options, &problem, &summary);                  // 开始优化
    if(summary.iterations.size()<=1){
        LOG(INFO) << "仅迭代一轮，优化异常";
        return false;
    }
    if(summary.iterations[summary.iterations.size()-1].iteration<=1){
        LOG(INFO) << "最终仅迭代一轮，优化异常";
        return false;
    }
    LOG(INFO) << "ceres solve total time: " << summary.total_time_in_seconds * 1000 << "ms";

    result.theta = t_vars[1] * 180.0f / M_PI;

    float sin_t = std::sin(t_vars[1]);
    float cos_t = std::cos(t_vars[1]);
    for (auto &point: in_cloud_ptr->points) {
        float x = (point.x - centroid.x()) * cos_t - (point.y - centroid.y()) * sin_t + t_vars[0] + centroid.x();
        float y = 1.0 / (1.0 + ceres::exp(-400 * (x + 0.03))) - 1.0 / (1.0 + ceres::exp(-400 * (x-0.025)));
        if (y > 0.7) {
            out_cloud_ptr->push_back(point);
        }
    }
    pcl::compute3DCentroid(*out_cloud_ptr, centroid);
    result.center.x = centroid.x();
    result.center.y = centroid.y();
    result.center.z = centroid.z();
    LOG(INFO) << "num: " << out_cloud_ptr->size() << " " << result.info();

    return true;
}

bool WheelCeresDetector::detect(pcl::PointCloud<pcl::PointXYZ>::Ptr left_in_cloud,
                                pcl::PointCloud<pcl::PointXYZ>::Ptr left_out_cloud,
                                WheelCeresDetector::WheelDetectResult &left_result,
                                pcl::PointCloud<pcl::PointXYZ>::Ptr right_in_cloud,
                                pcl::PointCloud<pcl::PointXYZ>::Ptr right_out_cloud,
                                WheelCeresDetector::WheelDetectResult &right_result) {
    if (left_in_cloud->empty() && right_in_cloud->empty()) {
        return false;
    } else if (left_in_cloud->empty() && !right_in_cloud->empty()) {
        return detect(right_in_cloud, right_out_cloud, right_result, false);
    } else if (!left_in_cloud->empty() && right_in_cloud->empty()) {
        return detect(left_in_cloud, left_out_cloud, left_result, true);
    }


    double t_vars[3] = {0, 3, 0};
    std::vector<float> linel,liner;
    doubleWheelModel(linel,liner);

    // 构建最小二乘问题
    ceres::Problem problem;

    problem.AddResidualBlock(new ceres::AutoDiffCostFunction<doubleWheelCostFunction, ceres::DYNAMIC, 3>(
                                     (new doubleWheelCostFunction(linel,liner, left_in_cloud, right_in_cloud)), left_in_cloud->size() + right_in_cloud->size()),
                             nullptr, // 核函数，这里不使用，为空
                             t_vars   // 待估计参数
    );

//    problem.SetParameterLowerBound(t_vars, 1, (-20) * M_PI / 180.0f);
//    problem.SetParameterUpperBound(t_vars, 1, (20) * M_PI / 180.0f);

    // 配置求解器
    ceres::Solver::Options options;                             // 这里有很多配置项可以填
    options.linear_solver_type = ceres::DENSE_QR;               // 增量方程如何求解
    options.minimizer_progress_to_stdout = false;               // 输出到cout
    options.max_num_iterations = 100;

    ceres::Solver::Summary summary;                             // 优化信息
    ceres::Solve(options, &problem, &summary);                  // 开始优化
    if(summary.iterations.size()<=1){
        LOG(INFO) << "仅迭代一轮，优化异常";
        return false;
    }
    if(summary.iterations[summary.iterations.size()-1].iteration<=1){
        LOG(INFO) << "最终仅迭代一轮，优化异常";
        return false;
    }
    LOG(INFO) << "ceres solve total time: " << summary.total_time_in_seconds * 1000 << "ms";

    left_result.theta = t_vars[0] * 180.0f / M_PI;
    right_result.theta = t_vars[0] * 180.0f / M_PI;

    LOG(INFO) << left_result.info();
    LOG(INFO) << t_vars[1] << " " << t_vars[2];
    float sin_t = std::sin(t_vars[0]);
    float cos_t = std::cos(t_vars[0]);

    for (auto &point: left_in_cloud->points) {
        left_out_cloud->emplace_back((point.x ) * cos_t - (point.y ) * sin_t  + t_vars[2] + t_vars[1]/2,
                                     (point.x ) * sin_t + (point.y) * cos_t ,
                                     point.z);
    }
    for (auto &point: right_in_cloud->points) {
        right_out_cloud->emplace_back((point.x ) * cos_t - (point.y) * sin_t  + t_vars[2] - t_vars[1]/2,
                                      (point.x) * sin_t + (point.y ) * cos_t,
                                      point.z);
    }
    return true;
}

bool WheelCeresDetector::detect(std::vector<pcl::PointCloud<pcl::PointXYZ>::Ptr> vct_wheels_cloud,
                                pcl::PointCloud<pcl::PointXYZ>::Ptr left_out_cloud,
                                WheelCeresDetector::WheelDetectResult *wheels_result) {
    std::vector<Eigen::Vector4f> vct_centroid;
    size_t all_wheel_size = 0;
    for (auto &wheel_cloud: vct_wheels_cloud) {
        Eigen::Vector4f centroid;
        pcl::compute3DCentroid(*wheel_cloud, centroid);
        vct_centroid.push_back(centroid);
        all_wheel_size += wheel_cloud->size();
    }

    std::vector<ceres::CubicInterpolator<ceres::Grid1D<float>>*> vct_model;
    std::vector<float> left_model, right_model;
    left_model = singleLeftWheelModel();
    right_model = singleRightWheelModel();
    ceres::Grid1D<float> left_grid(left_model.data(), 0, left_model.size());
    ceres::Grid1D<float> right_grid(right_model.data(), 0, right_model.size());
    vct_model.push_back(new ceres::CubicInterpolator<ceres::Grid1D<float>>(left_grid));
    vct_model.push_back(new ceres::CubicInterpolator<ceres::Grid1D<float>>(right_grid));
    vct_model.push_back(new ceres::CubicInterpolator<ceres::Grid1D<float>>(left_grid));
    vct_model.push_back(new ceres::CubicInterpolator<ceres::Grid1D<float>>(right_grid));

    double t_vars[8] = {-vct_centroid[0].x(), 0, -vct_centroid[1].x(), 0, -vct_centroid[2].x(), 0, -vct_centroid[3].x(), 0};

    // 构建最小二乘问题
    ceres::Problem problem;
    problem.AddResidualBlock(new ceres::AutoDiffCostFunction<WheelsTogetherCostFunction, ceres::DYNAMIC, 8>(
                                     (new WheelsTogetherCostFunction(vct_model, vct_wheels_cloud, vct_centroid)), all_wheel_size),
                             nullptr, // 核函数，这里不使用，为空
                             t_vars   // 待估计参数
    );

//    problem.SetParameterLowerBound(t_vars, 1, (-20) * M_PI / 180.0f);
//    problem.SetParameterUpperBound(t_vars, 1, (20) * M_PI / 180.0f);

    // 配置求解器
    ceres::Solver::Options options;                             // 这里有很多配置项可以填
    options.linear_solver_type = ceres::DENSE_QR;               // 增量方程如何求解
    options.minimizer_progress_to_stdout = false;               // 输出到cout
    options.max_num_iterations = 100;

    ceres::Solver::Summary summary;                             // 优化信息
    ceres::Solve(options, &problem, &summary);                  // 开始优化
    if(summary.iterations.size()<=1){
        LOG(INFO) << "仅迭代一轮，优化异常";
        return false;
    }
    if(summary.iterations[summary.iterations.size()-1].iteration<=1){
        LOG(INFO) << "最终仅迭代一轮，优化异常";
        return false;
    }
    LOG(INFO) << "ceres solve total time: " << summary.total_time_in_seconds * 1000 << "ms";

    wheels_result[0].theta = t_vars[0] * 180.0f / M_PI;
    wheels_result[1].theta = t_vars[2] * 180.0f / M_PI;
    wheels_result[2].theta = t_vars[4] * 180.0f / M_PI;
    wheels_result[3].theta = t_vars[6] * 180.0f / M_PI;
    LOG(INFO) << wheels_result[0].info();
    LOG(INFO) << wheels_result[1].info();
    LOG(INFO) << wheels_result[2].info();
    LOG(INFO) << wheels_result[3].info();
    return true;
}

std::vector<float> WheelCeresDetector::singleWheelModel() {
    float min = -2;
    float max = 0.5;
    float e = 0.001;
    std::vector<float> line;
    pcl::PointCloud<pcl::PointXYZ>::Ptr model(new pcl::PointCloud<pcl::PointXYZ>);
    for (int i = min / e; i < max / e; i++) {
        float x = i * e;
        float y = std::min<float>(exp(-x * x), 1.0f - 1.0f / (1.0f + exp(-400 * (x - 0.025))));
        // float y = std::max<float>(-100 * x * x + 1, 0);
        y = 1.0 / (1.0 + ceres::exp(-400 * (x + 0.025))) - 1.0 / (1.0 + ceres::exp(-400 * (x-0.025)));
        model->emplace_back(x, 1 - y, 0);
        line.emplace_back(y);
    }
    Point3D_tool::WriteTxtCloud(ETC_PATH PROJECT_NAME "/data/save/single_double_wheel_model.txt", model);
    return line;
}

std::vector<float> WheelCeresDetector::singleLeftWheelModel() {
    static std::vector<float> line;
    if (!line.empty()) {
        return line;
    }
    float min = -2;
    float max = 0.5;
    float e = 0.001;
    for (int i = min / e; i < max / e; i++) {
        float x = i * e;
        // float y = std::min<float>(exp(-x * x), 1.0f - 1.0f / (1.0f + exp(-400 * (x - 0.025))));
        // float y = std::max<float>(-100 * x * x + 1, 0);
        float y = 1.0 / (1.0 + ceres::exp(-400 * (x + 0.03))) - 1.0 / (1.0 + ceres::exp(-400 * (x-0.025)));
        line.emplace_back(y);
    }
    return line;
}

std::vector<float> WheelCeresDetector::singleRightWheelModel() {
    static std::vector<float> line;
    if (!line.empty()) {
        return line;
    }
    float min = -2;
    float max = 0.5;
    float e = 0.001;
    for (int i = min / e; i < max / e; i++) {
        float x = i * e;
        // float y = std::min<float>(exp(-x * x), 1.0f - 1.0f / (1.0f + exp(-400 * (x - 0.025))));
        // float y = std::max<float>(-100 * x * x + 1, 0);
        float y = 1.0 / (1.0 + ceres::exp(-400 * (x + 0.025))) - 1.0 / (1.0 + ceres::exp(-400 * (x-0.03)));
        line.emplace_back(y);
    }
    return line;
}

void WheelCeresDetector::doubleWheelModel( std::vector<float>& linel,std::vector<float>& liner) {
    float min = -2;
    float max = 0.5;
    float e = 0.001;

    pcl::PointCloud<pcl::PointXYZ>::Ptr model(new pcl::PointCloud<pcl::PointXYZ>);
    for (int i = min / e; i < max / e; i++) {
        float x = i * e;
        float y = std::min<float>(exp(-x * x), 1.0f - 1.0f / (1.0f + exp(-400 * (x - 0.025))));
        // float y = std::max<float>(-100 * x * x + 1, 0);
        model->emplace_back(x, 1 - y, 0);
        liner.emplace_back(y);
    }

    Point3D_tool::WriteTxtCloud(ETC_PATH PROJECT_NAME "/data/save/r_model.txt", model);

    min = -0.5;
    max = 2;
    e = 0.001;
    model->clear();
    for (int i = min / e; i < max / e; i++) {
        float x = i * e;
        float y = std::min<float>(exp(-x * x), 1.0f - 1.0f / (1.0f + exp(-400 * (-x - 0.025))));
        // float y = std::max<float>(-100 * x * x + 1, 0);
        model->emplace_back(x, 1 - y, 0);
        linel.emplace_back(y);
    }



    Point3D_tool::WriteTxtCloud(ETC_PATH PROJECT_NAME "/data/save/l_model.txt", model);
}