//OpenCV
#include <opencv2/opencv.hpp>

#include "progress_thread.h"

//create a custom classifier to access the number of classes (OpenCV protected variable)
class RFClass : public CvRTrees{
public:
	int get_nclasses(){ return CvRTrees::nclasses;}
};

/// Perform classification of the ENVI file using the current RF classifier
std::vector< stim::image<unsigned char> > predict_rf(stim::envi* E, RFClass* RF, unsigned char* MASK = NULL){

	size_t nC = RF->get_nclasses();							//get the number of classes
	size_t X = E->header.samples;
	size_t Y = E->header.lines;
	size_t B = E->header.bands;
	size_t XY = E->header.samples * E->header.lines;

	size_t tP = 0;											//calculate the total number of pixels
	if(MASK){
		for(size_t xy = 0; xy < XY; xy++){
			if(MASK[xy]) tP++;
		}
	}
	else
		tP = X * Y;

	std::vector< stim::image<unsigned char> > C;					//create an array of mask images
	C.resize(nC);

	for(unsigned long long c = 0; c < nC; c++){						//for each class mask
		C[c] = stim::image<unsigned char>(X, Y, 1);					//allocate space for the mask
		memset(C[c].data(), 0, X * Y * sizeof(unsigned char));		//initialize all of the pixels to zero
	}

	cv::Mat classF(1, (int)B, CV_32FC1);							//allocate space for a single feature vector


	double progress = 0;											//initialize the progress bar variable
	std::thread t1(progressbar_thread, &progress);					//start the progress bar thread

	unsigned long long t = 0;
	for(unsigned long long p = 0; p < XY; p++){						//for each pixel
		if(!MASK || MASK[p] > 0){
			E->spectrum<float>((float*)classF.data, (int)p);		//fill the classF matrix with a single spectrum
			float c = RF->predict(classF);							//classify the feature vector
			C[(size_t)c].data()[p] = 255;							//write a white pixel to the appropriate class image
			t++;
			progress = (double)(t+1) / (double)(tP) * 100.0;		//update the progress bar variable
		}
	}
	t1.join();														//finish the progress bar thread

	return C;
}

//function for training a random forest classifier
void train_rf(RFClass* RF, cv::Mat &trainF, cv::Mat &trainR, int tree_depth, int num_trees, bool VERBOSE = false){

	//unsigned tP = trainF.cols;								//calculate the number of individual measurements

    CvRTParams  params( tree_depth, // max_depth,
                        1, // min_sample_count,
                        0.f, // regression_accuracy,
                        false, // use_surrogates,
                        16, // max_categories,
                        0, // priors,
                        false, // calc_var_importance,
                        0, // nactive_vars,
                        num_trees, // max_num_of_trees_in_the_forest,
                        0, // forest_accuracy,
                        CV_TERMCRIT_ITER // termcrit_type
                       );
	
	if(VERBOSE) std::cout<<"Starting OpenCV CvRT training algorithm...";
	RF->train(trainF, CV_ROW_SAMPLE, trainR, cv::Mat(), cv::Mat(), cv::Mat(), cv::Mat(), params );		//train the classifier
	if(VERBOSE) std::cout<<"done"<<std::endl;
}