//OpenCV
#include <opencv2/opencv.hpp>
#include "progress_thread.h"

/// Calculates the squared Euclidean distance between two vectors
template<typename T>
T l2_norm_sq(T* v, size_t len){
	T sq_sum = 0;								//initialize the squared sum to zero
	for(size_t i = 0; i < len; i++)					//for each element of the vectors
		sq_sum += pow(v[i], 2);						//multiply the components
	return sq_sum;									//return the square root of the squared sum
}

template<typename T>
T euclidean_dist_sq(T* v0, T* v1, size_t len){
	T* diff = (T*)malloc(sizeof(T) * len);			//allocate space for the difference vector
	for(size_t i = 0; i < len; i++)
		diff[i] = v1[i] - v0[i];
	return l2_norm_sq(diff, len);
}

void train_kmeans(cv::Mat &centers, cv::Mat &F, int k, int attempts, int iters, double epsilon){
	cv::Mat labels;

	cv::kmeans(F, k, labels, cv::TermCriteria( CV_TERMCRIT_EPS+CV_TERMCRIT_ITER, iters, epsilon), attempts, cv::KMEANS_PP_CENTERS, centers);

	int* nk = (int*) malloc(sizeof(int) * k);			//create an array to store the number of pixels for each cluster
	memset(nk, 0, sizeof(int) * k);
	std::cout<<"size: "<<labels.size()<<std::endl;
	for(int n = 0; n < labels.size().height; n++)
		nk[labels.at<int>(n, 0)]++;

	for(int n = 0; n < k; n++)									//for each cluster
		std::cout<<"cluster "<<n<<" = "<<nk[n]<<std::endl;		//output the number of pixels in that cluster
	std::cout<<"total = "<<labels.rows<<std::endl;				//also output the total number of pixels clustered
}



//Predict a set of classes based on given centroid vectors
template<typename T>
std::vector< stim::image<unsigned char> > predict_centroids(stim::envi* E, std::vector< std::vector<T> > centroids, unsigned char* MASK = NULL){
	size_t nC = centroids.size();								//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(size_t 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
	}

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

	size_t t = 0;
	T* spectrum = (T*)malloc(sizeof(T) * B);			//allocate space to hold a spectrum
	double min_dist, new_dist;
	int nearest_centroid;
	for(unsigned long long p = 0; p < XY; p++){			//for each pixel
		if(!MASK || MASK[p] > 0){
			E->spectrum<T>(spectrum, p);						//get the spectrum at pixel p
			
			min_dist = euclidean_dist_sq<T>(centroids[0].data(), spectrum, B);		//set the minimum distance as the distance to the first centroid
			nearest_centroid = 0;												//set the nearest centroid to the first centroid
			for(size_t c = 1; c < nC; c++){										//for each centroid
				new_dist = euclidean_dist_sq<T>(centroids[c].data(), spectrum, B);	//calculate the distance to the current centroid
				if(new_dist < min_dist){										//if the current distance is smaller
					min_dist = new_dist;										//update the minimum distance
					nearest_centroid = (int)c;										//set the nearest centroid to the current centroid
				}
			}
			C[nearest_centroid].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;
}