#include "road_seg_statistics.h"
#include "omp.h"

RoadSegProcess::RoadSegProcess(){
    m_max_batchsize = 16;
}

RoadSegProcess::~RoadSegProcess(){
    release();
}


cv::Mat RoadSegProcess::mask_to_rgb(cv::Mat img, cv::Mat mask) {
    cv::Mat masks = img.clone();
    int reg_cls = 9;
    for (int i = 0; i < masks.rows; i++) {
        for (int j = 0; j < masks.cols; j++) {
            for (int k = 1; k < reg_cls; k++) {
                if (mask.at<int>(i,j) == k) {
                    masks.at<cv::Vec3b>(i,j)[0] = seg_colors[k][0];
                    masks.at<cv::Vec3b>(i,j)[1] = seg_colors[k][1];
                    masks.at<cv::Vec3b>(i,j)[2] = seg_colors[k][2];
                }
            }
        }
    }
    return masks;
}

float RoadSegProcess::contourArea(std::vector<cv::Point> contour, cv::Point2f& center) {
    cv::RotatedRect rect = cv::minAreaRect(contour);  // 最小外接矩形 rect[0]中心点 rect[1]宽 高  rect[2]旋转角度
    center = rect.center;
    return cv::contourArea(contour);
}

void RoadSegProcess::lanes_process(const rs_lane* lanes, int lane_count, std::vector<std::pair<std::vector<cv::Point>, int>>& combined, float scale_w, float scale_h) {   
    std::vector<std::vector<cv::Point> > lanes_xys;
    std::vector<int> lanes_cls;
	for (int i = 0; i < lane_count; i++) {
        std::vector<cv::Point> xys;
		for (int j = 0; j < lanes[i].num_points; j++) {
            int x = static_cast<int>(lanes[i].points[j].x_ * scale_w);
            int y = static_cast<int>(lanes[i].points[j].y_ * scale_h);
            if (x > 0 && y > 0) {
                xys.emplace_back(x, y);
            }
        }
        if (!xys.empty()) {
            lanes_xys.push_back(xys);
            lanes_cls.push_back(lanes[i].cls);
        }
    }

    for (size_t i = 0; i < lanes_xys.size(); ++i) {
        combined.push_back(std::make_pair(lanes_xys[i], lanes_cls[i]));
    }
    if (!combined.empty()) {
        //按车道线起点坐标排序,相应的类别顺序也会变化以保证标签对齐
		std::sort(combined.begin(), combined.end(), [](const std::pair<std::vector<cv::Point>, int>& a, const std::pair<std::vector<cv::Point>, int>& b) {
			return a.first[0].x < b.first[0].x;
		});
    }
}

cv::Mat RoadSegProcess::imshow_lanes(cv::Mat img, const rs_lane* lanes, int lane_count) {
	float scale_w = img.cols / 640.0;
	float scale_h = img.rows / 360.0;
    std::vector<std::pair<std::vector<cv::Point>, int>> combined;
    lanes_process(lanes, lane_count, combined, scale_w, scale_h);

    for (const auto& lane_info : combined) {
        const auto& xys = lane_info.first;
        int cls = lane_info.second;
		cv::Scalar color(lane_colors[cls][0],lane_colors[cls][1],lane_colors[cls][2]);
        for (size_t i = 1; i < xys.size(); ++i) {
            cv::line(img, xys[i - 1], xys[i], color, 4);
        }
    }
    
    return img;
}

int RoadSegProcess::Mask2LanePoints(const cv::Mat& pred, std::vector<std::vector<cv::Point>>&lanes, std::vector<int>& cats) {
    std::vector<int> labels = {9, 10, 11, 12, 13, 14};
    for(auto cat: labels) {
        cv::Mat b_masks, measure_labels, stats, centroids;
        cv::compare(pred, cat, b_masks, cv::CMP_EQ); //将pred元素逐个和cat比较，相同255，不同0
        // cv::threshold(pred, b_masks, 126, 255, cv::THRESH_OTSU);    //生成二值图
        // 连通域计算 b_masks:闭操作后的二值图像 measure_labels:和原图一样大的标记图 centroids:nccomps×2的矩阵,表示每个连通域的质心(x, y)
        // stats:nccomps×5的矩阵 表示每个连通区域的外接矩形和面积(x, y, w, h, area)，索引0是背景信息
        int nccomps = cv::connectedComponentsWithStats (b_masks, measure_labels, stats, centroids, 4); //8
        for(int cv_measure_id = 1; cv_measure_id < nccomps; cv_measure_id++ ) {	//跳过背景信息0
            cv::Mat cv_measure_mask;
            cv::compare(measure_labels, cv_measure_id, cv_measure_mask, cv::CMP_EQ);
            std::vector< std::vector< cv::Point> > contours;
            cv::findContours(cv_measure_mask,contours,cv::noArray(),cv::RETR_TREE,cv::CHAIN_APPROX_SIMPLE);
            std::vector<cv::Point> contours_poly;
            for (int j = 0; j < contours.size();j++) {
                float area = cv::contourArea(contours[j]);  if (area < 60) continue; //30
                cv::approxPolyDP(cv::Mat(contours[j]), contours_poly, 6, true); //减小epsilon提高拟合精度(需要调试，epsilon过小会使多边形不够简化，单条线变多条)
                if (contours_poly.size() == 1) continue;
                lanes.push_back(contours_poly);
                cats.push_back(cat);
            }
        }    
    }
    return 0;
}


cv::Mat RoadSegProcess::seg_post_process(bool large_resolution, unsigned char *seg_array, std::vector<std::pair<std::vector<cv::Point>, int>> combined, std::vector<std::vector<cv::Point>> &poly_masks, std::vector<int> &region_classes, std::vector<std::vector<cv::Point>> &lanes, std::vector<int> &cats, std::map<double, int> &x_sort) {
    std::vector<int> pred_cls; 
    int h = 360, w = 640;
    cv::Mat lanes_masks = cv::Mat_<int>(h,w); //正常分割结果
    cv::Mat mask_rmlane = cv::Mat_<int>(h,w); //车道线区域置为背景
    cv::Mat solve_masks = cv::Mat_<int>(h,w); //计算主行驶区域用(虚线及减速标线归入行车道)
    int step_size = h*w;
    int seg_min_region_area = 512; //1024 

	for (int i = 0; i < h; ++i) {
		for (int j = 0; j < w; ++j) {
			int max_cls = seg_array[(i * w + j)];
			lanes_masks.ptr<int>(i)[j] = max_cls;
            pred_cls.push_back(max_cls);
            mask_rmlane.ptr<int>(i)[j] = max_cls;
            solve_masks.ptr<int>(i)[j] = max_cls;
		}
	}

    for (const auto& lane_info : combined) {
        const auto& xys = lane_info.first;
        int cls = lane_info.second;
		if (cls == 1 || cls == 2) cls = 9; //实线
		if (cls == 3 || cls == 6) cls = 10; //虚线 
        if (cls == 4) cls = 11; /*停止线*/  if (cls == 5) cls = 12; /*减速标线*/
        if (cls == 7) cls = 13; /*道路方向、车道分割栏*/
        if (cls == 8) cls = 14; /*道理保护栏*/
        for (size_t i = 1; i < xys.size(); ++i) {
            cv::line(lanes_masks, xys[i - 1], xys[i], cls, 4); //绘制车道线用于求连通域
            if (cls == 10)   cls = 1; //求主行驶区域时将虚线归入行车道 
			cv::line(solve_masks, xys[i - 1], xys[i], cls, 4);
            
        }
    }

    /* 求背景区域--mask车道区域，场景变化用
    cv::compare(lanes_masks, 0, background_mask, cv::CMP_EQ); //将lanes_masks元素逐个和0比较，相同255，不同0*/
#if 0
    //--mask远处区域------------------------------------------------
    cv::Mat mask_black = mask_rmlane.clone();
    if (large_resolution)   mask_black(cv::Rect(0, 0, w, int(h * 0.14))) = 0;  
    else    mask_black(cv::Rect(0, 0, w, int(h * 0.22))) = 0;           
    mask_black(cv::Rect(0, int(h * 0.95), w, int(h * 0.05))) = 0;
    mask_black(cv::Rect(0, 0, int(w * 0.02), h)) = 0;
    mask_black(cv::Rect(int(w * 0.95), 0, int(w * 0.05), h)) = 0;
    mask_rmlane = mask_black;
    //-------------------------------------------------------------------------
#endif
    //2.去重获取预测到的类别
    std::vector<int> labels(pred_cls);
    std::sort(labels.begin(),labels.end());
    labels.erase(std::unique(labels.begin(),labels.end()),labels.end());
  
    //4.求车道线区域
    int flag = Mask2LanePoints(lanes_masks, lanes, cats); ///////////////////////// 车道线如何与mask结合
 
    //5.求道路区域
    int count = 0;
    for(auto cat: labels) {
        // std::cout << cat << std::endl;
        cv::Mat b_masks, measure_labels, stats, centroids;
        cv::Mat n_masks = cv::Mat_<int>(h,w);
        n_masks = cv::Scalar(255);
        cv::compare(mask_rmlane, cat, b_masks, cv::CMP_EQ); //将mask_rmlane元素逐个和cat比较，相同255，不同0
        //连通域计算
        int nccomps = cv::connectedComponentsWithStats (b_masks, measure_labels, stats, centroids, 8);
        for(int i = 1; i < nccomps; i++ ) {	//跳过背景信息0
            //移除过小的区域，并将对应位置置为0
            if (stats.at<int>(i, cv::CC_STAT_AREA) < seg_min_region_area) {
                cv::Mat comparison_result;
                cv::compare(measure_labels, cv::Scalar(i), comparison_result, cv::CMP_EQ); //相等为255不等为0
                n_masks.setTo(0, comparison_result); // 将comparison_result中非零区域置为0

                cv::bitwise_and(mask_rmlane,n_masks,mask_rmlane);
                continue;
            }

            double centr_x = centroids.at<double>(i, 0);
            double centr_y = centroids.at<double>(i, 1);

            int region_class = cat;
            // printf("region_class: %d\n", region_class);
            cv::Mat region_mask;
            cv::bitwise_and(measure_labels,n_masks,measure_labels);/////
            cv::compare(measure_labels, i, region_mask, cv::CMP_EQ);

            std::vector< std::vector< cv::Point> > contours;
            cv::findContours(region_mask,contours,cv::noArray(),cv::RETR_TREE,cv::CHAIN_APPROX_SIMPLE);
            std::vector<cv::Point> contours_poly;
            for (int j = 0; j < contours.size();j++) {
                cv::approxPolyDP(cv::Mat(contours[j]), contours_poly, 10, true);
                if (contours_poly.size() <= 2) continue;
                poly_masks.push_back(contours_poly);
                region_classes.push_back(cat);
                if (x_sort.count(centr_x)) centr_x += 0.0001;
                x_sort.insert(std::make_pair(centr_x, count));
                ++ count;
                // for (auto iter: contours_poly) {
                //     std::cout << "contour " << iter << " " << cat << std::endl;
                // }
    
            }
        }    
    }

#if 0
    //6.draw lanes
    for (int i = 0; i < lanes.size(); ++i) {
        int thickness = 4;
        for (int j = 0; j < lanes[i].size()-1; ++j) {
            cv::line(mask_rmlane, lanes[i][j], lanes[i][j+1], cats[i], thickness);
            // std::cout << lanes[i][j] << " " << lanes[i][j+1] << " " << cats[i] << std::endl;
        }
    }
#endif
    return mask_rmlane;
}


/* 算法初始化 */
int RoadSegProcess::init(int gpu_id, string models_dir) {
    string model_path = models_dir + "/models/village/tzroad_seg_310p.om" ;

    LOG_INFO("road_seg 版本：{}  模型路径：{}", rs_get_version(), model_path);

    rs_param param;
    char modelNames[100];
    strcpy(modelNames, model_path.c_str());
    param.modelNames = modelNames;
    param.thresld = 0.25;
    param.devId = gpu_id;

    m_devId = param.devId;
    ACL_CALL(aclrtSetDevice(m_devId), ACL_SUCCESS, -1);
    ACL_CALL(aclrtCreateContext(&m_algorthim_ctx, m_devId), ACL_SUCCESS, -1);

    int ret = rs_init(&m_seg_handle, param); 
    if(ret != 0) {
        LOG_ERROR("Init RoadSegSdk failed error code is {}", ret);
        return -1;
    }

    return 0;
}

/* 算法计算 */
int RoadSegProcess::process_gpu(sy_img * batch_img, vector<DeviceMemory*> vec_segMem, vector<string>& tasklist,
					            vector<vector<vector<int>>>& traffic_region, vector<vector<int>>& labels)
{
    int batchsize = tasklist.size();
    traffic_region.resize(batchsize), labels.resize(batchsize);

    rs_result result[batchsize];
	int fea_size = SEG_IMG_RES_W * SEG_IMG_RES_H;
	for (int b = 0; b < batchsize; b++) {
		result[b].seg_array = new unsigned char[fea_size];
	}

    /* 路数太多时 按照最大batchsize数 拆批次运行 */
    int ret = aclrtSetCurrentContext(m_algorthim_ctx);
    if (ACL_SUCCESS != ret) {
        return false;
    }
	int cur_batch_size = m_max_batchsize;
	int cycleTimes = batchsize / cur_batch_size + (batchsize % cur_batch_size == 0 ? 0 : 1);
	for (int c = 0; c < cycleTimes; c++) {
		int real_batchsize = c == cycleTimes - 1 ? (batchsize - cur_batch_size*c) : cur_batch_size;
		int startbatch = c*cur_batch_size;

		rs_result *real_res = result + startbatch;
		ret = rs_batch(m_seg_handle, batch_img + startbatch, real_batchsize, real_res);
		
	}

    for (int b = 0; b < batchsize; b++) {			
		std::vector<std::pair<std::vector<cv::Point>, int>> combined;
		lanes_process(result[b].reg_array, result[b].lane_count, combined);
		std::vector<std::vector<cv::Point>> poly_masks, lanes;
        std::vector<int> region_classes, cats;
		std::map<double, int> x_sort;
		bool large_resolution = false;
		if (batch_img[b].w_ > 1920) large_resolution = true;

        cv::Mat seg_output = seg_post_process(large_resolution, result[b].seg_array, combined, poly_masks, region_classes, lanes, cats, x_sort); //m_masks:mask前的结果  poly_masks后的结果               
	
        float width_ratio = batch_img[b].w_*1.0 / SEG_IMG_RES_W;
        float height_ratio = batch_img[b].h_*1.0 / SEG_IMG_RES_H;

        for (int i = 0; i < lanes.size(); ++i) {
            std::vector<cv::Point> points = lanes[i];
            int lane_cls = cats[i];
            vector<int> cur_region;
            for (int j = 0; j < points.size(); ++j) {
                int px = points[j].x*width_ratio;   int py = points[j].y*height_ratio;
                cur_region.push_back(px);  cur_region.push_back(py);
            }
            traffic_region[b].push_back(cur_region); // 存储区域
            labels[b].push_back(lane_cls); // 存储类别
        } 
        
        for (int i = 0; i < poly_masks.size(); ++i) {
            std::vector<cv::Point> points = poly_masks[i];
            int seg_cls = region_classes[i];
            vector<int> cur_region;
            for (int j = 0; j < points.size(); ++j) {
                int px = points[j].x*width_ratio;   int py = points[j].y*height_ratio;
                cur_region.push_back(px);  cur_region.push_back(py);
            }
            traffic_region[b].push_back(cur_region); // 存储区域
            labels[b].push_back(seg_cls); // 存储类别
        }

   
#if 1
	{
        jpegUtil.jpeg_init(m_devId);
        std::string file_path = "/data/shzhao/vpt_ascend/bin/res/seg/";
		auto time_now = std::chrono::system_clock::now();
		std::string cur_timestamp_us =  std::to_string(std::chrono::duration_cast<std::chrono::microseconds>(time_now.time_since_epoch()).count());
		std::string img_filename = file_path + cur_timestamp_us + "_" + tasklist[b] + ".jpg";
		vpc_img_info src_img_info = VPCUtil::vpc_devMem2vpcImg(vec_segMem[b]);
        bool bSaved = jpegUtil.jpeg_encode(src_img_info.pic_desc, img_filename);
        std::cout << bSaved << " " << img_filename << std::endl;
        cv::Mat cur_img = cv::imread(img_filename);
        cv::resize(cur_img, cur_img, cv::Size(640,360), 0, 0, cv::INTER_LINEAR);
        cv::Mat vis_image = mask_to_rgb(cur_img, seg_output);
        cv::imwrite(img_filename, vis_image);
		VPCUtil::vpc_img_release(src_img_info);
        jpegUtil.jpeg_release();
	}
#endif
   
        std::vector<std::vector<cv::Point>>().swap(poly_masks);
        std::vector<std::vector<cv::Point>>().swap(lanes);
        std::vector<int>().swap(region_classes);
        std::vector<int>().swap(cats);
    }

    for (int b = 0; b < batchsize; b++) {
		delete[] result[b].seg_array; result[b].seg_array = NULL;
	}

    return 0;
}


/* 算法句柄 资源释放 */
void RoadSegProcess::release(){
    if (m_seg_handle){
        rs_release(&m_seg_handle);
        m_seg_handle = NULL;
    }
    if(m_algorthim_ctx){
        aclrtDestroyContext(m_algorthim_ctx);
    }
}