#ifndef __GENERAL_UTILS_H__
#define __GENERAL_UTILS_H__
#include <iostream>
#include <string>
#include <chrono>
#include <cfloat>
#include "NvInfer.h"
#include <cuda_runtime_api.h>


using namespace std;

// Simple implementation for the TRT logger.
class Logger : public nvinfer1::ILogger           
{
public:	
     void log(Severity severity, const char* msg) override
     {
         // suppress info-level messages
         //if (severity != Severity::kINFO)
         cout << msg << endl;
     }
 };
 
// Define a simple profiler call back
class RVLayerProfiler : public nvinfer1::IProfiler
{
private:	
	int m_index;
	float m_total;
	float m_min_time;
	float m_max_time;
	string m_min_layer_name;
	string m_max_layer_name;
	string m_profiler_name;
	bool m_print_stats_per_layer;
	
public:	
	RVLayerProfiler() : m_index(1), m_total(0.f), m_max_time(0.f), m_min_time(FLT_MAX), m_profiler_name("..."), m_print_stats_per_layer(true) {}
	~RVLayerProfiler() 
	{
		cout << "[RVLayerProfiler - " << m_profiler_name << "]: Total time for " << m_index - 1 << " layers: " << m_total << "ms" << endl;
		cout << "[RVLayerProfiler - " << m_profiler_name << "]: Shortest layer: [" << m_min_layer_name << "] ran for [" << m_min_time << "ms]" << endl;
		cout << "[RVLayerProfiler - " << m_profiler_name << "]: Longest  layer: [" << m_max_layer_name << "] ran for [" << m_max_time << "ms]" << endl;
		cout << "[RVLayerProfiler - " << m_profiler_name << "]: Average time  : [" << m_total / m_index << "ms]" << endl;
	}
	
	virtual void reportLayerTime (const char *layerName, float ms)
	{
//        if (0 != strcmp(layerName, "(Unnamed Layer* 0) [Convolution]"))
//            return;
		if (m_print_stats_per_layer)
			cout << "[RVLayerProfiler - " << m_profiler_name << "]: Layer [" << m_index << "]: [" << layerName << "]: " << ms << "ms" << endl;
		m_index++;
		m_total += ms;
		// Update the maximum time consuming layer.
		if (m_max_time < ms) 
		{
			m_max_time = ms;
			m_max_layer_name = string(layerName);
		}
		if (m_min_time > ms)
		{
			m_min_time = ms;
			m_min_layer_name = string(layerName);
		}
	}
	
	int GetLayerCount() { return m_index - 1; }
	void SetPrintStatsPerLayer(bool print_stats_per_layer) { m_print_stats_per_layer = print_stats_per_layer; }
	void SetName(const char *name)
	{
		m_profiler_name = string(name);
	}
};


/**
 * Following code taken from /usr/src/tensorrt/samples/common/common.h
 * 
 */
 
#define CHECK(status)                                          \
    do                                                         \
    {                                                          \
        auto ret = (status);                                   \
        if (ret != 0)                                          \
        {                                                      \
            std::cout << "Cuda failure: " << ret << std::endl; \
            abort();                                           \
        }                                                      \
    } while (0)


#endif	// __GENERAL_UTILS_H__