#include "alg_jd.h"
#include <direct.h> //所需的库文件

extern SysConf g_sys_conf;

// Initialize the parameters
static float confThreshold = g_sys_conf.ConfThreshold*0.01; // Confidence threshold
static float nmsThreshold = 0.4;  // Non-maximum suppression threshold
static int inpWidth = 416;        // Width of network's input image
static int inpHeight = 416;       // Height of network's input image

static std::vector<std::string> classes;

bool AlgJd::init(QString model_path, QString model_name)
{
    // Load names of classes
    std::string classesFile;
    cv::String modelConfiguration;
    cv::String modelWeights;

    if (model_path.length() > 0 && model_name.length() > 0) {
        // 拼凑的模型文件路径
        modelWeights = model_path.toStdString() + "/" + model_name.toStdString();
        modelConfiguration = model_path.toStdString() + "/jd.cfg";
        classesFile = model_path.toStdString() + "/jd.names";
    }
    else {
        modelWeights = "D:/model/jd.weights";
        classesFile = "D:/model/jd.names";
        // Give the configuration and weight files for the model
        modelConfiguration = "D:/model/jd.cfg";
    }

    std::ifstream classNamesFile(classesFile.c_str());
    if (classNamesFile.is_open())
    {
        std::string className = "";
        while (std::getline(classNamesFile, className))
            classes.push_back(className);
    }
    else{
		return false;
    }

    // Load the network
    net = cv::dnn::readNetFromDarknet(modelConfiguration, modelWeights);

    net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
    net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
    
	//cv::Mat image = cv::imread("alg_jd.bmp");
    cv::Mat image;
    if (g_sys_conf.jpg_path.length() > 0) {
        image = cv::imread(g_sys_conf.jpg_path.toStdString());
    }
    else {
        image = cv::imread("D:/Release/alg_jd.bmp");
    }

	//识别一张图，测试模型是否正确，并且完成GPU数据加载
	if (!image.data)	return false;	//判断测试图片是否正常读取
	std::vector<std::pair<int, cv::Rect> > results;
	detect(image, image, results);
	if (results.size() > 0)
		return true;	//检测到目标，则初始化成功
	else
		return false;	//否则初始化失败
}

bool AlgJd::test_detect()
{
    cv::Mat m1;
    m1 = cv::Mat(544, 728, CV_8UC3, cv::Scalar(0, 0, 0));
	std::vector<std::pair<int, cv::Rect> > results;
    double t = (double)cv::getTickCount();
	detect(m1, m1, results);
    t = ((double)cv::getTickCount() - t) / cv::getTickFrequency();
    DEBUG("    test_detect time process:%f\n", t);
    return true;
}

bool AlgJd::test_detect_batcht(int shoot)
{
    std::vector<cv::Mat> vec_in;
    std::vector<cv::Mat> vec_out;
    std::vector<std::vector<std::pair<int, cv::Rect> > > vec_results;
    cv::Mat m1, m2, m3;
    m1 = cv::Mat(544, 728, CV_8UC3, cv::Scalar(0, 0, 0));
    m2 = cv::Mat(544, 728, CV_8UC3, cv::Scalar(0, 0, 0));
    m3 = cv::Mat(544, 728, CV_8UC3, cv::Scalar(0, 0, 0));

    double t = (double)cv::getTickCount();
    switch(shoot){
        case 1:{
            std::vector<std::pair<int, cv::Rect> > results;
            detect(m1, m1, results);
            break;
        }
        case 3:vec_in.push_back(m3);
        case 2:vec_in.push_back(m2);
        default:{
            vec_in.push_back(m1);
            detect_batch(vec_in, vec_out, vec_results);
        }
    }
    t = ((double)cv::getTickCount() - t) / cv::getTickFrequency();
    DEBUG("    test_detect_batcht time process:%f\n", t);
    return true;
}

int AlgJd::detect(cv::Mat& _frame, cv::Mat &out, std::vector<std::pair<int, cv::Rect> > &results)
{
    cv::Mat frame = _frame.clone();
    cv::Mat image_clone=frame.clone();
    // Process frames.
    // Create a 4D blob from a frame.
    cv::Mat blob;
    cv::dnn::blobFromImage(frame, blob, 1/255.0, cv::Size(inpWidth, inpHeight), cv::Scalar(0,0,0), true, false);

    //Sets the input to the network
    net.setInput(blob);


    // Runs the forward pass to get output of the output layers
    std::vector<cv::Mat> outs;
    net.forward(outs, getOutputsNames(net));


    // Remove the bounding boxes with low confidence
    post_process(frame, outs, results);


    // Write the frame with the detection boxes
    cv::Mat detectedFrame;
    frame.convertTo(detectedFrame, CV_8U);
    //show detectedFrame
    out=detectedFrame.clone();
	return results.size();
}

void AlgJd::analyse(cv::Mat vec_in, std::vector<std::pair<int, cv::Rect> >& vec_results)
{
    for (int i = 0; i < vec_results.size(); i++)
    {
        std::pair<int, cv::Rect> element;
        cv::Rect TempRect = vec_results[i].second;
        cv::Point position = cv::Point(TempRect.x + (TempRect.width / 2), TempRect.y + (TempRect.height / 2));
        cv::Mat topography = vec_in(TempRect);
        cv::imshow("analyse", topography);
        cv::waitKey(1);
    }
    
}

// Get the names of the output layers
std::vector<cv::String> AlgJd::getOutputsNames(const cv::dnn::Net& net)
{
    std::vector<cv::String> names;
    if (names.empty())
    {
        //Get the indices of the output layers, i.e. the layers with unconnected outputs
        std::vector<int> outLayers = net.getUnconnectedOutLayers();

        //get the names of all the layers in the network
        std::vector<cv::String> layersNames = net.getLayerNames();

        // Get the names of the output layers in names
        names.resize(outLayers.size());
        for (size_t i = 0; i < outLayers.size(); ++i)
        names[i] = layersNames[outLayers[i] - 1];
    }
    return names;
}

void AlgJd::post_process(cv::Mat& frame, std::vector<cv::Mat>& outs, std::vector<std::pair<int, cv::Rect> > &results)
{
    std::vector < std::vector<int>> classIds(classes.size());
    std::vector < std::vector<float>> confidences(classes.size());
    std::vector < std::vector<cv::Rect>> boxes(classes.size());
    for (size_t i = 0; i < outs.size(); ++i)
    {
        // Scan through all the bounding boxes output from the network and keep only the
        // ones with high confidence scores. Assign the box's class label as the class
        // with the highest score for the box.
        float* data = (float*)outs[i].data;
        for (int j = 0; j < outs[i].rows; ++j, data += outs[i].cols)
        {
            cv::Mat scores = outs[i].row(j).colRange(5, outs[i].cols);
            cv::Point classIdPoint;
            double confidence;
            // Get the value and location of the maximum score
            cv::minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);

            if (confidence > confThreshold)
            {
                int centerX = (int)(data[0] * frame.cols);
                int centerY = (int)(data[1] * frame.rows);
                int width = (int)(data[2] * frame.cols);
                int height = (int)(data[3] * frame.rows);
                int left = centerX - width / 2;
                int top = centerY - height / 2;

                classIds[classIdPoint.x].push_back(classIdPoint.x);
                confidences[classIdPoint.x].push_back((float)confidence);
                boxes[classIdPoint.x].push_back(cv::Rect(left, top, width, height));
            }
        }
    }


    // Perform non maximum suppression to eliminate redundant overlapping boxes with
    // lower confidences

    for (size_t i = 0; i < classes.size(); ++i)
    {
        std::vector<int> indices;
        cv::dnn::NMSBoxes(boxes[i], confidences[i], confThreshold, nmsThreshold, indices);
        for (size_t j = 0; j < indices.size(); ++j)
        {
            int idx = indices[j];
            cv::Rect box = boxes[i][idx];
            drawPred(classIds[i][idx], confidences[i][idx], box.x, box.y,
                    box.x + box.width, box.y + box.height, frame);
            results.push_back(std::make_pair(classIds[i][idx], box));
        }
    }
}

// Draw the predicted bounding box
void AlgJd::drawPred(int classId, float conf, int left, int top, int right, int bottom, cv::Mat& frame)
{
	cv::Scalar color;
	if(classId==0)
        color = cv::Scalar(0,255,0);
	else if(classId == 1)
		color = cv::Scalar(255,0,0);
	else
		color = cv::Scalar(0, 255, 0);
    //Draw a rectangle displaying the bounding box
    cv::rectangle(frame, cv::Point(left, top), cv::Point(right, bottom), color, 4);
   
    //Get the label for the class name and its confidence
    std::string label = cv::format("%.2f%%", (conf*100));///
    if (!classes.empty())
    {
        CV_Assert(classId < (int)classes.size());
        label = classes[classId] + ": " + label;
    }
    else
    {
        std::cout<<"classes is empty..."<<std::endl;
    }

    //Display the label at the top of the bounding box
    int baseLine;
    cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 1, 1, &baseLine);
    top = std::max(top, labelSize.height);
    cv::putText(frame, label, cv::Point(left, top-5), cv::FONT_HERSHEY_SIMPLEX, 0.8, color, 1);
}

void AlgJd::detect_batch(std::vector<cv::Mat>& vec_in, std::vector<cv::Mat> &vec_out, std::vector<std::vector<std::pair<int, cv::Rect> > > &vec_results)
{
	cv::Mat blobs;
    std::vector<cv::Mat> image;
    
	cv::dnn::blobFromImages(vec_in, blobs, 1 / 255.0, cv::Size(inpWidth, inpHeight), cv::Scalar(0, 0, 0), true, false);

	//Sets the input to the network
	net.setInput(blobs);

	// Runs the forward pass to get output of the output layers
	std::vector<cv::Mat> outs;
	net.forward(outs, getOutputsNames(net));

	for (int i = 0; i < vec_in.size(); i++)
	{
		image.push_back(vec_in[i].clone());
	}

	// Remove the bounding boxes with low confidence
	post_process_batch(image, outs, vec_results);


	// Write the frame with the detection boxes
	for (int i = 0; i < vec_in.size(); i++)
	{
		cv::Mat detectedFrame, out;
		image[i].convertTo(detectedFrame, CV_8U);
		out = detectedFrame.clone();
		vec_out.push_back(out);
	}
}

void AlgJd::post_process_batch(std::vector<cv::Mat>& vec_frame, std::vector<cv::Mat>& outs, std::vector<std::vector<std::pair<int, cv::Rect> > > &vec_results)
{
	int batch = vec_frame.size();
    double confidence;///
	for (int k = 0; k < batch; k++)
	{
        std::vector < std::vector<int>> classIds(classes.size());
        std::vector < std::vector<float>> confidences(classes.size());
        std::vector < std::vector<cv::Rect>> boxes(classes.size());

        //std::cout << "outs.size()\t" << outs.size() << std::endl;
		//std::cout << "Type\t" << outs[0].type() << std::endl;

		for (size_t i = 0; i < outs.size(); ++i)
		{
			//std::cout << "outs.dims\t" << outs[i].dims << std::endl;
			//std::cout << "outs[" << i << "].rows\t" << outs[i].size[0] << std::endl;
			//std::cout << "outs[" << i << "].cols\t" << outs[i].size[1] << std::endl;
			//std::cout << "outs[" << i << "].cols\t" << outs[i].size[2] << std::endl;
			// Scan through all the bounding boxes output from the network and keep only the
			// ones with high confidence scores. Assign the box's class label as the class
			// 
			cv::Mat m0(outs[i].size[1], outs[i].size[2], CV_32F, outs[i].data + outs[i].step[0] * k);
			// with the highest score for the box.
			float* data = (float*)m0.data;
			for (int j = 0; j < m0.rows; ++j, data += m0.cols)
			{
				cv::Mat scores = m0.row(j).colRange(5, m0.cols);
				cv::Point classIdPoint;
				//double confidence;
				// Get the value and location of the maximum score
				cv::minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);

				if (confidence > confThreshold)
				{
					int centerX = (int)(data[0] * vec_frame[k].cols);
					int centerY = (int)(data[1] * vec_frame[k].rows);
					int width = (int)(data[2] * vec_frame[k].cols);
					int height = (int)(data[3] * vec_frame[k].rows);
					int left = centerX - width / 2;
					int top = centerY - height / 2;

					classIds[classIdPoint.x].push_back(classIdPoint.x);
					confidences[classIdPoint.x].push_back((float)confidence);
					boxes[classIdPoint.x].push_back(cv::Rect(left, top, width, height));
				}
			}
		}
		std::vector<std::pair<int, cv::Rect> > results;
        // Perform non maximum suppression to eliminate redundant overlapping boxes with
        // lower confidences
        for (size_t i = 0; i < classes.size(); ++i)
        {
            std::vector<int> indices;
            cv::dnn::NMSBoxes(boxes[i], confidences[i], confThreshold, nmsThreshold, indices);
            for (size_t j = 0; j < indices.size(); ++j)
            {
                int idx = indices[j];
                cv::Rect box = boxes[i][idx];
                if (confidences[i][idx] > g_sys_conf.ConfThreshold * 0.01)///识别度低于阈值NG处理
                {
                    if (box.width > 15)
                    {//识别框宽度大于15显示识别，小于认为无胶点，NG处理
                        drawPred(classIds[i][idx], confidences[i][idx], box.x, box.y,
                        box.x + box.width, box.y + box.height, vec_frame[k]);
                        results.push_back(std::make_pair(classIds[i][idx], box));
                    }
                }
            }
		}
		vec_results.push_back(results);
	}
}
