#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <fstream>
#include <stim/math/vector.h>
#include <stim/parser/arguments.h>
#include <stim/parser/filename.h>
#include <stim/image/image.h>
#include <ANN/ANN.h>

//load a set of points from an ASCII file
std::vector< stim::vec<float> > points_from_list(std::string filename){

	std::vector< stim::vec<float> > result;		//create an array to store the result

	std::string str;

	//load the ground truth
	std::ifstream gfile(filename.c_str());
	while(std::getline(gfile, str)){
		stim::vec<float> v(str);
		result.push_back(v);
	}

	return result;			//return the list of points
}

//load a set of points from an image (nonzero values are point positions)
std::vector< stim::vec<float> > points_from_image(std::string filename){

	stim::image<unsigned char> I;		//create an image object

	I.load(filename);					//load the image file

	std::vector<unsigned long> idx = I.sparse_idx();		//get the 1D indices for each nonzero pixel

	std::vector< stim::vec<float> > result;					//generate an array to store the 2D pixel positions
	result.resize(idx.size());								//allocate the space to save time

	unsigned long w = I.width();							//get the image width (used to resolve the 2D pixel position from the 1D index)

	unsigned long x, y;
	for(unsigned long i = 0; i < idx.size(); i++){
		y = idx[i] / w;
		x = idx[i] - y * w;

		result[i] = stim::vec<float>(x, y);
	}


	return result;											//return the array of points

}

/// This function calculates the number of hits (
void calc_hitlist(std::vector< unsigned int >& hits,
					std::vector< float >& dists,
				    std::vector< stim::vec<float> > data,
				    std::vector< stim::vec<float> > query){

	int n_data = data.size();
	int n_query = query.size();

	//allocate space in the hit lists
	hits.resize(n_query);
	dists.resize(n_query);

	//calculate the spatial dimension
	unsigned int dim = 1;
	for(unsigned int i = 0; i < data.size(); i++)
		if(data[i].size() > dim) dim = data[i].size();

										//number of data points
	ANNpointArray dataPts = annAllocPts(n_data, dim);				//array of data points
	ANNpoint queryPt = annAllocPt(dim);							//query point
	ANNidxArray nnIdx = new ANNidx[1];							//indices of nearest neighbors
	ANNdistArray sq_dist = new ANNdist[1];						//array of squared distances to nearest neighbors

	//load the ground truth points into the ANN point array
	for(unsigned int i = 0; i < data.size(); i++){
		for(unsigned int d = 0; d < dim; d++){
			dataPts[i][d] = data[i][d];
		}
	}

	//create a KD-tree
	ANNkd_tree* kdTree;											//KD tree search structure
	kdTree = new ANNkd_tree(dataPts, n_data, dim);				//create the KD tree search structure


	//submit each query point to the KD tree to find the nearest neighbor
	for(unsigned int q = 0; q < n_query; q++){

		for(unsigned int d = 0; d < dim; d++){
			queryPt[d] = query[q][d];
		}

		//query the KD-tree
		kdTree->annkSearch(queryPt, 1, nnIdx, sq_dist);

		//put the result in the hit array
		hits[q] = nnIdx[0];
		dists[q] = sqrt(sq_dist[0]);
	}
}

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

	//output advertisement
	std::cout<<std::endl<<std::endl;
	std::cout<<"========================================================================="<<std::endl;
	std::cout<<"Thank you for using the PointMets validation tool!"<<std::endl;
	std::cout<<"Scalable Tissue Imaging and Modeling (STIM) Lab, University of Houston"<<std::endl;
	std::cout<<"Developers: David Mayerich"<<std::endl;
	std::cout<<"Source: https://git.stim.ee.uh.edu/validation/pointmets"<<std::endl;
	std::cout<<"========================================================================="<<std::endl<<std::endl;


	stim::arglist args;

#ifdef _WIN32
	args.set_ansi(false);
#endif

	//add arguments
	args.add("help", "prints this help");
	args.add("r", "maximum possible distance to consider a point 'detected'", "10", "[positive value]");

	//parse the command line arguments
	args.parse(argc, argv);

	//display the help text if requested
	if(args["help"].is_set()){				
		std::cout<<std::endl<<"usage: pointmets ground_truth test_case --option [A B C ...]"<<std::endl;
		std::cout<<std::endl<<std::endl
				  << "examples: pointmets blue_gt.txt blue_t.txt"<<std::endl;
		std::cout << "          pointmets blue_gt.bmp blue_t.bmp" << 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() < 2){
		std::cout<<"ERROR: Two files must be specified for comparison, enter pointmets --help for options."<<std::endl<<std::endl;
		exit(1);
	}

	//get the ground truth and test file
	stim::filename gfilename(args.arg(0));
	stim::filename tfilename(args.arg(1));

	//set the radius
	float r = args["r"].as_float();

	std::vector< stim::vec<float> > G;
	std::vector< stim::vec<float> > T;
	

	//if an ASCII list of points is specified
	if(gfilename.get_extension() == "bmp")
		G = points_from_image(gfilename.str());
	else
		G = points_from_list(gfilename.str());

	if(tfilename.get_extension() == "bmp")
		T = points_from_image(gfilename.str());
	else
		T = points_from_list(tfilename.str());

	//calculate the hit list and distances for the test array
	std::vector< unsigned int > T_hits;
	std::vector< float > T_dists;
	calc_hitlist(T_hits, T_dists, G, T);

	//calculate the hit list and distances for the ground truth array
	std::vector< unsigned int > G_hits;
	std::vector< float > G_dists;
	calc_hitlist(G_hits, G_dists, T, G);

	//count up metrics for each test point
	std::vector< unsigned int > TP;
	std::vector< unsigned int > FP;
	unsigned int g, t;
	for(t = 0; t < T.size(); t++){

		//a true positive has a matching hit in both hit lists, with a distance less than the radius
		g = T_hits[t];
		if(G_hits[g] == t && T_dists[t] <= r) TP.push_back(t);
		else FP.push_back(t);
	}

	std::vector< unsigned int > FN;
	for(g = 0; g < G.size(); g++){
		t = G_hits[g];
		if(T_hits[t] != g || G_dists[g] > r) FN.push_back(g);
	}

	std::cout<<"N: "<<G.size()<<std::endl;
	std::cout<<"TP: "<<TP.size()<<"   ( "<<(double)TP.size() / (double)G.size() * 100<<"% detected)"<<std::endl;
	std::cout<<"FN: "<<FN.size()<<"   ( "<<(double)FN.size() / (double)G.size() * 100<<"% undetected)"<<std::endl;
	std::cout<<"FP: "<<FP.size()<<"   ( "<<(double)FP.size() / (double)T.size() * 100<<"% incorrectly detected)"<<std::endl;



}