#include <iostream>
#include <algorithm>
#include <unordered_map>
#include <stim/parser/arguments.h>
#include <stim/envi/agilent_binary.h>
#include <stim/parser/filename.h>
#include <stim/envi/envi.h>

stim::arglist args;								//generate a list of accepted arguments

#include <thread>
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());*/
}

bool sperosort(speroFile i, speroFile j) {
	return i.x < j.x;
}

//This function test to make sure that the given file names are valid for a SPERO reconstruction
bool verify_filenames(std::vector<stim::filename> files) {
	size_t idx = files[0].prefix().find('[');
	if (idx == std::string::npos) return false;
	for (size_t i = 0; i < files.size() - 1; i++) {
		if (files[i].prefix().substr(0, idx) != files[i + 1].prefix().substr(0, idx)) return false;
	}
	return true;
}

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

	//test for wild-cards in the directory
	if (directory.find('*') != std::string::npos || directory.find('?') != std::string::npos) {
		std::cout << "ERROR: no wild-cards are allowed, just specify the source directory containing SPERO files" << std::endl;
		exit(1);
	}
	
	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();
	if (hdr_list.size() == 0) {
		std::cout << "ERROR: no files found for assembly" << std::endl;
		exit(1);
	}
	std::cout << hdr_list.size() << " files found" << std::endl;

	if (!verify_filenames(hdr_list)) {
		std::cout << "ERROR: incorrect file prefix found - all files should have the same prefix" << std::endl;
		exit(1);
	}

	if (hdr_list.size() == 0) {
		std::cout << "ERROR: no header files found in directory '" + directory << std::endl;
		exit(1);
	}

	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();
	}

	//sort the SPERO files by x coordinate
	std::sort(speroFiles.begin(), speroFiles.end(), sperosort);

	//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
		if (first_x == speroFiles[i].x) height++;
		else break;
	}
	size_t width;
	if (height == 0) width = speroFiles.size();
	else width = speroFiles.size() / height;			//the number of files should be divisible by the width to get the height

	if (width * height != speroFiles.size()) {
		std::cout << "ERROR: calculated width and height do not match the number of input files: " << width << " x " << height << " = " << width * height << " found,  " << speroFiles.size() << " needed" << std::endl;
		exit(1);
	}

	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];

		size_t idx = yi * width + xi;
		M[idx].open(speroFiles[i].headername.str_noext(), speroFiles[i].headername.str());	//open the ENVI file and set parameters
		M[idx].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;

	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;

	//reconstruct the array
	for (size_t b = 0; b < h.bands; b++) {								//for each band
		for (size_t yi = 0; yi < height; yi++) {						//for each tile in X and Y
			for (size_t xi = 0; xi < width; xi++) {
				M[yi * width + xi].open();								//open the tile
				M[yi * width + xi].band_index(tile_band, b);			//fill the tile_band pointer with the tile image at band b
				for (size_t y = 0; y < tile_height; y++) {				//copy the tile image to the destination band image
					//dst = (yi * tile_height + tile_height - y - 1) * h.samples + xi * tile_width;
					dst = ((height - yi - 1) * tile_height + y) * h.samples + xi * tile_width;
					src = y * tile_width;
					memcpy(band + dst, tile_band + src, sizeof(float) * tile_width);
				}

				M[yi * width + xi].close();								//close the tile

				progress = (double)(b * height * width + yi * width + xi + 1) / (width * height * h.bands) * 100;	//update the progress bar
			}
		}
		target.write((char*)band, h.samples * h.lines * sizeof(float));	//write the entire band to the output image
	}

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

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

	free(band);
}

void advertise() {
	std::cout << std::endl << std::endl;
	std::cout << "=========================================================================" << std::endl;
	std::cout << "Thank you for using the SIPROC spectroscopic image processing toolkit!" << std::endl;
	std::cout << "Scalable Tissue Imaging and Modeling (STIM) Lab, University of Houston" << std::endl;
	std::cout << "Developers: Sam Saki, Ziqi He, David Mayerich, Brad Deutsch" << std::endl;
	std::cout << "Source: https://github.com/stimlab/hsiproc.git" << std::endl;
	std::cout << "This version has been compiled on " << __DATE__ << " at " << __TIME__ << std::endl;
	std::cout << "=========================================================================" << std::endl << std::endl;

	std::cout << "Daylight SPERO mosaic assembly usage:" << std::endl;
	std::cout << "     >> spero directory/to/files mosaic" << std::endl;
}

int main(int argc, char* argv[]) {

	args.add("help", "print help");
	args.parse(argc, argv);								//parse the argument list

	if (args["help"]) {
		std::cout << args.str() << std::endl;
		exit(1);
	}
	std::string indir = ".";
	std::string outfile = "mosaic";
	if (args.nargs() == 1)	outfile = args.arg(0);				//if only one argument is given, assume it is the output file
	else if (args.nargs() >= 2) {									//if two arguments are given
		indir = args.arg(0);									//the first argument is the directory
		outfile = args.arg(1);									//the second argument is the output file
	}
	std::cout << "SPERO mosaic assembly:    " << indir << " --> " << outfile << std::endl;
	mosaic_spero(indir, outfile);								//call the SPERO mosaic function
}