// Copyright (C) 2018 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

/**
* \brief The entry point for the Inference Engine object_detection demo application
* \file object_detection_demo_yolov3_async/main.cpp
* \example object_detection_demo_yolov3_async/main.cpp
*/
#include <gflags/gflags.h>
#include <functional>
#include <iostream>
#include <fstream>
#include <random>
#include <memory>
#include <chrono>
#include <vector>
#include <string>
#include <algorithm>
#include <iterator>
#include <math.h>
using namespace std;
#include <opencv2/opencv.hpp> 
#include <inference_engine.hpp>
using namespace InferenceEngine;

#include <samples/ocv_common.hpp>
#include <samples/common.hpp>
#include <samples/slog.hpp>

using namespace gflags;

#include "tiny_yolov3_async.hpp"
#include <ext_list.hpp>

#include <ie_layers.h>

typedef std::chrono::high_resolution_clock Time;
typedef std::chrono::duration<double, std::ratio<1, 1000>> ms;
typedef std::chrono::duration<float> fsec;

bool ParseAndCheckCommandLine(int argc, char *argv[]) {
    // ---------------------------Parsing and validating the input arguments--------------------------------------
    ParseCommandLineNonHelpFlags(&argc, &argv, true);
    if (FLAGS_h) {
        showUsage();
        return false;
    }
    cout << "Parsing input parameters.\n"  ;

    if (FLAGS_i.empty()) {
        throw logic_error("Parameter -i is not set");
    }

    if (FLAGS_m.empty()) {
        throw logic_error("Parameter -m is not set");
    }
    return true;
}

static int EntryIndex(int side, int lcoords, int lclasses, int location, int entry) {
    int n = location / (side * side);
    int loc = location % (side * side);
    return n * side * side * (lcoords + lclasses + 1) + entry * side * side + loc;
}

struct DetectionObject {
    int xmin, ymin, xmax, ymax, class_id;
    float confidence;

    /*DetectionObject(double x, double y, double h, double w, int class_id, float confidence, float h_scale, float w_scale) {
        this->xmin = static_cast<int>((x - w * 0.5) * w_scale);
        this->ymin = static_cast<int>((y - h  * 0.5) * h_scale);
        this->xmax = static_cast<int>(this->xmin + w * w_scale);
        this->ymax = static_cast<int>(this->ymin + h * h_scale);
        this->class_id = class_id;
        this->confidence = confidence;
    }*/

    bool operator<(const DetectionObject &s2) const {
        return this->confidence < s2.confidence;
    }
};
void fp16_to_fp32_2(float* dest, const short* src, int length) { 
	unsigned int* dest_int = reinterpret_cast<unsigned int*>(dest);
	for (int i = 0; i < length; i++) {
		int s = (src[i] >> 15); // 
		int e = ((src[i] & 0x7fff) >> 10); // e - 0 : +-0 , 0x1f: (s== 0) +-infinite, (s != 0) NaN
		int f = src[i] & 0x3ff;
		if (e != 0 && e != 0x1f) {
			if (s)
				dest[i] = 0.0 - (1 + f * pow(2, -10)) * pow(2, e - 15);
			else
				dest[i] = (1 + f * pow(2, -10)) * pow(2, e - 15);

		}
		else if (0 == e) { //非规格化2e−14 
			if (0 == f) dest[i] = 0.0;
			else if (s)
				dest[i] = 0 - pow(2, -24) * f;
			else
				dest[i] = pow(2, -24) * f;
		}
		else {
			if (0 == f) dest_int[i] = 0xffffffff;
			else
				dest_int[i] = 0xff800000;
		}
	}
}
void fp16_to_fp32(float* dest, const short* src, int length) {
	unsigned int* dest_int = reinterpret_cast<unsigned int*>(dest);
	for (int i = 0; i < length; i++) {		
		int s = (src[i] >> 15); // 
		int e = ((src[i] & 0x7fff) >> 10); // e - 0 : +-0 , 0x1f: (s== 0) +-infinite, (s != 0) NaN
		int f = src[i] & 0x3ff;
		if (e != 0 && e != 0x1f) {
			e -= 15;
			e += 127;
			dest_int[i] = (s << 31) | (e << 23) | (f << 13);
		}
		else if (0 == e) { //非规格化2e−14 
			if (0 == f) dest[i] = 0.0;
			else if (s)
				dest[i] = 0 - pow(2, -24) * f;
			else
				dest[i] = pow(2, -24) * f;
		}
		else {
			if (0 == f) dest_int[i] = (s << 31) | 0x7f800000;
			else dest_int[i] = (s << 31) | 0x7f800000 | f;
		}
		
	}
}
double IntersectionOverUnion(const DetectionObject &box_1, const DetectionObject &box_2) {
    double width_of_overlap_area = fmin(box_1.xmax, box_2.xmax) - fmax(box_1.xmin, box_2.xmin);
    double height_of_overlap_area = fmin(box_1.ymax, box_2.ymax) - fmax(box_1.ymin, box_2.ymin);
    double area_of_overlap;
    if (width_of_overlap_area < 0 || height_of_overlap_area < 0)
        area_of_overlap = 0;
    else
        area_of_overlap = width_of_overlap_area * height_of_overlap_area;
    double box_1_area = (box_1.ymax - box_1.ymin)  * (box_1.xmax - box_1.xmin);
    double box_2_area = (box_2.ymax - box_2.ymin)  * (box_2.xmax - box_2.xmin);
    double area_of_union = box_1_area + box_2_area - area_of_overlap;
    return area_of_overlap / area_of_union;
}
inline double logistic_activate(double x) { return x; }
void ParseYOLOV3Output(const CNNLayerPtr& layer, const float* output_data, const SizeVector& dims , int img_w, int img_h, 
	double threshold, vector<DetectionObject> &objects) {
    // --------------------------- Validating output parameters -------------------------------------

	const int out_blob_h = dims[2];
	const int out_blob_w = dims[3];
    
	if (layer->type != "RegionYolo")
		throw runtime_error("Invalid output type: " + layer->type + ". RegionYolo expected"); 
	if (out_blob_h != out_blob_w)
		throw runtime_error("Invalid size of output " + layer->name +
			" It should be in NCHW layout and H should be equal to W. Current H = " + to_string(out_blob_h) +
			", current W = " + to_string(out_blob_h));
	// --------------------------- Extracting layer parameters -------------------------------------
	std::vector<float> anchors = { 10,14,23,27,37,58,81,82,135,169,344,319 };
	try { anchors = layer->GetParamAsFloats("anchors"); }
	catch (...) {}

	cout << "Parsing Yolo results ....\n";
	int num = 3;
	try { num = layer->GetParamAsInts("mask").size(); }
	catch (...) {}
	int coords = 4;
	coords = layer->GetParamAsInt("coords");
	int classes = 1;
	classes = layer->GetParamAsInt("classes");
	int channels = num * (1 + classes + coords);
	if (channels != dims[1]) {
		throw runtime_error("Invalid size of channels " + layer->name +
			" It should be " + to_string(channels) + ", Current channels = " + to_string(dims[1])  );
	}

	int anchor_offset = (13 == out_blob_h) ? 6 : 0;
	int map_size = out_blob_h * out_blob_w; 
	double fw = 1.0 / out_blob_w;
	double fh = 1.0 / out_blob_h;
	DetectionObject obj;
	int anchor_channels = coords + classes + 1;
	for (int i = 0; i < map_size; i++) {
		int row = i / out_blob_w;
		int col = i % out_blob_w;
		for (int a = 0, c_offset = 0; a < num;
			a++, c_offset += anchor_channels) {
			int x_index = anchor_channels * map_size + i;
			int y_index = x_index + map_size;
			int w_index = y_index + map_size;
			int h_index = w_index + map_size;
			int conf_index = h_index + map_size;
			obj.confidence = logistic_activate(output_data[conf_index]);
			if (obj.confidence < threshold) continue;

			//cool. we get predictions 
			double x = (col + logistic_activate(output_data[x_index])) * fw;
			double y = (row + logistic_activate(output_data[y_index])) * fh;
			double h = exp(logistic_activate(output_data[h_index])) *
				anchors[anchor_offset + 2 * a + 1] / 416.0;
			double w = exp(logistic_activate(output_data[w_index]))
				* anchors[anchor_offset + 2 * a] / 416.0;
			int cls_index = conf_index + map_size;
			obj.xmin = (x - 0.5 * w) * img_w;
			if (obj.xmin < 0) obj.xmin = 0;
			obj.xmax = (x + 0.5 * w) * img_w;
			if (obj.xmax >= img_w)
				obj.xmax = img_w - 1;
			obj.ymin = (y - 0.5 * h) * img_h;
			if (obj.ymin < 0) obj.ymin = 0;
			obj.ymax = (y + 0.5 * h) * img_h;
			if (obj.ymax >= img_h) obj.ymax = img_h - 1;
			cout << "Anchor " << a << " found object(" << obj.xmin << "," << obj.xmax <<
				"," << obj.ymin << "," << obj.ymax << ") in cell (" << col << "," << row << ")\n";
			obj.class_id = -1;
			for (int c = 0; c < classes; c++) {
				double cls_probability = logistic_activate(output_data[cls_index]);
				cls_index += map_size;
				if (cls_probability * obj.confidence < threshold) continue;
				obj.class_id = c;
			}
			if (obj.class_id != -1) {
				objects.push_back(obj);
			}
		}
	}
     
}

int main(int argc, char *argv[]) {
    try {
        /** This demo covers a certain topology and cannot be generalized for any object detection **/
        cout << "InferenceEngine: " << GetInferenceEngineVersion() << endl;

        // ------------------------------ Parsing and validating the input arguments ---------------------------------
        if (!ParseAndCheckCommandLine(argc, argv)) {
            return 0;
        } 
		auto t0 = Time::now();
#if 0     
		cv::Mat img = cv::imread(FLAGS_i.c_str(), cv::IMREAD_COLOR); 
		if (img.empty()) {
			cerr << "Read image `" << FLAGS_i << "` failed! \n";
			return -1;
		}

		const size_t width = img.size().width;
		const size_t height = img.size().height;
		float* f32_data = new float[IMAGE_ELEMENTS];
 
		cv::Mat temp = cv::Mat::zeros(416, 416, CV_8UC3);
		cv::resize(img, temp, temp.size()); 

		uchar* u8_data = temp.data ;	
	  
		
		int n = 0; 
		for (int h = 0; h < 416; h++) {
			for (int w = 0; w < 416; w++) {
				for (int c = 0; c < 3; c++, n++) {
					int m = c * 173056 + h * 416 + w;
					f32_data[m] = u8_data[n] ;
				}
			}
		}   
#endif
		       
        InferencePlugin plugin = PluginDispatcher({""}).getPluginByDevice("MYRIAD");
		//plugin.AddExtension(std::make_shared<Extensions::Cpu::CpuExtensions>());
        printPluginVersion(plugin, cout);

        /** Per-layer metrics **/
        if (FLAGS_pc) {
            plugin.SetConfig({ { PluginConfigParams::KEY_PERF_COUNT, PluginConfigParams::YES } });
        }
        // -----------------------------------------------------------------------------------------------------

        // --------------- 2. Reading the IR generated by the Model Optimizer (.xml and .bin files) ------------
		
		t0 = Time::now();
        CNNNetReader netReader; 

        netReader.ReadNetwork(FLAGS_m);   
		
        netReader.ReadWeights(fileNameNoExt(FLAGS_m) + ".bin");
		auto t1 = Time::now(); 
		fsec fs = t1 - t0;
		double dur = chrono::duration_cast<ms>(fs).count();


		std::cout << " Network and weights loaded in " << fixed << setprecision(2) << dur << "ms.\n";

		CNNNetwork network = netReader.getNetwork(); 
		
        
        vector<string> labels;
		ifstream f("coco.names");
		if (f.is_open()) {
			char label[100];
			while (f.getline(label, 100)) {
				labels.push_back(label);
			}
			f.close();
		}
		cv::Mat orig_image = cv::imread(FLAGS_i.c_str());
		if (orig_image.empty()) {
			slog::warn << "Image " << FLAGS_i << " cannot be read!" << slog::endl;
			return 0;
		}
		int img_w = orig_image.cols;
		int img_h = orig_image.rows;
		


        // -----------------------------------------------------------------------------------------------------

        /** TINY_YOLOV3-based network should have one input and 2 outputs **/
        // --------------------------- 3. Configuring input and output -----------------------------------------
        // --------------------------------- Preparing input blobs ---------------------------------------------
 
        InputsDataMap inputs_info(network.getInputsInfo());
        if (inputs_info.size() != 1) {
            throw logic_error("This demo accepts networks that have only one input");
        }
        InputInfo::Ptr& input = inputs_info.begin()->second;

        string inputName = inputs_info.begin()->first;
		input->setPrecision(Precision::U8); 
        


        // --------------------------------- Preparing output blobs -------------------------------------------
 
        OutputsDataMap outputs_info(network.getOutputsInfo());

		for (auto& output : outputs_info) { 
			output.second->setPrecision(Precision::FP32);
		}
         
		t0 = Time::now();
        // --------------------------- 4. Loading model to the plugin ------------------------------------------
       
        ExecutableNetwork executable_network = plugin.LoadNetwork(network, {});

        // --------------------------- 5. Creating infer request -----------------------------------------------
		InferRequest infer_request = executable_network.CreateInferRequest();

		Blob::Ptr network_input = infer_request.GetBlob(inputName);

		const SizeVector& in_dims = network_input->getTensorDesc().getDims();
		int net_w = in_dims[3];
		int net_h = in_dims[2];
		int net_c = in_dims[1];
		int image_size = net_h * net_w;  

		cv::Mat resized = cv::Mat::zeros(net_h, net_w, CV_8UC3);
		resize(orig_image, resized, resized.size());

		t1 = Time::now();

	
		unsigned char* data = static_cast<unsigned char*>(network_input->buffer());
		
		for (int i = 0; i < image_size; i++) { 
			for (size_t c = 0; c < net_c; c++) { 
				data[c * image_size + i] = resized.data[i * net_c + c];
			}
		}
	 

        // --------------------------- 6. Doing inference ------------------------------------------------------
       
		dur = chrono::duration_cast<ms>(t1 - t0).count();
		std::cout << "Model loaded to the plugin in " << fixed << setprecision(2) << dur <<  "ms.\nStart inference...\n";
			
		t0 = Time::now();
		infer_request.Infer();		
		t1 = Time::now();
		dur = chrono::duration_cast<ms>(t1 - t0).count();
		std::cout << "Inference time: " << fixed << setprecision(2) << dur << " ms\n";
		vector<DetectionObject> objects;

		for (auto& output : outputs_info) { 
			const Blob::Ptr output_blob = infer_request.GetBlob(output.first);
			const float* detection = static_cast<PrecisionTrait<Precision::FP32>::value_type*>(output_blob->buffer());
			const SizeVector& out_dims = output_blob->getTensorDesc().getDims();
			ParseYOLOV3Output(network.getLayerByName(output.first.c_str()),detection,out_dims, img_w,img_h, FLAGS_t,objects);
		}
		std::sort(objects.begin(), objects.end());
		for (int i = 0; i < objects.size(); ++i) {
			if (objects[i].confidence == 0) continue;
			for (int j = i + 1; j < objects.size(); ++j)
				if (IntersectionOverUnion(objects[i], objects[j]) >= 0.3)
					objects[j].confidence = 0;
		}
		char text[100];
		for (DetectionObject &object : objects) {
			if (object.confidence < FLAGS_t) continue;
			int label = object.class_id;
			float confidence = object.confidence;
			if (FLAGS_r) {
				std::cout << "[" << label << "] element, prob = " << confidence <<
					"    (" << object.xmin << "," << object.ymin << ")-(" << object.xmax << "," << object.ymax << ")"
					<< ((confidence > FLAGS_t) ? " WILL BE RENDERED!" : "") << endl;
			}
 
			if (confidence > FLAGS_t) {
				/** Drawing only objects when >confidence_threshold probability **/
				if (object.class_id <  labels.size()) {
					sprintf(text, "%s : %.2f%%", labels[object.class_id].c_str(), object.confidence* 100.0);
				}
				else {
					sprintf(text, "label #%d : %.2f%%", label, object.confidence * 100.0);
				}
				cv::putText(orig_image, text, cv::Point2i(object.xmin, object.ymin - 5),
					cv::FONT_HERSHEY_COMPLEX_SMALL, 1, cv::Scalar(0, 0, 255));
				cv::rectangle(orig_image, cv::Point2i(object.xmin, object.ymin),
					cv::Point2i(object.xmax, object.ymax), cv::Scalar(0, 0, 255));
			}
		} 
		if (FLAGS_pc) {
			printPerformanceCounts(infer_request, std::cout);
		} 
		vector<int> params;
		params.push_back(cv::IMWRITE_JPEG_QUALITY);
		params.push_back(90);
		imwrite("predictions.jpg", orig_image, params);
            
    }
    catch (const exception& error) {
        cerr << "[ ERROR ] " << error.what() << endl;
        return 1;
    }
    catch (...) {
        cerr << "[ ERROR ] Unknown/internal exception happened.\n" ;
        return 1;
    }
	
    cout << "Execution successful,\n" ;
    return 0;
}