#include "RoadSegAnalysis.h"

uint8_t seg_colors[][3] = { {0, 0, 0}, {0, 255, 255}, {128, 255, 0}, {255, 128, 0}, {128, 0, 255}, {255, 0, 128}, {0, 128, 255}, {0, 255, 128}, {128, 255, 255}};
uint8_t lane_colors[][3] = { {0, 0, 0}, {255, 0, 0}, {0, 255, 0}, {0, 0, 255}, {255, 255, 0}, {255, 0, 255}, {0, 255, 255}, {128, 255, 0}, {255, 128, 0}};

void lanes_process(const rs_lane* lanes, int lane_count, std::vector<std::pair<std::vector<cv::Point>, int>>& combined, float scale_w = 1.0, float scale_h = 1.0) {   
    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;
		});
    }
}

RoadSegAnalysis::RoadSegAnalysis(/* args */)
{
}

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

int RoadSegAnalysis::init(int devId){
    ACL_CALL(aclrtCreateContext(&ctx, devId), SY_SUCCESS, SY_FAILED);

    rs_param param;
    param.modelNames = "./models/road_seg/tzroad_seg240108_310p.om";
    param.thresld = 0.25;
    param.devId = devId;

    cout << "rs_init start " << endl;
    int ret = rs_init(&m_handle, param);
    if (ret != 0) {
        return -1;
	}

    cout << "rs_init success " << endl;

    return SY_SUCCESS;
}

std::vector<RoadInfo> RoadSegAnalysis::detect(vector<sy_img> vec_img){

    std::vector<RoadInfo> vec_road;

    const int batchsize = vec_img.size();
    rs_result results[batchsize];
    const int fea_size = 360*640;
    for (int b = 0; b < batchsize; b++) {
        results[b].seg_array = new unsigned char[fea_size];
        results[b].direct_seg = new unsigned char[fea_size];
    }

    int ret = SY_FAILED;

    do
    {
        ret = aclrtSetCurrentContext(ctx);
        if (SY_SUCCESS != ret) {
            printf("aclrtSetCurrentContext failed!");
            break;
        }

        ret = rs_batch(m_handle, vec_img.data(), batchsize, results);
        if (SY_SUCCESS != ret) {
            printf("rs_batch failed!");
            break;
        }
        
        for (int b = 0; b < batchsize; b++) {
            auto one_result = results[b];
            RoadInfo one_road;
            for (size_t i = 0; i < one_result.lane_count; i++)
            {
                auto one_lane = one_result.reg_array[i];
                LineInfo info;
                for (size_t j = 0; j < one_lane.num_points; j++)
                {
                    info.vec_pt.push_back(one_lane.points[j]);
                }
                info.line_type = one_lane.cls;
                one_road.vec_line.push_back(info);
            }

            one_road.vec_road = parse_seg(one_result, vec_img[b]);
            
            vec_road.push_back(one_road);
        }

    } while (0);

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

    return vec_road;
}

int RoadSegAnalysis::release() {

    ACL_CALL(aclrtSetCurrentContext(ctx), SY_SUCCESS, SY_FAILED);	

    if (m_handle) {
        rs_release(&m_handle);
    }

    if(ctx){
        aclrtDestroyContext(ctx);
        ctx = nullptr;
    }

    return SY_SUCCESS;
}

std::vector<SegInfo> RoadSegAnalysis::parse_seg(rs_result one_result, sy_img src) {
    int w = 640;
    int h = 360;
    float alpha = 0.75;
    cv::Mat overlayed_img(cv::Size(w,h), CV_8UC3, cv::Scalar(0, 0, 0));
    // 将车道线标签转换为彩色图像
    for (int i = 0; i < h; ++i) {
        for (int j = 0; j < w; ++j) {
            int idx = one_result.seg_array[(i * w + j)];
            overlayed_img.at<cv::Vec3b>(i, j)[0] = seg_colors[idx][0]; // R通道;
            overlayed_img.at<cv::Vec3b>(i, j)[1] = seg_colors[idx][1]; // G通道
            overlayed_img.at<cv::Vec3b>(i, j)[2] = seg_colors[idx][2]; // B通道
        }
    }

    cv::resize(overlayed_img, overlayed_img, cv::Size(src.w_,src.h_), 0, 0, cv::INTER_LINEAR); 
    // 将原始图像和彩色车道线图进行混合
    // cv::addWeighted(cvImg, alpha, overlayed_img, 1 - alpha, 0, overlayed_img);

    float scale_w = src.w_ / 640.0;
	float scale_h = src.h_ / 360.0;
    std::vector<std::pair<std::vector<cv::Point>, int>> combined;
    lanes_process(one_result.reg_array, one_result.lane_count, combined, scale_w, scale_h);

    std::vector<SegInfo> vec_seg;
    for (const auto& lane_info : combined) {
        const auto& xys = lane_info.first;
        int cls = lane_info.second;
        SegInfo info;
        info.seg_type = cls;
		// 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(overlayed_img, xys[i - 1], xys[i], color, 4);
            sy_point pt;
            pt.x_ = xys[i].x;
            pt.y_ = xys[i].y;
            info.vec_pt.push_back(pt);
        }
        vec_seg.push_back(info);
    }

    return vec_seg;
}