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

#include "progress_thread.h"

double progress = 0.0;

//create a custom classifier to access the number of classes (OpenCV protected variable)
class BAYClass : public CvNormalBayesClassifier{
public:
	int get_nclasses(){
		return CvNormalBayesClassifier::cls_labels->cols;
	}
};

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

	size_t nC = BAY->get_nclasses();
	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
		C[c] = 0;													//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 = BAY->predict(classF);							//classify the feature vector
			if((size_t)c < C.size())								//if the class returned is valid
				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 bayesian classifier
void train_bay(BAYClass* BAY, cv::Mat &trainF, cv::Mat &trainR, bool VERBOSE = false){

	unsigned tP = trainF.cols;								//calculate the number of individual measurements
	
	if(VERBOSE) std::cout<<"Starting OpenCV CvBAY training algorithm...";
	BAY->train(trainF, trainR);								//train the classifier
	if(VERBOSE) std::cout<<"done"<<std::endl;
}