#ifndef STIM_CUDA_KMEANS_CUH
#define STIM_CUDA_KMEANS_CUH

#include <stim/image/image.h>
//#include <cmath>
#include <stim/visualization/colormap.h>
//#include <stim/image/image_contour_detection.h>
#include <opencv2/opencv.hpp>
#include <iostream>

/// This function use cvkmeans to cluster given textons  

/// @param testons is a multi-channel image
/// @param K is the number of clusters

stim::image<float> kmeans(stim::image<float> textons, unsigned int K){

	unsigned int w = textons.width();              // get the width of picture
	unsigned int h = textons.height();             // get the height of picture
	unsigned int feature_n = textons.channels();   // get the spectrum of picture
	unsigned int N = w * h;						   // get the number of pixels

	float* sample1 = (float*) malloc(sizeof(float) * N * feature_n);  //allocate the space for textons

	//reallocate a multi-channel texton image to a single-channel image
	for(unsigned int c = 0; c < feature_n; c++){

		stim::image<float> temp;
		temp = textons.channel(c);

		for(unsigned int j = 0; j < N; j++){

			sample1[c + j * feature_n] = temp.data()[j];
		}
	}
	
	
	cv::Mat sample2(N, feature_n, CV_32F, sample1); //copy image to cv::mat

	//(optional) show the test result
	//imshow("sample2", sample2);


	cv::TermCriteria criteria(CV_TERMCRIT_EPS+CV_TERMCRIT_ITER, 10, 0.1);	// set stop-criteria for kmeans iteration
	cv::Mat labels(N, 1, CV_8U, cvScalarAll(0));							// allocate space for kmeans output
	cv::Mat centers;														// allocate space for kmeans output

	unsigned int test_times = 2;  // set the number of times of trying kmeans, it will return the best result

	cv::kmeans(sample2, K, labels, criteria, test_times, cv::KMEANS_PP_CENTERS, centers); // kmeans clustering

	//(optional) show the test result
	//imwrite( "data_output/labels_1D.bmp", labels);

	stim::image<float> texture(w, h, 1, 1);						// allocate space for texture
	 
	for(unsigned int i = 0; i < N; i++){                        // reshape the labels from iD array to image

		texture.data()[i] = labels.at<int>(i);
	
	}

	//texture.save("data_output/kmeans_test0924_2.bmp");

	//(optional) show the test result
	//stim::cpu2image(texture.data(), "data_output/kmeans_test.bmp", w, h, stim::cmBrewer);  

	return texture;

}

#endif