#include "laneatt.hh"

#include <fstream>
#include <NvInferPlugin.h>
#include <cuda_runtime_api.h>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <map>

#define INPUT_H 360
#define INPUT_W 640
#define N_OFFSETS 72
#define N_STRIPS (N_OFFSETS - 1)
#define MAX_COL_BLOCKS 1000

//std::vector<cv::Point2f> warped_point, origin_point;

//float METER_PER_PIXEL;

//float BOARD_SIDE_LENGTH = 0.5;

//cv::Mat get_pers_mat(std::string pers_file)
//{
//    cv::Point2f src[4];
//    cv::Point2f dst[4];
//    cv::FileStorage pf(pers_file, cv::FileStorage::READ);
//    pf["board_left_top"]>>src[0];
//    pf["board_left_bottom"]>>src[1];
//    pf["board_right_top"]>>src[2];
//    pf["board_right_bottom"]>>src[3];

//    dst[1].x = src[0].x;
//    dst[1].y = src[1].y;

//    dst[3].x = src[2].x;
//    dst[3].y = src[3].y;

//    int side_length = dst[3].x - dst[1].x;

//    METER_PER_PIXEL = BOARD_SIDE_LENGTH / side_length;

//    dst[0].x = src[0].x;
//    dst[0].y = dst[1].y - side_length;

//    dst[2].x = src[2].x;
//    dst[2].y = dst[3].y - side_length;

//    cv::Mat pers_mat = getPerspectiveTransform(dst, src);
//    cv::Mat pers_mat_inv = pers_mat.inv();

//    return pers_mat_inv;
//}

//float DX (float x1, float x2, float dis) {
//    return dis / (x2 - x1);
//}

//float DY (float y1, float y2, float dis) {
//    return dis / (x2 - x1);
//}



//std::string PERS_FILE = "./yaml/pers.yaml";
//cv::Mat pers_mat_inv = get_pers_mat(PERS_FILE);  // 得到透视变换关系

LaneATT::LaneATT(const std::string& plan_path) {
    buffer_size_[0] = 3 * INPUT_H * INPUT_W;
    buffer_size_[1] = MAX_COL_BLOCKS * (5 + N_OFFSETS);

    cudaMalloc(&buffers_[0], buffer_size_[0] * sizeof(float));
    cudaMalloc(&buffers_[1], buffer_size_[1] * sizeof(float));
    image_data_.resize(buffer_size_[0]);
    detections_.resize(MAX_COL_BLOCKS);

    cudaStreamCreate(&stream_);
    LoadEngine(plan_path);
}

LaneATT::~LaneATT() {
    cudaStreamDestroy(stream_);
    for (auto& buffer : buffers_) {
        cudaFree(buffer);
    }
    if (context_ != nullptr) {
        context_->destroy();
        engine_->destroy();
    }
}

std::vector<std::vector<cv::Point2f>> LaneATT::DetectLane(const cv::Mat& raw_image) {
    // Preprocessing
    cv::Mat img_resize;
    cv::resize(raw_image, img_resize, cv::Size(INPUT_W, INPUT_H), cv::INTER_LINEAR);
    // img_resize.convertTo(img_resize, CV_32FC3, 1.0);
    uint8_t* data_hwc = reinterpret_cast<uint8_t*>(img_resize.data);
    float* data_chw = image_data_.data();
    for (int c = 0; c < 3; ++c) {
        for (unsigned j = 0, img_size = INPUT_H * INPUT_W; j < img_size; ++j) {
            data_chw[c * img_size + j] = data_hwc[j * 3 + c] / 255.f;
        }
    }

    // Do inference
    cudaMemcpyAsync(buffers_[0], image_data_.data(), buffer_size_[0] * sizeof(float), cudaMemcpyHostToDevice, stream_);
    context_->execute(1, &buffers_[0]);
    cudaMemcpyAsync(detections_.data(), buffers_[1], buffer_size_[1] * sizeof(float), cudaMemcpyDeviceToHost, stream_);

    // NMS and decoding
    cv::Mat rst;
    std::vector<std::vector<cv::Point2f>> lanes;
    lanes = PostProcess(img_resize, raw_image);
    return lanes;
}

void LaneATT::LoadEngine(const std::string& engine_file) {
    std::ifstream in_file(engine_file, std::ios::binary);
    if (!in_file.is_open()) {
        std::cerr << "Failed to open engine file: " << engine_file << std::endl;
        return;
    }
    in_file.seekg(0, in_file.end);
    int length = in_file.tellg();
    in_file.seekg(0, in_file.beg);
    std::unique_ptr<char[]> trt_model_stream(new char[length]);
    in_file.read(trt_model_stream.get(), length);
    in_file.close();

    // FIXME: getPluginCreator could not find plugin: ScatterND version: 1
    initLibNvInferPlugins(&g_logger_, "");
    nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(g_logger_);
    assert(runtime != nullptr);
    engine_ = runtime->deserializeCudaEngine(trt_model_stream.get(), length, nullptr);
    assert(engine_ != nullptr);
    context_ = engine_->createExecutionContext();
    assert(context_ != nullptr);

    runtime->destroy();
}

float LaneIoU(const Detection& a, const Detection& b) {
    int start_a = static_cast<int>(a.start_y *  + 0.5f);
    int start_b = static_cast<int>(b.start_y *  + 0.5f);
    int start = std::max(start_a, start_b);
    int end_a = start_a + static_cast<int>(a.length + 0.5f) - 1;
    int end_b = start_b + static_cast<int>(b.length + 0.5f) - 1;
    int end = std::min(std::min(end_a, end_b), N_STRIPS);
    // if (end < start) {
    //     return 1.0f / 0.0f;
    // }
    float dist = 0.0f;
    for (int i = start; i <= end; ++i) {
        dist += fabs(a.lane_xs[i] - b.lane_xs[i]);
    }
    dist /= static_cast<float>(end - start + 1);
    return dist;
}

std::vector<std::vector<cv::Point2f>> LaneATT::PostProcess(cv::Mat& lane_image, cv::Mat raw_image, float conf_thresh, float nms_thresh, int nms_topk) { //conf_thresh=0.4 nms_thresh=50 nms_topk=4
    // 1.Do NMS
    std::vector<Detection> candidates;
    std::vector<Detection> proposals;
    for (auto det : detections_) {
        if (det.score > conf_thresh) {
            candidates.push_back(det);
        }
    }
    // std::cout << candidates.size() << std::endl;
    std::sort(candidates.begin(), candidates.end(), [=](const Detection& a, const Detection& b) { return a.score > b.score; });
    for (int i = 0; i < candidates.size(); ++i) {
        if (candidates[i].score < 0.0f) {
            continue;
        }
        proposals.push_back(candidates[i]);
        if (proposals.size() == nms_topk) {
            break;
        }
        for (int j = i + 1; j < candidates.size(); ++j) {
            if (candidates[j].score > 0.0f && LaneIoU(candidates[j], candidates[i]) < nms_thresh) {
                candidates[j].score = -1.0f;
            }
        }
    }

    // 2.Decoding
    std::vector<float> anchor_ys;
    for (int i = 0; i < N_OFFSETS; ++i) {
        anchor_ys.push_back(1.0f - i / float(N_STRIPS));
    }
    std::vector<std::vector<cv::Point2f>> lanes;
    for (const auto& lane: proposals) { 
        int start = static_cast<int>(lane.start_y * N_STRIPS + 0.5f);
        int end = start + static_cast<int>(lane.length + 0.5f) - 1;
        end = std::min(end, N_STRIPS);
        std::vector<cv::Point2f> points;
        for (int i = start; i <= end; ++i) {
            points.push_back(cv::Point2f(lane.lane_xs[i], anchor_ys[i] * INPUT_H));
        }
        lanes.push_back(points);
    }

    std::vector<std::vector<cv::Point2f>> lanes1;
    float y_ratio =  float(raw_image.size().height) / float(INPUT_H);
    float x_ratio =  float(raw_image.size().width) / float(INPUT_W);
    cv::Point2f pp;
    for (const auto& lane_points : lanes) {
        std::vector<cv::Point2f> points;
        for (int i=0; i<lane_points.size(); i++)  {
            pp.x =float(lane_points[i].x) *  x_ratio;
            pp.y =float(lane_points[i].y) * y_ratio;
            points.push_back(pp);
        }
        lanes1.push_back(points);
    }
    return lanes1;


//    float y_ratio =  float(raw_image.size().height) / float(INPUT_H);
//    float x_ratio =  float(raw_image.size().width) / float(INPUT_W);

//    // 3.PerspectiveTransform

//    std::vector<cv::Point2f> leftLane;
//    std::vector<cv::Point2f> rightLane;
//    std::vector<std::vector<cv::Point2f>> warped_lanes;

//    for (int i = 0; i < lanes.size(); i++) {
//        cv::perspectiveTransform(lanes[i], warped_point , pers_mat_inv); //  坐标变换(输入点，输出点，变换矩阵)
//        warped_lanes.push_back(warped_point);
//        double sum = 0.0;
//        for (int i = 0 ; i < 5 ; i++){
//            sum += warped_point[i].x;
//        }
//        double mean = sum / 5;
//        if (mean < INPUT_W / 2){
//            leftLane.push_back(cv::Point2f(mean, i));
//        }else{
//            rightLane.push_back(cv::Point2f(mean, i));
//        }
//    }
//    float left_min = 0;
//    float right_max = INPUT_W;
//    int left_num = -1;
//    int right_num = -1;
//    int left_left = -1;
//    int right_right = -1;
//    for (const auto& points : leftLane){
//        if(points.x > left_min){
//            left_num = (int) round(points.y);
//            left_min = points.x;
//        }
//    }
//    for (const auto& points : leftLane){
//        if((int)round(points.y) != left_num){
//            left_left = points.y;
//        }
//    }
//    if (left_num == -1){
//        std::cout<<"left lane failed"<<std::endl;
//    }
//    for (const auto& points : rightLane){
//        if(points.x < right_max){
//            right_num = (int) round(points.y);
//            right_max = points.x;
//        }
//    }
//    for (const auto& points : rightLane){
//        if((int)round(points.y) != right_num){
//            right_right = points.y;
//        }
//    }
//    if (right_num == -1){
//        std::cout<<"right lane failed"<<std::endl;
//    }

//    // 4.Visualize

//    cv::Point2f pp;
//    cv::Point3f ppp;
//    int flag = 0;
//    cv::Mat img_warp;
//    int DX, DY;
//    std::vector<std::vector<cv::Point3f>> world_lanes;
//    cv::warpPerspective(raw_image, img_warp, pers_mat_inv, raw_image.size());        //  图像变换(输入图像，输出图像，变换矩阵)
//    for (const auto& lane_points : warped_lanes) {
//        if(flag == left_num || flag == right_num){
//            std::vector<cv::Point3f> world_points;
//            for (const auto& point : lane_points) {
//                pp.x =float(point.x);
//                pp.y =float(point.y);
//                cv::circle(img_warp, pp, 1, cv::Scalar(0, 255, 0), -1);
//                ppp.x = (point.x - INPUT_W / 2) * DX;
//                ppp.y = float((INPUT_H - point.y) * DY);
//                ppp.z = float(0);
//                world_points.push_back(ppp);
//            }
//            world_lanes.push_back(world_points);
//        }else{
//            std::vector<cv::Point3f> world_points;
//            for (const auto& point : lane_points) {
//                pp.x =float(point.x);
//                pp.y =float(point.y);
//                cv::circle(img_warp, pp, 1, cv::Scalar(0, 0, 255), -1);
//                ppp.x = (point.x - INPUT_W / 2) * DX;
//                ppp.y = float((INPUT_H - point.y) * DY);
//                ppp.z = float(0);
//                world_points.push_back(ppp);
//            }
//            world_lanes.push_back(world_points);
//        }
//        flag++;
//    }
//    cv::imshow("warp",img_warp);



//    flag = 0;
//    for (const auto& lane_points : lanes) {
//        if(flag == left_num || flag == right_num){
//            for (int i=0; i<lane_points.size()-1; i++)  {
//                pp.x =float(lane_points[i].x) *  x_ratio;
//                pp.y =float(lane_points[i].y) * y_ratio;
//                cv::circle(raw_image, pp, 1, cv::Scalar(0, 255, 0), -1);

//            }
//        }else{
//            for (int i=0; i<lane_points.size()-1; i++) {
//                pp.x =float(lane_points[i].x) *  x_ratio;
//                pp.y =float(lane_points[i].y) * y_ratio;
//                cv::circle(raw_image, pp, 1, cv::Scalar(0, 0, 255), -1);
//            }
//        }
//        flag++;
//    }
//    cv::Point2f p1, p2;
//    flag = 0;
//    for (const auto& lane_points : lanes) {
//        if(flag == left_num || flag == right_num){
//            for (int i=0; i<lane_points.size()-1; i++) {
//                p1.x =float(lane_points[i].x) *  x_ratio;
//                p1.y =float(lane_points[i].y) * y_ratio;
//                p2.x =float(lane_points[i+1].x) *  x_ratio;
//                p2.y =float(lane_points[i+1].y) *  y_ratio;
//                cv::line(raw_image, p1, p2, cv::Scalar(0, 255, 0), 1);

//            }
//        }else{
//            for (int i=0; i<lane_points.size()-1; i++) {
//                p1.x =float(lane_points[i].x) *  x_ratio;
//                p1.y =float(lane_points[i].y) * y_ratio;
//                p2.x =float(lane_points[i+1].x) *  x_ratio;
//                p2.y =float(lane_points[i+1].y) *  y_ratio;
//                cv::line(raw_image, p1, p2, cv::Scalar(0, 0, 255), 1);
//            }
//        }
//        flag++;
//    }

//    return raw_image;
}