#ifndef MASK_MANIPULATION
#define MASK_MANIPULATION

#include <numeric>

#include <stim/image/image.h>

///This file contains functions designed to process, subsample, and manipulate masks images

void push_training_image(std::vector< stim::image<unsigned char> >& C, std::vector<size_t>& nP, size_t& tP, stim::image<unsigned char> I) {
	C.push_back(I);
	size_t npixels = I.nnz();
	nP.push_back(npixels);									//push the number of pixels onto the pixel array
	tP += npixels;											//add to the running total of pixels
//	nC++;
}

//re-calculates the total number of pixels in an image list
size_t calc_tP(std::vector< stim::image<unsigned char> >& C, std::vector< size_t > nP) {
	size_t tP = 0;
	for (size_t i = 0; i < C.size(); i++) {
		tP += nP[i];
	}
	return tP;
}

size_t calc_tP(std::vector< stim::image<unsigned char> >& C) {
	std::vector<size_t> nP(C.size());
	for (size_t ci = 0; ci < C.size(); ci++)
		nP[ci] = C[ci].nnz();
	return calc_tP(C, nP);
}

//load a mask given a mask filename
void AND_mask(std::vector< stim::image<unsigned char> >& C, std::vector<size_t>& nP, size_t& tP, std::string maskfile){
	stim::image<unsigned char> mask_image(maskfile);				//load the image
	stim::image<unsigned char> mono_image = mask_image.channel(0);	//retrieve the first channel

	if(C.size() == 0){
		push_training_image(C, nP, tP, mono_image);
	}
	else{
		size_t N = C[0].width() * C[0].height();									//precompute the number of pixels in the mask
		tP = 0;
		for(size_t xy = 0; xy < N; xy++){					//for each pixel in the mask
			if(!C.back().data()[xy] || !mono_image.data()[xy]){			//if both the current and new mask is true
				C.back().data()[xy] = 0;
			}
		}
		nP.back() = C.back().nnz();
	}
	calc_tP(C, nP);	
	std::cout<<"masked pixels: "<<tP<<std::endl;
}

void push_full_mask(std::vector< stim::image< unsigned char > >& C, size_t X, size_t Y) {
	stim::image<unsigned char> img(X, Y, 1);				//create an empty image
	img = 255;												//set all values to 255

	C.push_back(img);											//push the image onto the image array
}

void push_fully_masked_image(std::vector< stim::image< unsigned char > >& C, std::vector<size_t>& nP, size_t& tP, size_t X, size_t Y){
	
	push_full_mask(C, X, Y);

	size_t npixels = C[C.size() - 1].nnz();
	nP.push_back(npixels);									//push the number of pixels onto the pixel array
	tP += npixels;											//add to the running total of pixels
	//nC++;													//increment the number of classes
}

void subsample_masks(std::vector< stim::image<unsigned char> >&I, size_t n) {
	size_t tP = calc_tP(I);
	if (tP <= n) return;										//if there are less than n pixels already, nothing changes

	std::vector< stim::image<unsigned char> > new_C;		//create a new training image array
	new_C.resize(I.size());										//set the correct size of the array

	std::vector<size_t> idx(tP);							//create an array that will store the indices of each pixel
	size_t ip = 0;											//index into the index array

	size_t X = I[0].width();
	size_t Y = I[0].height();

	for (size_t c = 0; c < I.size(); c++) {							//for each training image
		new_C[c] = stim::image<unsigned char>(X, Y, 1);		//create a new image that matches the size of the original
		new_C[c] = 0;										//set all values in the image to zero
		std::vector<size_t> c_idx = I[c].sparse_idx();		//get the indices of each non-zero pixel

		for (size_t i = 0; i < c_idx.size(); i++) {			//for each index
			idx[ip + i] = c_idx[i] + c * X * Y;				//scale the index by the image number and store it in the index array
		}
		ip += c_idx.size();									//increment the global index
	}

	std::random_device rd;											//create a random number generation function
	std::mt19937 g(rd());
	std::shuffle(idx.begin(), idx.end(), g);						//shuffle the index values

	size_t c, ix;
	for (size_t i = 0; i < n; i++) {									//for each of the first n random pixels
		c = idx[i] / (X * Y);										//calculate the image that the current pixel is in
		ix = idx[i] - c * X * Y;									//calculate the 1D index of the current pixel
		new_C[c].data()[ix] = I[c].data()[ix];						//store the value of the pixel in the new C array
	}

	I = new_C;														//replace the class image array
}

/// This function picks a random set of N samples from among the training masks
void set_random_samples(std::vector< stim::image<unsigned char> >& C, std::vector<size_t>& nP, size_t& tP, size_t n){

	subsample_masks(C, n);

	tP = 0;															//reset the total number of pixels to zero
	for(size_t c = 0; c < C.size(); c++){									//for each class image
		nP[c] = C[c].nnz();											//re-calculate the number of pixels in each image
		tP += nP[c];												//re-calculate the total number of pixels
	}
}

/// Calculate a new image composed of N random pixels from the input image
stim::image<unsigned char> random_subset(stim::image<unsigned char> in, size_t n) {
	if (n >= in.nnz()) return in;									//if the number of pixels requested is less than the number of masked pixels, return the original image
	std::vector<size_t> idx = in.sparse_idx();					//get the indices of each non-zero pixel
	std::random_device rd;											//create a random number generation function
	std::mt19937 g(rd());
	std::shuffle(idx.begin(), idx.end(), g);					//shuffle the index values

	stim::image<unsigned char> out(in.width(), in.height(), 1);		//create an output image the same size as the input image
	out = 0;														//set all pixels to zero

	for (size_t i = 0; i < n; i++)								//go through each pixel index until the number of requested pixels is reached
		out.data()[idx[i]] = 255;								//set the pixel value at that location to TRUE

	return out;													//return the output image
}


/// This function picks a random set of N samples from among the training masks. The samples are selected to make the training
///	set as balanced as possible (equal numbers of training samples from each class).
void set_random_samples_balanced(std::vector< stim::image< unsigned char > >& C, size_t n, std::string output_mask) {

	//size_t tP = calc_tP(C, nP);
	//if (tP <= n) return;									//if there are less than n pixels already, nothing changes

	std::vector< stim::image<unsigned char> > new_C;		//create a new training image array
	new_C.resize(C.size());										//set the correct size of the array

	//calculate the number of pixels in each training image
	std::vector<size_t> num_pix_per_class(C.size());				//create a vector to store the number of pixels in each class
	for (size_t c = 0; c < C.size(); c++)							//for each class
		num_pix_per_class[c] = C[c].nnz();					//calculate the number of pixels

	std::vector<int> c_idx(num_pix_per_class.size());		//create a vector that will store indices into the sorted class array
	std::iota(c_idx.begin(), c_idx.end(), 0);				//fill the index array with a list of indices [0, c)

	//create a function that compares two values in the num_pix_per_class array and sorts the indices
	auto comparator = [&num_pix_per_class](int a, int b) { return num_pix_per_class[a] < num_pix_per_class[b]; };
	std::sort(c_idx.begin(), c_idx.end(), comparator);		//sort the indices, which can be used to look up the sorted version of the number of pixels per class

	size_t ppc = n / C.size();									//calculate the ideal number of pixels per class
	size_t pixels_used = 0;									//number of pixels that have already been accepted into the mask array

	for (size_t ci = 0; ci < C.size(); ci++) {					//for each class (starting with the smallest)
		ppc = (n - pixels_used) / (C.size() - ci);				//determine the number of pixels per class that have to be used to balance out the rest of the classes
		size_t pix = C[c_idx[ci]].nnz();					//get the number of pixels in this class
		new_C[c_idx[ci]] = random_subset(C[c_idx[ci]], ppc);	//get a random subset of the class
		pixels_used += new_C[c_idx[ci]].nnz();				//get the actual number of pixels in the subsample and use it to increment the number of pixels used in the new mask images			
	}

	C = new_C;														//replace the class image array

	if (output_mask != "") {											//if the user specifies a mask filename, save each subsampled image
		stim::filename fmask = output_mask;							//create a file name object from the mask string
		for (size_t c = 0; c < C.size(); c++) {
			stim::filename f = fmask.insert(c, 4);					//generate a numbered file name
			C[c].save(f.str());										//convert the file name to a string and save the image
		}
	}

	/*//tP = 0;															//reset the total number of pixels to zero
	for (size_t c = 0; c < C.size(); c++) {								//for each class image
		//nP[c] = C[c].nnz();											//re-calculate the number of pixels in each image
		//tP += nP[c];												//re-calculate the total number of pixels
		if (maskfile != "") {
			std::stringstream ss;
			ss << "resample_class_" << c << ".bmp";
			C[c].save(ss.str());
		}
	}*/
}


#endif