#include <iostream>
#include <string>
#include <fstream>
#include <cuda_runtime.h>
#include <stim/math/vector.h>
#include <stim/parser/arguments.h>
#include <stim/parser/filename.h>
#include <stim/visualization/colormap.h>
#include <stim/image/image.h>
#include <stim/math/constants.h>

stim::arglist args;
int iter;
unsigned int rmax;
unsigned int nlmax;
float t;
float sigma;
float phi;
size_t x, y, z;


void ivote3(float* img, float std[], float phi, float d_phi, unsigned int r[], int iter, float t, unsigned int conn[], 
			unsigned int x, unsigned int y, unsigned int z);
void lmax(float* center, float* vote, float t1, unsigned int conn[], unsigned int x, unsigned int y, unsigned int z);

void invert_data(float* cpuI, unsigned int x, unsigned int y, unsigned int z){
		for(int ix = 0; ix < x; ix++){
			for (int iy = 0; iy < y; iy++){
				for (int iz = 0; iz < z; iz++){
					int idx = iz * x * y + iy * x + ix;
					cpuI[idx] = 255 - cpuI[idx];
				}
			}
		}
	}
	


void advertise() {
	std::cout << std::endl << std::endl;
	std::cout << "=========================================================================" << std::endl;
	std::cout << "Thank you for using the ivote3 segmentation tool!" << std::endl;
	std::cout << "Scalable Tissue Imaging and Modeling (STIM) Lab, University of Houston" << std::endl;
	std::cout << "Developers: Laila Saadatifard and David Mayerich" << std::endl;
	std::cout << "Source: https://git.stim.ee.uh.edu/segmentation/ivote3" << std::endl;
	std::cout << "=========================================================================" << std::endl << std::endl;

}

void init_args(int argc, char* argv[]) {
#ifdef _WIN32
	args.set_ansi(false);
#endif

	//add arduments
	args.add("help", "prints this help");
	args.add("x", "size of the dataset along X axis", "positive value");
	args.add("y", "size of the dataset along Y axis", "positive value");
	args.add("z", "size of the dataset along Z axis", "positive value");
	args.add("t", "threshold value for the final result", "positive valu");
	args.add("invert", "to invert the input data set", "string");
	args.add("rmax", "maximum possible radius of the cells in the input image", "10", "[positive value]");
	args.add("phi", "starting angle for the vote region (in degrees)", "25.0", "0 <= phi < 180");
	args.add("iter", "number of iterations for voting", "8", "i > 0");
	args.add("sigma", "the gaussian blur standard deviation", "5", "s >=0 (s = 0, no blurring)");
	args.add("conn", "the number of connected neighbors for calculating the local maxima", "5", "[positive value]");
	//parse the command line arguments.
	args.parse(argc, argv);

	//display the help text if requested
	if (args["help"].is_set()) {
		advertise();
		std::cout << std::endl << "usage: ivote input_image output_list --option [A B C ...]" << std::endl;
		std::cout << std::endl << std::endl
			<< "examples: ivote blue.bmp list.txt " << std::endl;

		std::cout << std::endl << std::endl;
		std::cout << args.str() << std::endl;
		exit(1);
	}

	//if the input and output files aren't specified, throw an error and exit
	if (args.nargs() < 3) {
		std::cout << "ERROR: three files must be specified for segmentation, enter ivote --help for options." << std::endl << std::endl;
		exit(1);
	}

	//set the x, y, z.
	x = (size_t)args["x"].as_int();
	y = (size_t)args["y"].as_int();
	z = (size_t)args["z"].as_int();
	iter = args["iter"].as_int();
	rmax = (unsigned int)args["rmax"].as_int();
	nlmax = (unsigned int)args["conn"].as_int();
	t = args["t"].as_float();
	sigma = args["sigma"].as_float();
	phi = (float)args["phi"].as_float() * (float)stim::PI / 180;
	
}
int main(int argc, char** argv){


	cudaDeviceProp prop;
	int count;
	cudaGetDeviceCount(&count);
	for (int i=0; i<count; i++){
		cudaGetDeviceProperties(&prop, i);
		printf("current device ID: %d\n", i);
		printf("device name: %s\n", prop.name);
		printf("total global mem: %lu\n", prop.totalGlobalMem);
		printf("shared memory per block: %lu\n", prop.sharedMemPerBlock);
	}
		
	init_args(argc, argv);

	unsigned int r[3] = { rmax  , rmax, rmax};
	
	float sigma3[3] = { sigma, sigma, sigma};
	unsigned int conn[3] = { nlmax, nlmax, nlmax};
	float d_phi = phi/(iter+2);
	
	size_t bytes = x*y*z*sizeof(float);

	//allocate space on the cpu for the input data
	float* cpuI = (float*) malloc(bytes);

	//load the input file into the cpuI
	std::ifstream nissl(args.arg(0), std::ios::in | std::ios::binary);
	nissl.read((char*)cpuI, bytes);
	nissl.close();
	if(args["invert"].is_set())
		invert_data(cpuI, x, y, z);
	
	//write a new file from the cpuI.
	std::ofstream original("0-inv-128.vol", std::ofstream::out | std::ofstream::binary);
	original.write((char*)cpuI, bytes);
	original.close();
	
	
	ivote3(cpuI, sigma3, phi, d_phi, r, iter, t, conn, x, y, z);			// call the ivote function
		
	std::ofstream fvote("0-vote8.vol", std::ofstream::out | std::ofstream::binary);
	fvote.write((char*)cpuI, bytes);
	fvote.close();
	
	//allocate space on the cpu for the output result
	float* cpu_out = (float*)malloc(bytes * 3);

	//write the output file.
	//for (int t0=0; t0<=5000; t0+=100){
	//	float t1 = t0;
		int t0 = t;
			lmax(cpu_out, cpuI, t, conn, x, y, z);
			//std::ofstream fo("shared2D-v8/" + OutName.str(), std::ofstream::out | std::ofstream::binary);
			std::ofstream fo( args.arg(1), std::ofstream::out | std::ofstream::binary);
			fo.write((char*)cpu_out, bytes);
			fo.close();
	
	// creat a file for saving the list centers
	
		std::ofstream list(args.arg(2));
		// set the number of detected cells to zero.
		int nod = 0;
		if (list.is_open()){

				for (int iz=0; iz<z; iz++){
					for (int iy=0; iy<y; iy++){
						for (int ix=0; ix<x; ix++){

							int idx = iz * x * y + iy * x + ix;
							if (cpu_out[idx]>0){
								nod++;
								list << ix << " " << iy << " "<< iz << " " << cpu_out[idx] << '\n' ;
					
							}
						}
					}
				}
				
		list.close();
		}


	//}
		cudaDeviceReset();       
	
}