#include "road_seg.h"
#include "sy_errorinfo.h"
#include "road_extractor.h"
#include "dvpp_processx.h"
#include <time.h>
#include <sys/time.h>
#include "stream_data.h"
#include <algorithm>
#include <map>
#include <cstring>
#include <cmath>


using namespace atlas_utils;
using namespace std;

struct Resource {
    aclrtContext ctx;
    aclrtStream stream;
};

typedef struct Tools {
    Resource src;
    ROADExtract* roadsegExtract;
    DvppProcessx* dvpp;

}Tools;


int rs_init(void **handle, rs_param param){

    int ret = SY_SUCCESS;
    Tools* tools = new Tools;
    // init resource
    ACL_CALL(aclrtCreateStream(&tools->src.stream), ACL_SUCCESS, SY_FAILED);

    // head_shoulder detection init
    tools->roadsegExtract = new ROADExtract();
    tools->roadsegExtract->config.confThr = param.thresld;

    ret = tools->roadsegExtract->Init(param.modelNames);
    if (ret != SY_SUCCESS) {
        delete tools->roadsegExtract;
        tools->roadsegExtract = nullptr;
        return SY_FAILED;
    }

    tools->dvpp = new DvppProcessx();
    tools->dvpp->InitResource(tools->src.stream);

    *handle = tools;
    
    return SY_SUCCESS;
}

namespace road_seg_7cls
{
    int lane_cls_classes = 6;   //5 + 1
    int road_seg_classes = 9;   //8类别+背景
    int direct_seg_classes = 5; //4类别+背景
    std::vector<float> lane_softmax(const std::vector<float>& x) {
        std::vector<float> result(x.size());
        // 实现softmax逻辑
        float max_val = *std::max_element(x.begin(), x.end());
        float sum = 0.0;
        for (size_t i = 0; i < x.size(); i++) {
            result[i] = std::exp(x[i] - max_val);
            sum += result[i];
        }
        for (size_t i = 0; i < x.size(); i++) {
            result[i] /= sum;
        }

        return result;
    }

    std::vector<float> lane_linspace(float start, float end, int n) {
        std::vector<float> result;
        float step = (end - start) / (n - 1);
        for (int i = 0; i < n; i++) {
            result.push_back(start + step * i);
        }
        return result;
    }

    bool Lane_IOU(const std::vector<float>& parent_box, const std::vector<float>& compared_box, int threshold) {
        // 实现Lane_IOU逻辑
        int y = lane_cls_classes, length = y + 2, offset_start = y + 3;

        int n_offsets = 72;
        int n_strips = n_offsets - 1;

        int start_a = static_cast<int>((parent_box[y] * n_strips) + 0.5);
        int start_b = static_cast<int>((compared_box[y] * n_strips) + 0.5);
        int start = std::max(start_a, start_b);
        int end_a = start_a + static_cast<int>(parent_box[length]) - 1 + 0.5 - ((parent_box[length] - 1 < 0) ? 1 : 0);
        int end_b = start_b + static_cast<int>(compared_box[length]) - 1 + 0.5 - ((compared_box[length] - 1 < 0) ? 1 : 0);
        int end = std::min(std::min(end_a, end_b), n_strips);
        if (end - start < 0) {
            return false;
        }
        int dist = 0;
        for (int i = offset_start + start; i < offset_start + end; i++) {
            if (i > (offset_start + end)) {
                break;
            }
            if (parent_box[i] < compared_box[i]) {
                dist += (compared_box[i] - parent_box[i]);
            } else {
                dist += (parent_box[i] - compared_box[i]);
            }
        }
        return dist < (threshold * (end - start + 1));
    }

    void Lane_nms(const std::vector<std::vector<float>>& proposals, std::vector<size_t> &keep_index, size_t &num_to_keep, const std::vector<float>& scores, int overlap = 50, int top_k = 4) {
        // 实现Lane_nms逻辑
        /*std::vector<float> sorted_score = scores;
        std::sort(sorted_score.begin(), sorted_score.end(), std::greater<float>());*/
        std::vector<int> indices(scores.size());
        for (size_t i = 0; i < indices.size(); i++) {
            indices[i] = i;
        }
        std::sort(indices.begin(), indices.end(), [&](int a, int b) { return scores[a] > scores[b]; });

        std::vector<int> r_filters(scores.size(), 0);

        for (size_t i = 0; i < indices.size(); i++) {
            int indice = indices[i];
            if (r_filters[i] == 1) {  // continue if this proposal is filtered by nms before
                continue;
            }
            keep_index.push_back(indice);
            if (keep_index.size() > static_cast<size_t>(top_k)) {  // break if more than top_k
                break;
            }
            if (i == indices.size() - 1) {  // break if indice is the last one
                break;
            }
            std::vector<int> sub_indices(indices.begin() + i + 1, indices.end());
            for (size_t sub_i = 0; sub_i < sub_indices.size(); sub_i++) {
                int sub_indice = sub_indices[sub_i];
                bool r_filter = Lane_IOU(proposals[indice], proposals[sub_indice], overlap);
                if (r_filter) {
                    r_filters[i + 1 + sub_i] = 1;
                }
            }
        }

        num_to_keep = static_cast<int>(keep_index.size());
    }


    int predictions_to_pred(const std::vector<std::vector<float>>& predictions, rs_lane* reg_array, int n_strips, int n_offsets) {
        std::vector<float> prior_ys = lane_linspace(1.0, 0.0, n_offsets);
        int lane_count = 0;
        int start_x = lane_cls_classes + 1, start_y = lane_cls_classes, offset_start = lane_cls_classes + 4, length_ = lane_cls_classes + 3;
        for (const auto& lane : predictions) {
            std::vector<float> lane_xs(lane.begin() + offset_start, lane.end());  // normalized value
            int start = std::min(std::max(0, static_cast<int>(std::round(lane[start_y] * n_strips))), n_strips);
            int length = static_cast<int>(std::round(lane[length_]));
            int end = start + length - 1;
            end = std::min(end, static_cast<int>(prior_ys.size()) - 1);
            /*
            std::vector<bool> mask(prior_ys.size(), false);
            for (int i = 0; i < start; i++) {
                if (lane_xs[i] >= 0.0 && lane_xs[i] <= 1.0) {
                    mask[i] = true;
                }
            }
            for (int i = end + 1; i < static_cast<int>(prior_ys.size()); i++) {
                if (lane_xs[i] >= 0.0 && lane_xs[i] <= 1.0) {
                    mask[i] = true;
                }
            }*/

            std::fill(lane_xs.begin() + end + 1, lane_xs.end(), -2);
            std::fill(lane_xs.begin(), lane_xs.begin() + start, -2);
            
            std::vector<float> lane_ys;
            for (size_t i = 0; i < prior_ys.size(); i++) {
                // if (!mask[i]) {
                if (lane_xs[i] >= 0) {
                    lane_ys.push_back(prior_ys[i]);
                }
            }

            std::vector<float> lane_xs_new;
            for (size_t i = 0; i < prior_ys.size(); i++) {
                // if (!mask[i]) {
                if (lane_xs[i] >= 0) {
                    lane_xs_new.push_back(lane_xs[i]);
                }
            }


            if (lane_xs_new.empty()) {
                continue;
            }
                
            int lane_xs_new_size = lane_xs_new.size();
            // if (lane_xs_new_size <= 1) {
            //     continue;
            // }

            if (lane_xs_new_size > 180) {
                lane_xs_new_size = 180;
            }

            int count = 0;
    #if 1
            for (size_t i = 0; i < lane_xs_new_size; i++) {
                // float lane_x = lane_xs_new[i],	lane_y = lane_ys[i];
                float lane_x = lane_xs_new[i] * 640;
                float lane_y = lane_ys[i] * 360;
                reg_array[lane_count].points[count].x_ = lane_x;
                reg_array[lane_count].points[count].y_ = lane_y;
                // printf("lane_x:%f lane_y:%f %d %d\n", lane_x, lane_y, reg_array[lane_count].points[count].x_, reg_array[lane_count].points[count].y_);
                count ++;
            } 
    #endif

            reg_array[lane_count].start_x = lane[start_x]* 640;
            reg_array[lane_count].start_y = lane[start_y]* 360;

            
            std::vector<float> lane_info(lane.begin()+1, lane.begin()+lane_cls_classes); // 去除背景类别
            std::vector<float> softmax_res = lane_softmax(lane_info);
            reg_array[lane_count].cls = std::distance(softmax_res.begin(),std::max_element(softmax_res.begin(),softmax_res.end())) + 1;
            reg_array[lane_count].conf = *(std::max_element(softmax_res.begin(),softmax_res.end()));
            reg_array[lane_count].num_points = count;
            lane_count ++;
        }
        return lane_count;
    }

    // 统一求nms
    int get_lanes(const vector<float> output, rs_lane* reg_array, float conf_threshold, int nms_thres, int max_lanes) {

        int stride1 = 192, stride2 = 82, num_points = 72, input_width = 640, input_height = 360;
        int theta = lane_cls_classes + 2, length = lane_cls_classes + 3;
        int n_strips = num_points - 1, n_offsets = num_points;
        std::vector<std::vector<float>> predictions, nms_predictions;
        std::vector<float> scores;
        for (int i = 0; i < stride1; i++) {
            std::vector<float> pred_score(output.begin() + i*stride2, output.begin() + i*stride2+lane_cls_classes); // 存储类别得分
            std::vector<float> softmax_res = lane_softmax(pred_score);
            std::vector<float> selected_scores(softmax_res.begin()+1, softmax_res.begin()+lane_cls_classes); // 选择目标类别的概率值，去除背景类别 
            float max_score = *std::max_element(selected_scores.begin(), selected_scores.end());
            if (max_score >= conf_threshold) {
                std::vector<float> pred(output.begin() + i*stride2, output.begin() + (i+1)*stride2);
                nms_predictions.push_back(pred);
                predictions.push_back(pred);
                scores.push_back(max_score);
            }
        }
        
        for (auto& row : nms_predictions) {
            row[theta] *= n_strips;
            for (size_t i = length; i < row.size(); i++) {
                row[i] *= (input_width - 1);
            }
        }

        std::vector<size_t> keep;
        size_t num_to_keep;
        Lane_nms(nms_predictions, keep, num_to_keep, scores, nms_thres, max_lanes);
        keep.resize(num_to_keep);

        std::vector<std::vector<float>> final_predictions;
        for (auto& index : keep) {
            predictions[index][length] = round(predictions[index][length] * n_strips);
            final_predictions.push_back(predictions[index]);
        }

        int lane_count = predictions_to_pred(final_predictions, reg_array, n_strips, n_offsets);
        return lane_count;
    }
        
    // 对每个类别分别求nms
    int get_lanesV2(const vector<float> output, rs_lane* reg_array, float conf_threshold, int nms_thres, int max_lanes) {

        int stride1 = 192, stride2 = 82, num_points = 72, input_width = 640, input_height = 360;
        int theta = lane_cls_classes + 2, length = lane_cls_classes + 3;
        int n_strips = num_points - 1, n_offsets = num_points;

        std::vector<std::vector<float>> final_predictions;
        for (int cls = 0; cls < lane_cls_classes; cls++) {
            std::vector<std::vector<float>> predictions, nms_predictions;
            std::vector<float> scores;
            for (int i = 0; i < stride1; i++) {
                std::vector<float> pred_score(output.begin() + i*stride2, output.begin() + i*stride2+lane_cls_classes); // 存储类别得分
                std::vector<float> softmax_res = lane_softmax(pred_score);
                std::vector<float> selected_scores(softmax_res.begin()+1, softmax_res.begin()+lane_cls_classes); // 选择目标类别的概率值，去除背景类别 
                
                auto iter = std::max_element(selected_scores.begin(), selected_scores.end());
                int indice = std::distance(selected_scores.begin(), iter);
                float max_score = *iter;
                if (indice == cls && max_score >= conf_threshold) {
                    std::vector<float> pred(output.begin() + i*stride2, output.begin() + (i+1)*stride2);
                    nms_predictions.push_back(pred);
                    predictions.push_back(pred);
                    scores.push_back(max_score);
                }
            }
            
            for (auto& row : nms_predictions) {
                row[theta] *= n_strips;
                for (size_t i = length; i < row.size(); i++) {
                    row[i] *= (input_width - 1);
                }
            }

            std::vector<size_t> keep;
            size_t num_to_keep;
            Lane_nms(nms_predictions, keep, num_to_keep, scores, nms_thres, max_lanes);
            keep.resize(num_to_keep);

            
            for (auto& index : keep) {
                predictions[index][length] = round(predictions[index][length] * n_strips);
                final_predictions.push_back(predictions[index]);
            }
        }

        int lane_count = predictions_to_pred(final_predictions, reg_array, n_strips, n_offsets);
        return lane_count;
    }


    // 获取道路分割结果中最大值的索引（类别）
    // int road_seg_show(const float* inference_features, uint8_t *road_seg) {
    int road_seg_show(const vector<float> inference_features, uint8_t *road_seg, int seg_num_cls) {
        int h = 360, w = 640;
        int step_size = h*w;
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                float max_score = 0; int max_cls = 0;
                for (int k = 0; k < seg_num_cls; k++) {
                    int cur_steps = i*w+j+k*step_size;
                    if (inference_features[cur_steps] > max_score) {
                        max_score = inference_features[cur_steps];
                        max_cls = k;
                    }   
                }
                // if (max_cls == 5 || max_cls == 6) max_cls = 0;	// 将分割的车道线变成背景，防止混淆================
                road_seg[i*w+j] = max_cls;
            }
        }

        return 0;
    }
}

int rs_process(void *handle, sy_img img_data, rs_result result)
{
	return rs_batch(handle, &img_data, 1, &result);
}

int rs_batch(void * handle, sy_img *image_data_array, int batchsize, rs_result *result){
    Tools* tools = (Tools*) handle;

    int inputW = tools->roadsegExtract->GetInputWidth();
    int inputH = tools->roadsegExtract->GetInputHeight();
    
    //printf("debug inputw:%d,inputh:%d\n",inputW,inputH);

    for (int b = 0; b < batchsize; b++) {
        if (image_data_array[b].data_ == NULL || image_data_array[b].w_ == 0 || image_data_array[b].h_ == 0) {
            ERROR_LOG(" RoadSegExtract get null input ptr!");
	        return SY_FAILED;
        }
        
        ImageData resizeImg, src;
        // Utils::CopysyImageDataToDvpp(src, image_data_array[b]);
		ACL_CALL(Utils::CopysyImageDataToDvppV2(src, image_data_array[b]), SY_SUCCESS, SY_FAILED);
        ACL_CALL(tools->dvpp->CropAndPadding(resizeImg, src, inputW, inputH), SY_SUCCESS, SY_FAILED);
        // forward
        // double t1, t2;
        // t1 = msecond();
        int ret = tools->roadsegExtract->Inference(resizeImg);
        if (ret != SY_SUCCESS) {
            return SY_MODEL_FORWARD_ERROR;
        }
        // t2 = msecond();
        // printf("debug infer time: %.2f\n", t2 - t1);

        vector<float> detRes;
        ret = tools->roadsegExtract->PostProcess(detRes);
        if (ret != SY_SUCCESS) {
            return SY_MODEL_GETRESULT_ERROR;
        }

		int lane_length = 192 * 82;
        int roadseg_length = 360 * 640 * road_seg_7cls::road_seg_classes;
		std::vector<float> clane_detection(detRes.begin(), detRes.begin() + lane_length); 
		std::vector<float> croad_segmentation(detRes.begin() + lane_length, detRes.begin() + lane_length + roadseg_length);
        std::vector<float> cdirect_segmentation(detRes.begin() + lane_length + roadseg_length, detRes.end());

        int res = road_seg_7cls::road_seg_show(croad_segmentation, result[b].seg_array, road_seg_7cls::road_seg_classes);
        res = road_seg_7cls::road_seg_show(cdirect_segmentation, result[b].direct_seg, road_seg_7cls::direct_seg_classes);
		float conf_threshold = 0.4;
		int nms_thres = 50, max_lanes = 21;
		int lane_count = road_seg_7cls::get_lanesV2(clane_detection, result[b].reg_array, conf_threshold, nms_thres, max_lanes);
		result[b].lane_count = lane_count;
	
        vector<float>().swap(detRes);  
		vector<float>().swap(clane_detection);
		vector<float>().swap(croad_segmentation);  
        vector<float>().swap(cdirect_segmentation);        
    }

    return SY_SUCCESS;
}

void rs_release(void **handle) {
    if (*handle) {
        Tools* tools = (Tools*) *handle;
        if (tools->roadsegExtract) {
            delete tools->roadsegExtract;
            tools->roadsegExtract = nullptr;
        }
        if (tools->dvpp) {
            delete tools->dvpp;
            tools->dvpp = nullptr;
        }
        // aclFinalize();
        delete tools;
        tools = NULL;
    }
}

const char * rs_get_version() {
    return "road_seg_vdec_arm_v310p_0.0.2.20240110_without_timelimit";
}