#ifndef		TORCHDEFINE_H_
#define		TORCHDEFINE_H_

#include <stdio.h>
#include "malloc.h"

#define LENGTH 256
#define MAX_GPU_COUNT 8
#define CLS_MODEL_BASE (10)
#define SEG_MODEL_BASE (10)
#define DET_MODEL_BASE (0)
#define	PER_MODEL_MAX_COUNT	(10)


enum DNN_MODEL_TYPE
{
	DNN_MODEL_TYPE_CLS=0,
	DNN_MODEL_TYPE_DET=1,
	DNN_MODEL_TYPE_SEG=2,
	DNN_MODEL_TYPE_DET_YOLO = 3,
	DNN_MODEL_TYPE_SEG_FILTER = 4,
};

enum DNN_MODEL_FILE_TYPE
{
	DNN_MODEL_CAFFE=0,
	DNN_MODEL_TORCH=1,
	DNN_MODEL_TENSORRT=2,
	DNN_MODEL_ONNX=3,
};


#ifndef _TrainInfo_
#define _TrainInfo_
struct TrainInfo {
	float	loss;
	float	accuracy;//用户指定iou阈值，置信度阈值下的Precison * Recall
	int		record_index; //训练执行的轮数
	bool	is_detection; //是否是训练定位模型
	float	mAP;//coco 计算方式下的 map@0.5:0.95
	float	mAP50;//coco 计算方式下的 map@0.5
	float	mAP75;//coco 计算方式下的 map@0.75
	float	mIOU; // 检出定位的与目标的平均iou
	float	accuracy_bg;
	float	accuracy_fg;
	float	accuracy_objectness;
	bool	is_multilabel;
	float	accuracy_fg_all;
};
#endif

#ifndef _DetResult_
#define _DetResult_
struct DetResult
{
	int image_id;
	int label;
	char name[16];
	float score;
	int left, top, height, width;
};
#endif

#ifndef _GpuDeviceInfo_
#define _GpuDeviceInfo_
struct GpuDeviceInfo
{
	int	 gpu_count;
	char gpu_name[MAX_GPU_COUNT][LENGTH];
	int	 gpu_memory_size[MAX_GPU_COUNT];//单位MB
	int  gpu_sharedMemPerBlock[MAX_GPU_COUNT];//单位KB
	int  gpu_maxThreadsPerBlock[MAX_GPU_COUNT];
};
#endif

#ifndef _YYDLL_ERRORMSG_
#define _YYDLL_ERRORMSG_
struct YYErrorMsg
{
	char	strErrorMsg[1024];
};
#endif

#ifndef _DnnInferInfo_
#define _DnnInferInfo_
struct DnnInferInfo
{
	char	strModelFile[256];
	int		iBatchSize;
	int		iInferMode;
	int		iDeviceIndex;
};
#endif

#ifndef _Shape_
#define _Shape_
struct Shape
{
	Shape() { N = 1; C = 1; H = 1; W = 1; }
	Shape(unsigned int n, unsigned int c, unsigned int h, unsigned int w)
	{
		N = n; C = c; H = h; W = w;
	}
	unsigned int count()
	{
		return N * C * H * W;
	}
	unsigned int N;
	unsigned int C;
	unsigned int H;
	unsigned int W;
};
#endif // !_Shape_

#ifndef _SegResult_
struct SegResult
{
	unsigned char* pLabel;
	float* pProb;
	Shape  inferShape;
	int offsetX;
	int offsetY;
	double dScale;
};
#endif

#ifndef _Datum_
#define _Datum_
struct Datum
{
	Datum()
	{
		data = NULL;
		data_item_size = 4;
		alloc = false;
	}

	Datum(int iDataItemSize)
	{
		data = NULL;
		data_item_size = iDataItemSize;
		alloc = false;
	}

	Datum(void* pData, Shape tshape, int itemSize)
	{
		data = pData;
		this->shape = tshape;
		data_item_size = itemSize;
		alloc = false;
	}

	int		GetDataItemSize()
	{
		return data_item_size;
	}
	void*	Getdata()
	{
		if (data != NULL)
		{
			return data;
		}
		else
		{
			unsigned int size = shape.count();
			data = malloc(size * data_item_size);
			alloc = true;
			return data;
		}
	}

	void	SetDataItemSize(int iDataItemSize)
	{
		if (data != NULL && iDataItemSize > data_item_size )
		{
			free(data);
			data = NULL;
			alloc = false;
		}
		data_item_size = iDataItemSize;

		if (NULL == data)
		{
			unsigned int size = shape.count();
			data = malloc(size * data_item_size);
			alloc = true;
		}
	}

	void Reshape(Shape s)
	{
		if (data != NULL && s.count() > shape.count())
		{
			free(data);
			data = NULL;
			alloc = false;
		}
		shape = s;
		if (NULL == data)
		{
			unsigned int size = shape.count();
			data = malloc(size * data_item_size);
			alloc = true;
		}
	}
	~Datum()
	{
		if (data != NULL)
		{
			if (alloc)
			{
				free(data);
				data = NULL;
			}
		}
	}
	Shape	shape;
	void*	data;
	int		data_item_size;
	bool	alloc;
};
#endif // !_Datum_




#ifndef _TrainingParameter_
#define _TrainingParameter_
struct TrainingParameter 
{
	//enhance relate
	bool flip;
	bool blur;
	bool hsv;
	bool transform;
	bool crop;
	//train information
	int train_size;//100000
	int kLogInterval;//100
	int save_model_iterval;//1000
	int batch_size;
	//model relate need to write to file
	//如果参数不相同将进行重新训练
	int use_model_type;//0: small model; 1: medium model; 2: big model
	int channel;
	int height;
	int width;
	
	char data_file_root[LENGTH];
	char val_data_file_root[LENGTH];
	//添加容易出错的图像，每个batch保证一定比率添加到训练中
	char harddata_file_root[LENGTH];
	float input_scale;//1

	float learn_rate;//0.01
	
	//save relate
	char pretrained_model_path[LENGTH];
	char save_model_path[LENGTH];

	int		  train_iter;		//训练迭代轮数
	bool	  save_interval_flag;	//是否保存训练过程中文件
	bool	  vflip;			//是否使用上下翻转
	float	  max_rotate_angle;	//使用旋转增强时最大旋转角度(+-范围)
	int		  loss_type;		//训练loss_type
	float	  image_scale;		//图像训练前预缩放比例
	int		  infer_height;		//推理时模型输入尺寸高度
	int		  infer_width;		//推理时模型输入尺寸宽度
	double	  label_smooth;		//是否平滑类别标签，（即人为预估标记偏差）
	bool	  in_memory_train;	//是否在内存中训练
	int		  model_type;		//增加模型类别,0表示分类,1表示定位，2表示分割
	bool	  with_prob;		//推理时是否要输出概率
	double	  crop_min_scale;	//如果启用crop增强，最大的随机尺度
	double	  crop_max_scale;	//如果启用crop增强，最小的随机尺度

	/*------------------ for detect --------------------------------------*/
	char background_images_path[LENGTH];//加水印的背景图
	bool is_create_anchor;//
	float pos_overlap_low;//overlap大于这个值的被认为前景目标
	float neg_overlap_up;//overlap小于这个值的被认为背景目标，中间区域认为忽略的前景目标
	float gamma; //focal loss parameter
	float balance; //focal loss parameter
};
#endif

#ifndef _InferenceParameter_
#define _InferenceParameter_
struct InferenceParameterEx
{
	int use_model_type;
	int batch_size;
	int input_channel;
	int input_width;
	int input_height;
	int label_num;
	double input_scale;
	int model_type;	//增加模型类别,0表示分类,1表示定位，2表示分割
	float mean_values[4]; 
	float scale_value[4]; 
	int   input_count;
	char  input_name[3][64];
	int	  output_count;
	char  output_name[3][64];
	double image_scale;
};
#endif

#endif