/*
 * @Author: turboLIU
 * @Date: 2022-07-25 09:26:32
 * @LastEditTime: 2024-07-08 13:26:31
 * @Description: Do not edit
 * @FilePath: /Cudayolo_track/utils.cpp
 */
#include <iostream>
#include <fstream>
#include <algorithm>
#include <map>
#include <sys/stat.h>
#include "utils.h"
// #include "detection_type_api.h"
#define M_PI       3.14159265358979323846 

float sigmoid(float x){
	// return 1./(1+exp(-x));
	return x;
}

std::vector<float>  softmax(std::vector<float> scores){
	std::vector<float> newScores;
	float total;
	for(auto s : scores){
		float ss = exp(s);
		total += ss;
		newScores.push_back(ss);
	}
	for(int i=0; i<newScores.size(); i++){
		newScores[i] /= total;
	}
	return newScores;
}

bool InVector(int idx, std::vector<uint8_t> flags){
	for(auto f : flags){
		if(idx == f)
		return true;
	}
	return false;
};

void splitStr(const std::string& str, std::vector<std::string>& tokens, char delim){
    tokens.clear();
    
    std::istringstream iss(str);
    std::string tmp;
    while (std::getline(iss, tmp, delim)) {
        if (tmp != "") {
            // 如果两个分隔符相邻，则 tmp == ""，忽略。
            tokens.emplace_back(std::move(tmp));
        }
    }
}

std::string splitKeyname(std::string imgname){
	std::vector<std::string> strs;
    std::vector<std::string> keys;
    splitStr(imgname, strs, '/');
    printf("imgname: %s\n", strs[strs.size() -1].c_str());
    splitStr(strs[strs.size()-1], keys, '.');
    std::string keyname = keys[0];
	return keyname;
}


void nms(std::vector<objinfo> &input, std::vector<objinfo> &output, float nmsthreshold, int type)
{//  this function has bug~~~~~~~~ not bug ,,det's param has negtive, maybe
	if (input.empty()) {
		return;
	}
	std::sort(input.begin(), input.end(),
		[](const objinfo& a, const objinfo& b)
	{
		return a.score < b.score;
	});
	printf("box num: %zd\n", input.size());
	float IOU = 0;
	float maxX = 0;
	float maxY = 0;
	float minX = 0;
	float minY = 0;
	std::vector<int> vPick;
	int nPick = 0;
	std::multimap<float, int> vScores;
	const int num_boxes = input.size();
	vPick.resize(num_boxes);
	for (int i = 0; i < num_boxes; ++i) {
		vScores.insert(std::pair<float, int>(input[i].score, i));
	}
	printf("vScores size: %zd\n",vScores.size());
	// printf("checking %s, %d\n",__FILE__, __LINE__);
	while (vScores.size() > 0) {
		int last = vScores.rbegin()->second;
		vPick[nPick] = last;
		nPick += 1;
		for (std::multimap<float, int>::iterator it = vScores.begin(); it != vScores.end();) {
			int it_idx = it->second;
			maxX = std::max(input.at(it_idx).x1, input.at(last).x1);
			maxY = std::max(input.at(it_idx).y1, input.at(last).y1);
			minX = std::min(input.at(it_idx).x2, input.at(last).x2);
			minY = std::min(input.at(it_idx).y2, input.at(last).y2);
			//maxX1 and maxY1 reuse 
			maxX = ((minX - maxX + 1) > 0) ? (minX - maxX + 1) : 0;
			maxY = ((minY - maxY + 1) > 0) ? (minY - maxY + 1) : 0;
			//IOU reuse for the area of two bbox
			IOU = maxX * maxY;
			if (type == NMS_UNION){
				IOU = IOU / (input.at(it_idx).area + input.at(last).area - IOU);
			}else if (type == NMS_MIN){
				IOU = IOU / ((input.at(it_idx).area < input.at(last).area) ? input.at(it_idx).area : input.at(last).area);
			}
			if (IOU > nmsthreshold) {
				it = vScores.erase(it);
			}
			else {
				it++;
			}
			// printf("checking %s, %d\n",__FILE__, __LINE__);
		}
		// printf("checking %s, %d\n",__FILE__, __LINE__);

	}
	// printf("checking %s, %d\n",__FILE__, __LINE__);
	vPick.resize(nPick);
	output.resize(nPick);
	for (int i = 0; i < nPick; i++) {
		output[i] = input[vPick[i]];
	}
}

void nms1(std::vector<objinfo>& input, std::vector<objinfo>& output, float nmsthreshold, int type)
{
	std::sort(input.begin(), input.end(),
		[](const objinfo& a, const objinfo& b)
	{
		return a.score > b.score;
	});

	int box_num = input.size();

	std::vector<int> merged(box_num, 0);

	for (int i = 0; i < box_num; i++)
	{
		if (merged[i])
			continue;

		output.push_back(input[i]);

		float h0 = input[i].y2 - input[i].y1 + 1;
		float w0 = input[i].x2 - input[i].x1 + 1;

		float area0 = h0 * w0;


		for (int j = i + 1; j < box_num; j++)
		{
			if (merged[j])
				continue;

			float inner_x0 = input[i].x1 > input[j].x1 ? input[i].x1 : input[j].x1;//std::max(input[i].x1, input[j].x1);
			float inner_y0 = input[i].y1 > input[j].y1 ? input[i].y1 : input[j].y1;

			float inner_x1 = input[i].x2 < input[j].x2 ? input[i].x2 : input[j].x2;  //bug fixed ,sorry
			float inner_y1 = input[i].y2 < input[j].y2 ? input[i].y2 : input[j].y2;

			float inner_h = inner_y1 - inner_y0 + 1;
			float inner_w = inner_x1 - inner_x0 + 1;


			if (inner_h <= 0 || inner_w <= 0)
				continue;

			float inner_area = inner_h * inner_w;

			float h1 = input[j].y2 - input[j].y1 + 1;
			float w1 = input[j].x2 - input[j].x1 + 1;

			float area1 = h1 * w1;

			float score;

			score = inner_area / (area0 + area1 - inner_area);

			if (score > nmsthreshold)
				merged[j] = 1;
		}

	}
}

void nms_cls(std::vector<objinfo> &inputs, std::vector<objinfo> &outputs, float nmsThresh, int clsnum){
    for(int i=0; i<clsnum; ++i){
        std::vector<objinfo> inboxes, outboxes;
        for (int k=0; k<inputs.size(); ++k){
            if (inputs[k].classNum == i){
                inboxes.push_back(inputs[k]);
            }
        }
		// printf("go nms with cls==%d~~~~~~~\n",i);
        nms1(inboxes, outboxes, nmsThresh, 1);
		// printf("out nms with cls==%d~~~~~~~\n",i);
        outputs.insert(outputs.end(), outboxes.begin(), outboxes.end());
        // outputs.insert(outputs.end(), inboxes.begin(), inboxes.end());
    }
}

int nhwc2nchw(char* src, char* dst, int width, int height, int depth, bool rbswap){
    if(depth != 3){
        printf("function only support for channels=3\n");
        return 1;
    }
    int wstride = width * depth;
    int mapsize = width * height;
	char* r = nullptr;
	char* g = nullptr;
	char* b = nullptr;
	if(rbswap){
		b = dst;
		g = dst + mapsize;
		r = dst + mapsize*2;
	}else{
		r = dst;
		g = dst + mapsize;
		b = dst + mapsize*2;
	}
    int i,j;
    for(i=0; i<height; i++){
        for(j=0; j<width; j++){
            *r = src[i*wstride + 3*j];
            *g = src[i*wstride + 3*j + 1];
            *b = src[i*wstride + 3*j + 2];
            r++;
            g++;
            b++;
        }
    }
    return 0;
};


int loadbin2float(const char* binname, float* dst, int datalen){
	std::ifstream file(binname, std::ios::in|std::ios::binary|std::ios::ate);
	if(file.is_open())
	{
		int fdatalen = file.tellg();
		if(fdatalen/sizeof(float) != datalen)
		{
			printf("filesize=%d not match datalen =%d\n", fdatalen, datalen);
			file.close();
			return 1;
		}
	}
	file.seekg(0, std::ios::beg);

	file.read((char*)dst, datalen*sizeof(float));
	file.close();
	return 0;
};

int loadbin2uint8(const char* binname, unsigned char* dst, int datalen){
	std::ifstream file(binname, std::ios::in|std::ios::binary|std::ios::ate);
	if(file.is_open())
	{
		int fdatalen = file.tellg();
		if(fdatalen/sizeof(char) != datalen)
		{
			printf("filesize=%d not match datalen =%d\n", fdatalen, datalen);
			file.close();
			return 1;
		}
	}
	file.seekg(0, std::ios::beg);

	file.read((char*)dst, datalen*sizeof(unsigned char));
	file.close();
	return 0;
}

int savefloat2bin(float* src, int datalen, char* dstname){
	FILE* f = fopen(dstname, "wb");
	if(f==NULL){
		printf("file %s open failed\n", dstname);
		return 1;
	}
	fwrite((char*)src, sizeof(float), datalen, f);
	fclose(f);
	return 0;
};




std::vector<float> bboxes_iou(OBJTRACK objtracker, std::vector<objinfo> dets, int iouType){
	float eps = 1e-5;
	std::vector<float> ious;
	for(objinfo det : dets){
		float w1 = objtracker._x2 - objtracker._x1;
		float h1 = objtracker._y2 - objtracker._y1;
		float w2 = det.x2 - det.x1;
		float h2 = det.y2 - det.y1;
		float inter = std::max(float(0), std::min(objtracker._x2, det.x2)-std::max(objtracker._x1, det.x1)) * 
			std::max(float(0), std::min(objtracker._y2, det.y2)-std::max(objtracker._y1, det.y1));
		float unarea = 	w1 * h1 + w2 * h2 - inter + eps;
		float iou = inter / unarea;
		if(iouType > 0){
			//GIOU
			float cw = std::max(objtracker._x2, det.x2) - std::min(objtracker._x1, det.x1);
			float ch = std::max(objtracker._y2, det.y2) - std::min(objtracker._y1, det.y1);
			if(iouType > 1){
				// DIOU
				float c2 = cw*cw + ch*ch + eps;
				float rho2 = ((objtracker._x1+objtracker._x2 - det.x1 - det.x2)*(objtracker._x1+objtracker._x2 - det.x1 - det.x2) + 
				(objtracker._y1+objtracker._y2 - det.y1 - det.y2)*(objtracker._y1+objtracker._y2 - det.y1 - det.y2))*0.25;
				if(iouType > 2){
					// CIOU
					float v = (4/(M_PI*M_PI)) * pow(atan(w2/(h2+eps)) - atan(w1/(h1+eps)), 2);
					float alpha = v/(v-iou+(1+eps));
					ious.push_back((iou - (rho2/c2 + v*alpha)));
				}
				ious.push_back(iou - rho2/c2);
			}
			float c_area = cw*ch + eps;
			ious.push_back(iou - (c_area-unarea)/c_area);
		}
		ious.push_back(iou);	
	}
	return ious;
};


int bboxes_iou(std::shared_ptr<OBJTRACK> objtracker, std::vector<objinfo> dets, std::vector<float> &ious, int iouType){
	float eps = 1e-5;
	// std::vector<float> ious;
	float objx1 = objtracker->last_x1 + objtracker->m_vx;
	float objx2 = objtracker->last_x2 + objtracker->m_vx;
	float objy1 = objtracker->last_y1 + objtracker->m_vy;
	float objy2 = objtracker->last_y2 + objtracker->m_vy;
	for(objinfo det : dets){
		float w1 = objx2 - objx1;
		float h1 = objy2 - objy1;
		float w2 = det.x2 - det.x1;
		float h2 = det.y2 - det.y1;

		float inter = std::max(float(0), std::min(objx2, det.x2)-std::max(objx1, det.x1)) * 
			std::max(float(0), std::min(objy2, det.y2)-std::max(objy1, det.y1));
		float unarea = 	w1 * h1 + w2 * h2 - inter + eps;
		float iou = inter / unarea;
		if(iouType > 0){
			//GIOU
			float cw = std::max(objx2, det.x2) - std::min(objx1, det.x1);
			float ch = std::max(objy2, det.y2) - std::min(objy1, det.y1);
			if(iouType > 1){
				// DIOU
				float c2 = cw*cw + ch*ch + eps;
				float rho2 = ((objx1 + objx2 - det.x1 - det.x2)*(objx1 + objx2 - det.x1 - det.x2) + 
				(objy1 + objy2 - det.y1 - det.y2)*(objy1 + objy2 - det.y1 - det.y2))*0.25;
				if(iouType > 2){
					// CIOU
					float v = (4/(M_PI*M_PI)) * pow(atan(w2/(h2+eps)) - atan(w1/(h1+eps)), 2);
					float alpha = v/(v-iou+(1+eps));
					ious.push_back((iou - (rho2/c2 + v*alpha)));
				}
				ious.push_back(iou - rho2/c2);
			}
			float c_area = cw*ch + eps;
			ious.push_back(iou - (c_area-unarea)/c_area);
		}
		ious.push_back(iou);	
	}
	return 0;
};

//
//int bboxes_dists(OBJTRACK objtracker, std::vector<objinfo> dets, std::vector<float>& distances, int iouType)
//{
//	float eps = 1e-5;
//	// std::vector<float> ious;
//	float objx1 = objtracker.last_x1;
//	float objx2 = objtracker.last_x2;
//	float objy1 = objtracker.last_y1;
//	float objy2 = objtracker.last_y2;
//
//	float objcenterx = (objx1 + objx2) / 2;
//	float objcentery = (objy1 + objy2) / 2;
//
//	for (int i = 0; i < dets.size(); i++)
//	{
//		float boxx1 = dets[i].x1;
//		float boxy1 = dets[i].x2;
//		float boxx2 = dets[i].y1;
//		float boxy2 = dets[i].y2;
//
//		float boxcenterx = (boxx1 + boxx2) / 2;
//		float boxcentery = (boxy1 + boxy2) / 2;
//
//		float dist = sqrt(pow(objcenterx - boxcenterx, 2) + pow(objcentery - boxcentery, 2));
//
//		distances.push_back(dist);
//	}
//	return 0;
//};

int bboxes_dists(std::shared_ptr<OBJTRACK> objtracker, std::vector<objinfo> dets, std::vector<float>& distances, int iouType) {
	float eps = 1e-5;
	// std::vector<float> ious;
	float objx1 = objtracker->last_x1 + objtracker->m_vx;
	float objx2 = objtracker->last_x2 + objtracker->m_vx;
	float objy1 = objtracker->last_y1 + objtracker->m_vy;
	float objy2 = objtracker->last_y2 + objtracker->m_vy;
	float fcenterX = (objx2 + objx1) / 2;
	float fcenterY = (objy2 + objy1) / 2;

	for (objinfo det : dets) {

		float fObjcenterX = (det.x2 + det.x1) / 2;
		float fObjcenterY = (det.y2 + det.y1) / 2;
		float dist = std::sqrt(std::pow(fcenterX - fObjcenterX, 2) + std::pow((fcenterY - fObjcenterY), 2));
		distances.push_back(dist);
	}
	return 0;
}


int find_arg_max(std::vector<float> datas, int &idx, float &Maxdata){
	if(datas.size() == 0){
		return 1;
	}
	idx = 0;
	Maxdata = 0; //__FLT_MIN__
	for(int i=0; i<datas.size(); i++){
		if(datas[i] > Maxdata){
			idx = i;
			Maxdata = datas[i];
		}
	}
	return 0;
};

int find_arg_min(std::vector<float> datas, int& idx, float& Mindata) {
	if (datas.size() == 0) {
		return 1;
	}
	idx = 0;
	Mindata = 1e9; //__FLT_MIN__
	for (int i = 0; i < datas.size(); i++) {
		if (datas[i] < Mindata) {
			idx = i;
			Mindata = datas[i];
		}
	}
	return 0;
};


std::vector<int> get_color(int cls){
    cls += 1;
    int stride = 200;
    std::vector<int> clr = {0,0,0};
    for(int k=0; k<cls; k++){
        int i = k % 3;
        clr[i] += stride;
    }
    for(int k = 0; k< 3; k++){
        clr[k] = clr[k] %255;
    }
    return clr;
}

#ifdef WITH_OPENCV

void swapYUV_I420toNV12(unsigned char* i420bytes, unsigned char* nv12bytes, int width, int height)
{
    int nLenY = width * height;
    int nLenU = nLenY / 4;

    memcpy(nv12bytes, i420bytes, width * height);

    for (int i = 0; i < nLenU; i++)
    {
        nv12bytes[nLenY + 2 * i] = i420bytes[nLenY + i];                    // U
        nv12bytes[nLenY + 2 * i + 1] = i420bytes[nLenY + nLenU + i];        // V
    }
}

void BGR2YUV_nv12(cv::Mat src, cv::Mat &dst)
{
    int w_img = src.cols;
    int h_img = src.rows;
    dst = cv::Mat(h_img*1.5, w_img, CV_8UC1, cv::Scalar(0));
    cv::Mat src_YUV_I420(h_img*1.5, w_img, CV_8UC1, cv::Scalar(0));  //YUV_I420
    cvtColor(src, src_YUV_I420, CV_BGR2YUV_I420);
    swapYUV_I420toNV12(src_YUV_I420.data, dst.data, w_img, h_img);
}

#endif


int pad_img(ImgMat img, int padleft, int padright, int padtop, int padbot, unsigned char *dst, std::vector<int> values){
	int imgW = img.width;
	int imgH = img.height;
	int imgC = img.depth;

	int dstW = imgW + padleft + padright;
	int dstH = imgH + padtop + padbot;

	unsigned char* data = img.data;

	int wstride = imgW*imgC;
	int dstwstride = dstW*imgC;

	for(int j=0; j<padtop; j++){
		// printf("padding top\n");
		for(int i=0; i<dstW; i++){
			for(int k=0; k<imgC; k++){
				dst[j*dstwstride + i*imgC + k] = values.at(k);
			}
		}
	}
	for(int j=padtop; j<(padtop+imgH); j++){
		for(int i=0; i<padleft; i++){
			for(int k=0; k<imgC; k++){
				dst[j*dstwstride + i*imgC + k] = values.at(k);
			}
		}
		// for(int i=padleft; i<(padleft+imgW); i++){
		unsigned char* pdst = dst + j * dstwstride + padleft * imgC;
		unsigned char* psrc = data + (j - padtop) * wstride;
		memcpy(pdst, psrc, wstride*sizeof(unsigned char));
		// }
		for(int i=(padleft+imgW); i<dstW; i++){
			for(int k=0; k<imgC; k++){
				dst[j*dstwstride + i*imgC + k] = values.at(k);
			}
		}
	}
	for(int j=(padtop + imgH); j<dstH; j++){
		// printf("padding bot\n");
		for(int i=0; i<dstW; i++){
			for(int k=0; k<imgC; k++){
				dst[j*dstwstride + i*imgC + k] = values.at(k);
			}
		}
	}

	return 0;
};


//int add_TensorMat(std::shared_ptr<TensorMat> a, std::shared_ptr<TensorMat> b){
//	return 0;
//};

void crop_img(unsigned char* src, int srcw, int srch, int srcc, int x1, int y1, int x2, int y2, unsigned char* dst){
    int cropx1 = std::max(0, x1);
    int cropx2 = std::min(srcw, x2);
    int cropy1 = std::max(0, y1);
    int cropy2 = std::min(srch, y2);

    int cropw = cropx2 - cropx1;
    int croph = cropy2 - cropy1;
    int elenum = cropw*croph;
    if(elenum <=0){
        return ;
    }
    unsigned char* dstp = (unsigned char*)dst;
    unsigned char* srcp = (unsigned char*)src + cropy1 * srcw*srcc + cropx1 * srcc;
    for(int i=0; i<croph; i++){
        memcpy(dstp, srcp, cropw*srcc*sizeof(unsigned char));
        dstp += cropw*srcc;
        srcp += srcw * srcc;
    }
    return;
};

void crop_yuv(unsigned char* srcYUV, int srcw, int srch, int x1, int y1, int x2, int y2, unsigned char* dstYUV){
	int cropx1 = std::max(0, x1);
    int cropx2 = std::min(srcw, x2);
    int cropy1 = std::max(0, y1);
    int cropy2 = std::min(srch, y2);

    int cropw = cropx2 - cropx1;
    int croph = cropy2 - cropy1;
    int elenum = cropw*croph;
    if(elenum <=0){
        return ;
    }

    //src point
	int srcWstride = srcw * sizeof(char);
    uint8_t* pY1 = srcYUV + cropy1 * srcWstride + cropx1;//
    uint8_t* pUV1 = srcYUV + srch*srcWstride + cropy1*srcWstride/2 + cropx1;// int(boxes.at(0).y0*wstride/2 + boxes.at(0).x0 * sizeof(char));

    //target point
	int dstWstride = cropw * sizeof(char);
    uint8_t * pTarY = dstYUV;
    uint8_t * pTarUV = pTarY + cropw*croph*sizeof(char);
    
	// memcpy Y
    for(int i=0; i<croph; ++i){
        memcpy(pTarY + i*dstWstride, pY1 + i*srcWstride, dstWstride);
    }
    // memcpy UV
    for(int i=0; i<int(croph/2); ++i){
        memcpy(pTarUV + i*dstWstride, pUV1 + i*srcWstride, dstWstride);
    }
	return ;
};



float seg_IOU(const objinfo &b1, const objinfo &b2)
{
	/*float x1 = max(b1.x - b1.w / 2, b2.x - b2.w / 2);
	float y1 = max(b1.y - b1.h / 2, b2.y - b2.h / 2);
	float x2 = min(b1.x + b1.w / 2, b2.x + b2.w / 2);
	float y2 = min(b1.y + b1.h / 2, b2.y + b2.h / 2);*/

	float x1 = std::max(b1.x1 , b2.x1);
	float y1 = std::max(b1.y1 , b2.y1);
	float x2 = std::min(b1.x2, b2.x2);
	float y2 = std::min(b1.y2,  b2.y2);

	float w = std::max(0.0f, x2 - x1 + 1);
	float h = std::max(0.0f, y2 - y1 + 1);
	float area = w * h;
	//return area / (b1.w * b1.h + b2.w * b2.h - area);
	return area / ((b2.x2 - b2.x1) * (b2.y2- b2.y1));
}

bool SortAreaScore(objinfo box1, objinfo box2)
{
	return box1.area > box2.area;
}

void seg_NMS(std::vector<objinfo> &boxes, float thr, std::vector<objinfo> &result)
{
	result.clear();
	std::sort(boxes.begin(), boxes.end(), SortAreaScore);  //按照置信度从大到小排序

	while (boxes.size() != 0) 
	{
		result.push_back(boxes[0]);
		size_t index = 1;
		while (boxes.size() > index) 
		{
			float iou = seg_IOU(boxes[0], boxes[index]);
			if (iou > thr && boxes[0].classNum == boxes[index].classNum)
		    //if (iou > 0)
			{
				boxes.erase(boxes.begin() + index);
				continue;
			}
			++index;
			//cout<<index<<endl;
		}
		boxes.erase(boxes.begin());
	}
}


void makeRect(int& x1, int& y1, int& x2, int& y2, int srcW, int srcH){
	int width = x2-x1;
	int height = y2-y1;
	int MaxSide = MAX(width, height);

	int cx = (x1+x2)*0.5;
	int cy = (y1+y2)*0.5;


	x1 = MAX(0, cx - 0.5*MaxSide);
	y1 = MAX(0, cy - 0.5*MaxSide);
	x2 = MIN(srcW-1, cx + 0.5*MaxSide);
	y2 = MIN(srcH-1, cy + 0.5*MaxSide);

	return;
};