#include "stim/parser/filename.h"
#include "stim/envi/envi.h"
#include "stim/visualization/colormap.h"

#include <thread>
#include <algorithm>
#include <unordered_map>
#include <fstream>
void progress_thread_double(double* e);		//progress bar threaded function

struct speroFile{
	int x;
	int y;
	stim::filename headername;
};

void spero_coord(int &x, int &y, stim::filename filename){

	//get the file prefix
	std::string prefix = filename.get_prefix();

	//these files are in the format:
	//		???????????[xxxxx_yyyyy][??????????].hdr

	//first find the x coordinate

	//find the position of the first bracket
	size_t start = prefix.find_first_of("[") + 1;				//find the first element of the y coordinate
	prefix = prefix.substr(start, prefix.length() - start);
	size_t end = prefix.find_first_of("_");					//find the index of the last element of the y coordinate

	std::string x_string = prefix.substr(0, end);

	x = atoi(x_string.c_str());

	//crop out the x coordinate and the last underscore
	prefix = prefix.substr(end + 1, prefix.length() - end);
	end = prefix.find_first_of("]");
	std::string y_string = prefix.substr(0, end);

	y = atoi(y_string.c_str());
	//find the y coordinate using the same method as above
/*	size_t x_start = prefix.find_last_of("_");
	std::string x_string = prefix.substr(x_start + 1);
	x = atoi(x_string.c_str());*/
}

void mosaic_spero(std::string directory, std::string outfile){

	std::cout<<"Building a mosaic from SPERO frames..."<<std::endl;

	//generate a list of dmd files
	stim::filename f = directory + "/*.hdr";
	std::vector<stim::filename> hdr_list = f.get_list();

	std::vector<speroFile> speroFiles;											//create a list of spero files
	speroFiles.resize(hdr_list.size());											//pre-allocate to save time (we know how many there are)
	for(size_t i = 0; i < hdr_list.size(); i++){								//for each header name
		spero_coord(speroFiles[i].x, speroFiles[i].y, hdr_list[i]);
		speroFiles[i].headername = hdr_list[i].str();
		//speroFiles[i].file.open(hdr_list[i].str_noext(), hdr_list[i].str());
		//speroFiles[i].file.close();												//close the file (can only keep so many open at once)
	}

	//calculate the height of the image by looking at how many files have the same x coordinate
	size_t height = 0;									//stores the width of the array
	int first_x;										//the first x value encountered
	for(size_t i = 0; i < speroFiles.size(); i++){		//for each file
		if(i == 0) first_x = speroFiles[i].x;			//if this is the first file, store the x coordinate
		else{											//otherwise
			if(speroFiles[i].x != first_x){				//if the x-coord of this file is not the same as the first file
				height = (int)i;							//we now know the number of files
				break;									//exit the loop
			}
		}
	}
	size_t width = speroFiles.size() / height;			//the number of files should be divisible by the width to get the height

	
	std::vector<int> yvals(height);						//stores the height values
	for(size_t i = 0; i < height; i++)					//go through the first "width" files and store their y values
		yvals[i] = speroFiles[i].y;

	std::vector<int> xvals(width);							//stores the width values
	for(size_t i = 0; i < speroFiles.size(); i += height)	//go through every "width" file and store the x coordinate
		xvals[i/height] = speroFiles[i].x;

	std::sort(&yvals[0], &yvals[0] + yvals.size());			//sort both arrays of coordinates
	std::sort(&xvals[0], &xvals[0] + xvals.size());
	
	//create a mapping from the x and y coordinates to their indices
	std::unordered_map<int, size_t> x2idx;					
	std::unordered_map<int, size_t> y2idx;

	for(size_t i = 0; i < height; i++) y2idx.insert(std::pair<int, size_t>(yvals[i], i));	//unordered maps now give an index for a coordinate
	for(size_t i = 0; i < width; i++)  x2idx.insert(std::pair<int, size_t>(xvals[i], i));

	std::vector<stim::envi> M;							//create an array of ENVI files in mosaic order
	M.resize(speroFiles.size());

	size_t xi, yi;
	for(size_t i = 0; i < M.size(); i++){				//for each input file
		xi = x2idx[speroFiles[i].x];					//get the x and y indices for this file's position in the mosaic
		yi = y2idx[speroFiles[i].y];

		M[yi * width + xi].open(speroFiles[i].headername.str_noext(), speroFiles[i].headername.str());	//open the ENVI file and set parameters
		M[yi * width + xi].close();						//close the file, since we can only open so many
	}

	std::ofstream target(outfile.c_str(), std::ios::binary);	//create a file to store the output

	stim::envi_header h = M[0].header;			//create a header for the mosaic
	h.samples *= width;
	h.lines *= height;
	size_t tile_width = M[0].header.samples;
	size_t tile_height = M[0].header.lines;

	//h.bands = 10;

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

	float* band = (float*) malloc(h.samples * h.lines * sizeof(float));	//allocate space for a band image
	float* tile_band = (float*) malloc(tile_width * tile_height * sizeof(float));
	size_t dst, src;
	for(size_t b = 0; b < h.bands; b++){
		for(size_t yi = 0; yi < height; yi++){
			for(size_t xi = 0; xi < width; xi++){
				M[yi * width + xi].open();
				M[yi * width + xi].band_index(tile_band, b);
				for(size_t y = 0; y < tile_height; y++){
					dst = (yi * tile_height + tile_height - y - 1) * h.samples + xi * tile_width;
					src = y * tile_width;
					memcpy(band + dst, tile_band + src, sizeof(float) * tile_width);
				}
				//stim::cpu2image(tile_band, outfile + ".bmp", tile_width, tile_height, stim::cmBrewer);
				//exit(1);
				M[yi * width + xi].close();

				progress = (double)((b+1) * (yi+1) * (xi + 1)) / (width * height * h.bands) * 100;
			}
		}
		target.write((char*)band, h.samples * h.lines * sizeof(float));
		//stim::cpu2image(band, outfile + ".bmp", h.samples, h.lines, stim::cmBrewer);
		//exit(1);
	}

	t1.join();								//wait for the thread to terminate

	target.close();							//close the output file
	h.save(outfile + ".hdr");				//write the header

	free(band);
}