#ifndef RTSSKELETONIZER_H
#define RTSSKELETONIZER_H

#include "rts_itkVolume.h"
#include "objJedi.h"
#include <vector>
#include <queue>

using namespace std;

typedef vector<point3D<double>> filament;
typedef unsigned int indextype;

#define BACKGROUND_NODE	0
#define SKELETON_NODE	255
#define ENDPOINT_NODE	128
#define VISITED_NODE	64

class rtsSkeletonizer
{
private:
	rts_itkVolume<unsigned char> implicit_skeleton;
	/*Key:
	0 - background node
	255 - skeleton node
	128 - endpoint node
	*/

	vector<filament> explicit_filaments;
	vector<point3D<unsigned int>> explicit_endpoints;



public:
	rtsSkeletonizer(rts_itkVolume<unsigned char> input);
	unsigned int BackgroundComponents6(rts_itkVolume<unsigned char>* function,
									   indextype x, 
									   indextype y, 
									   indextype z, 
									   unsigned char threshold);
	unsigned int rtsSkeletonizer::Peel(unsigned char isovalue);
	int Neighbors26(rts_itkVolume<unsigned char>* function, 
					indextype x, 
					indextype y, 
					indextype z, 
					unsigned char isovalue);
	vector<point3D<unsigned int>> FloodFill6(rts_itkVolume<unsigned char>* function, indextype x, indextype y, indextype z, unsigned char target_value);
	void FloodFill26(rts_itkVolume<unsigned char>* function, int x, int y, int z, unsigned char target_value);
	unsigned int Neighbors6(rts_itkVolume<unsigned char>* function, indextype x, indextype y, indextype z, unsigned char threshold);
	rts_itkVolume<unsigned char> getImplicitResult(){return implicit_skeleton;}
	unsigned int NumConnectedComponents(indextype x, indextype y, indextype z, unsigned char threshold);

	void ComputeEndPoints();	//find all of the filament end points
	void TraceFilament(point3D<unsigned int> p0, point3D<unsigned int> p1);
	rts_itkVolume<unsigned char> getLocalRegion(unsigned int x, unsigned int y, unsigned int z);

	void ConstructExplicitSkeleton();
	void SmoothExplicitSkeleton(int iterations);
	void SaveExplicitSkeleton(const char* filename);


};

void rtsSkeletonizer::SmoothExplicitSkeleton(int iterations)
{
	/*This function smooths the vessels by repeatedly taking the dual of the skeleton.
	Branch points and end points are constrained to their initial positions.
	*/

	vector<filament>::iterator f;
	int num_verts, end_vert;
	int v;
	point3D<double> new_vert;

	for(int i = 0; i<iterations; i++)
	{
		for(f=explicit_filaments.begin(); f!=explicit_filaments.end(); f++)
		{
			filament new_filament;
			
			//find the number of vertices in the filament
			num_verts = f->size();
			end_vert = num_verts - 1;

			//take care of the first vertex
			new_filament.push_back((*f)[0]);

			for(v=1; v<end_vert; v++)
			{
				//find the position of the next vertex (midpoint between the last and current vertices)
				new_vert = (*f)[v-1] + ((*f)[v] - (*f)[v-1])*0.5;
				new_filament.push_back(new_vert);			
			}

			//take care of the last vertex
			new_filament.push_back((*f)[end_vert]);

			//delete the old filament and add the new one
			(*f) = new_filament;
		}
	}
		
}

void rtsSkeletonizer::SaveExplicitSkeleton(const char* filename)
{
	/*Save the filament file as an OBJ*/
	//create the OBJ object
	rtsOBJ output;

	unsigned int num_filaments = explicit_filaments.size();
	//cout<<"Number of filaments:  "<<num_filaments<<endl;
	vector<filament>::iterator i;
	filament::iterator f;
	for(i=explicit_filaments.begin(); i!=explicit_filaments.end(); i++)
	{
		output.objBegin(OBJ_LINE_STRIP);
		//cout<<"start filament"<<endl;
		for(f=i->begin(); f!=i->end(); f++)
		{
			output.objVertex3f((*f).x, (*f).y, (*f).z);
			//f->print();
		}
		//cout<<"end filament"<<endl<<endl;
		output.objEnd();
	}

	output.SaveFile(filename);
		
}

rts_itkVolume<unsigned char> rtsSkeletonizer::getLocalRegion(unsigned int x, unsigned int y, unsigned int z)
{
	rts_itkVolume<unsigned char> local(3, 3, 3);
	point3D<indextype> corner;
	indextype u, v, w;
	corner = point3D<indextype>(x-1, y-1, z-1);
	for(u=0; u<3; u++)
		for(v=0; v<3; v++)
			for(w=0; w<3; w++)
				local.SetPixel(u, v, w, implicit_skeleton.GetPixel(corner.x + u, corner.y + v, corner.z + w));
	return local;
}


void rtsSkeletonizer::TraceFilament(point3D<unsigned int> p0, point3D<unsigned int> p1)
{
	//create the new filament and start it with the two given points
	filament result;
	result.push_back(p0);
	result.push_back(p1);

	//if(p0 == point3D<unsigned int>(156, 433, 254))
	//	cout<<"Stop here"<<endl;

	//if p1 is an endpoint, go ahead and end the filament right away
	if(implicit_skeleton.GetPixel(p1.x, p1.y, p1.z) == ENDPOINT_NODE)
	{
		explicit_filaments.push_back(result);
		return;
	}

	rts_itkVolume<unsigned char> local;
	point3D<unsigned int> current = p1;
	point3D<unsigned int> next;
	point3D<unsigned int> possible_end;
	int x, y, z;
	int c;		//possible connections

	while(1)	//continue tracing until the next endpoint
	{
		c=0;
		local = getLocalRegion(current.x, current.y, current.z);
		//cout<<"Local Region at ("<<current.x<<","<<current.y<<","<<current.z<<")"<<endl;
		//local.toConsole();
		local.SetPixel(1, 1, 1, 0);
		for(x=0; x<3; x++)
			for(y=0; y<3; y++)
				for(z=0; z<3; z++)
				{
					//if the node is a skeleton node, store it as the next point
					if(local(x, y, z) == SKELETON_NODE)
					{
						c++;
						next = current - vector3D<unsigned int>(1, 1, 1) + vector3D<unsigned int>(x, y, z);
					}
					//if it is an endpoint node, check for a termination
					if(local(x, y, z) == ENDPOINT_NODE)
					{
						//if we are not backtracking to the start node
						possible_end = current - vector3D<unsigned int>(1, 1, 1) + vector3D<unsigned int>(x, y, z);
						if(possible_end != p0 || result.size() > 2)
						{
							//terminate the fiber
							result.push_back(current - vector3D<unsigned int>(1, 1, 1) + vector3D<unsigned int>(x, y, z));
							explicit_filaments.push_back(result);
							implicit_skeleton.SetPixel(current.x, current.y, current.z, VISITED_NODE);
							return;
						}
					}
				}
		//return when there is nowhere to go
		if(c == 0)
			return;
		//now we have the next node in the filament, go ahead and label the current
		//one as traversed
		implicit_skeleton.SetPixel(current.x, current.y, current.z, VISITED_NODE);
		//now insert it into the fiber and move to the next node
		result.push_back(next);
		if(current == next)
		{
			//cout<<"I've been waiting for you Obi-Wan"<<endl;
			local = getLocalRegion(current.x, current.y, current.z);
			//cout<<"Local Region at ("<<current.x<<","<<current.y<<","<<current.z<<")"<<endl;
			//local.toConsole();

			//cout<<"Number of neighbors: "<<Neighbors26(&implicit_skeleton, current.x, current.y, current.z, 1);
		}
		current = next;
	}
	
	
}

void rtsSkeletonizer::ConstructExplicitSkeleton()
{
	/*The implicit skeleton is constructed as follows:
	1)  Label all branch nodes.
	2)  For each branch, find the number of fibers connected to that branch.
		a)  Create each fiber as a seperate structure and store them in a queue.
		b)  Iteratively track each fiber.
	*/
	ComputeEndPoints();	//compute the endpoints for the explicit skeleton and
						//label them in the implicit skeleton
	
	//iterate through each end point
	vector<point3D<unsigned int>>::iterator i;
	rts_itkVolume<unsigned char> local;
	unsigned int x, y, z;
	for(i = explicit_endpoints.begin(); i != explicit_endpoints.end(); i++)
	{
		//cout<<"Tracing Filament..."<<endl;
		//i->print();
		local = getLocalRegion(i->x, i->y, i->z);
		local.SetPixel(1, 1, 1, 0);
		//local.toConsole();
		for(x=0; x<3; x++)
			for(y=0; y<3; y++)
				for(z=0; z<3; z++)
					if(local.GetPixel(x, y, z) > VISITED_NODE)
					{
						//cout<<(int)local.xyz(x, y, z)<<endl;
						//cout<<(int)implicit_skeleton.xyz(i->x - 1 + x,
						//							  i->y - 1 + y,
						//							  i->z - 1 + z)<<endl;
						TraceFilament((*i), point3D<unsigned int>(i->x - 1 + x,
																  i->y - 1 + y,
																  i->z - 1 + z));
					}
		//since we have tracked all filaments attached to an endpoint
		//set the endpoint as visited
		implicit_skeleton.SetPixel(i->x, i->y, i->z, VISITED_NODE);
	}

	//cout<<"Tracing complete"<<endl;


}


rtsSkeletonizer::rtsSkeletonizer(rts_itkVolume<unsigned char> input)
{
	implicit_skeleton = input;
}

void rtsSkeletonizer::ComputeEndPoints()
{
	/*This function finds all of the filament end points in the implicit skeleton.
	The end points are stored in the explicit_endpoints vector.
	An end point is defined as a node with (n == 1) or (n > 2) connected components.
	*/

	unsigned int max_x = implicit_skeleton.DimX();
	unsigned int max_y = implicit_skeleton.DimY();
	unsigned int max_z = implicit_skeleton.DimZ();
	unsigned int conn;

	indextype x, y, z;
	for(x=0; x<max_x; x++)		//go through each voxel in the data set
		for(y=0; y<max_y; y++)
			for(z=0; z<max_z; z++)
			{
				if(implicit_skeleton(x, y, z) == SKELETON_NODE)	//if the voxel is part of the skeleton
				{
					conn = Neighbors26(&implicit_skeleton, x, y, z, 1);	//test the number of connected components
					if(conn == 1 || conn > 2)	//if the number of connected components meets the
												//criteria, store the node in the explicit_endpoints vector.
					{
						explicit_endpoints.push_back(point3D<unsigned int>(x, y, z));
					}
				}

			}

	//tag the endpoints in the implicit function
	vector<point3D<unsigned int>>::iterator i;
	for(i=explicit_endpoints.begin(); i!=explicit_endpoints.end(); i++)
		implicit_skeleton.SetPixel(i->x, i->y, i->z, ENDPOINT_NODE);

			//cout<<"Number of end points:  "<<explicit_endpoints.size()<<endl;
}



unsigned int rtsSkeletonizer::NumConnectedComponents(indextype x, indextype y, indextype z, unsigned char threshold)
{
	unsigned int result = 0;

	rts_itkVolume<unsigned char> local(3, 3, 3);
	

	local = getLocalRegion(x, y, z);
//	local.toConsole();
	local.SetPixel(1, 1, 1, 0);
//	local.toConsole();

	indextype xi, yi, zi;
	//iterate through each voxel
	for(xi=0; xi<3; xi++)
		for(yi=0; yi<3; yi++)
			for(zi=0; zi<3; zi++)
				if(local(xi, yi, zi) >= threshold)
				{
					FloodFill6(&local, xi, yi, zi, 0);
					result++;
				}
	/*
	if(result == 1 || result > 2)
	{
		cout<<result<<endl;
		local = getLocalRegion(x, y, z);
		local.toConsole();
	}
	*/
	return result;


}

unsigned int rtsSkeletonizer::BackgroundComponents6(rts_itkVolume<unsigned char>* function, 
													   indextype x, 
													   indextype y, 
													   indextype z, 
													   unsigned char threshold)
{
	/**
	This function computes the number of 6-connected background components in the local region of (x, y, z).
	This computation is performed by testing all 6 possible voxels that can connect to the specified node.  If
	a background node is found, the entire background component associated with that node is filled and the counter
	is incremented by 1.  The value n specifies the connectivity domain for the flood fill.
	The definition of background components is that specified by He, Kischell, Rioult and Holmes.
	**/

	//see if there is at least one BG component
	if(Neighbors6(function, x, y, z, threshold) == 6)
		return 0;

	
	//retrieve the local region of the function
	rts_itkVolume<unsigned char> local(3, 3, 3);
	point3D<indextype> corner(x-1, y-1, z-1);
	indextype u, v, w;
	for(u=0; u<3; u++)
		for(v=0; v<3; v++)
			for(w=0; w<3; w++)
				local.SetPixel(u, v, w, function->GetPixel(corner.x + u, corner.y + v, corner.z + w));

	//threshold the background to find inside/outside points
	local.Binary(threshold, 1);
	//fill points that are not in the connectivity domain
	/*if(n == 18)
	{
		local(0, 0, 0) = 1;
		local(0, 0, 2) = 1;
		local(0, 2, 0) = 1;
		local(0, 2, 2) = 1;
		local(2, 0, 0) = 1;
		local(2, 0, 2) = 1;
		local(2, 2, 0) = 1;
		local(2, 2, 2) = 1;
	}*/
	//local.toConsole();

	//search all 6 possible connected points.  If a background node is found, fill the component
	unsigned int components = 0;
	if(local(0, 1, 1) == 0)
	{
		components++;
		FloodFill6(&local, 0, 1, 1, 1);
	}
	if(local(2, 1, 1) == 0)
	{
		components++;
		FloodFill6(&local, 2, 1, 1, 1);
	}
	if(local(1, 0, 1) == 0)
	{
		components++;
		FloodFill6(&local, 1, 0, 1, 1);
	}
	if(local(1, 2, 1) == 0)
	{
		components++;
		FloodFill6(&local, 1, 2, 1, 1);
	}
	if(local(1, 1, 0) == 0)
	{
		components++;
		FloodFill6(&local, 1, 1, 0, 1);
	}
	if(local(1, 1, 2) == 0)
	{
		components++;
		FloodFill6(&local, 1, 1, 2, 1);
	}

	return components;
}

unsigned int rtsSkeletonizer::Neighbors6(rts_itkVolume<unsigned char>* function, indextype x, indextype y, indextype z, unsigned char threshold)
{

	unsigned int neighbors = 0;
	if(function->GetPixel(x+1, y, z) >= threshold)
		neighbors++;
	if(function->GetPixel(x-1, y, z) >= threshold)
		neighbors++;
	if(function->GetPixel(x, y+1, z) >= threshold)
		neighbors++;
	if(function->GetPixel(x, y-1, z) >= threshold)
		neighbors++;
	if(function->GetPixel(x, y, z+1) >= threshold)
		neighbors++;
	if(function->GetPixel(x, y, z-1) >= threshold)
		neighbors++;

	return neighbors;
}

int rtsSkeletonizer::Neighbors26(rts_itkVolume<unsigned char>* function,
									indextype x, 
									indextype y, 
									indextype z, 
									unsigned char isovalue)
{
	int neighbors = 0;
	int u,v,w;

	for(u=-1; u<=1; u++)
		for(v=-1; v<=1; v++)
			for(w=-1; w<=1; w++)
				if(function->GetPixel(x+u, y+v, z+w) >= isovalue)
					neighbors++;
	if(function->GetPixel(x, y, z) > isovalue)
		neighbors--;

	return neighbors;
}

unsigned int rtsSkeletonizer::Peel(unsigned char isovalue)
{
	/**
	Removes one layer of pixels from the specified isosurface.
	**/

	unsigned int res_x = implicit_skeleton.DimX();
	unsigned int res_y = implicit_skeleton.DimY();
	unsigned int res_z = implicit_skeleton.DimZ();
	vector<point3D<indextype>> border_nodes;	//get the border nodes for the image
	indextype x, y, z;
	int condition;
	for(x=0; x<res_x; x++)
	for(y=0; y<res_y; y++)
	for(z=0; z<res_z; z++)
		//find the border nodes
		if(implicit_skeleton(x, y, z) >= isovalue && BackgroundComponents6(&implicit_skeleton, x, y, z, isovalue) == 1 && Neighbors26(&implicit_skeleton, x, y, z, isovalue) != 1)
		{
			condition = 0;
			//now find the border pairs.  A border point must meet two of the following conditions to be a border pair.
			//south border: s(p) is background
			if(implicit_skeleton(x, y-1, z) < isovalue)
				condition++;
			//north border: n(p) is background, s(p) and s(s(p)) are foreground
			if(implicit_skeleton(x, y+1, z) < isovalue && implicit_skeleton(x, y-1, z) >= isovalue && implicit_skeleton(x, y-2, z) >= isovalue)
				condition++;
			//west border: w(p) is background
			if(implicit_skeleton(x-1, y, z) < isovalue)
				condition++;
			//east border: e(p) is background, w(p) and w(w(p)) are foreground
			if(implicit_skeleton(x+1, y, z) < isovalue && implicit_skeleton(x-1, y, z) >= isovalue && implicit_skeleton(x-2, y, z) >= isovalue)
				condition++;
			//up border: u(p) is background
			if(implicit_skeleton(x, y, z-1) < isovalue)
				condition++;
			//down border: d(p) is background, u(p) and u(u(p)) are foreground
			if(implicit_skeleton(x, y, z+1) < isovalue && implicit_skeleton(x, y, z-1) >= isovalue && implicit_skeleton(x, y, z-2) >= isovalue)
				condition++;
			
			if(condition > 1)
				border_nodes.push_back(point3D<indextype>(x, y, z));
		}
	//cout<<"Number of border nodes: "<<border_nodes.size()<<endl;

	//determine if each edge node can be removed without changing the topology of the model
	//declare some initial variables
	rts_itkVolume<unsigned char> local(3, 3, 3);	//store the region local to the current voxel
	int nodes_before, nodes_after;		//number of neighboring nodes before and after the filling operation
	point3D<indextype> fill_start;
	vector<point3D<indextype>>::iterator i;

	/*
	Here we determine if a point can be removed by looking at the number of foreground connected
	components in the local region.  If there is more than one connected component
	*/
	unsigned int removed = 0;
	for(i=border_nodes.begin(); i<border_nodes.end(); i++)
	{
		//get the local region around the current point
		for(x=-1; x<=1; x++)
			for(y=-1; y<=1; y++)
				for(z=-1; z<=1; z++)
					local.SetPixel(x+1, y+1, z+1, implicit_skeleton((*i).x + x, (*i).y + y, (*i).z + z));

		local.SetPixel(1, 1, 1, 0);				//remove the center voxel
		local.Binary(isovalue, 1);
		nodes_before = Neighbors26(&local, 1, 1, 1, 1);


		//find an interior voxel to fill
		for(x=0; x<3; x++)
			for(y=0; y<3; y++)
				for(z=0; z<3; z++)
					if(local(x, y, z) > 0)
						fill_start = point3D<indextype>(x, y, z);

		//fill the local region
		FloodFill26(&local, fill_start.x, fill_start.y, fill_start.z, 2);
		//get the number of filled neighbors
		nodes_after = Neighbors26(&local, 1, 1, 1, 2);
		if(nodes_after == nodes_before)
		{
			implicit_skeleton.SetPixel((*i).x, (*i).y, (*i).z, 0);
			removed++;
		}

	}
	return removed;
}

vector<point3D<unsigned int>> rtsSkeletonizer::FloodFill6(rts_itkVolume<unsigned char>* function, indextype x, indextype y, indextype z, unsigned char target_value)
{
	//create a vector containing all of the filled nodes
	vector<point3D<unsigned int>> result;

	unsigned char old_value = function->GetPixel(x, y, z);					//find the old value (the value being flood-filled)
	if(target_value == old_value)				//if the target value is the same as the old value, nothing to do
		return result;

	queue<point3D<indextype>> Q;				//create a queue for neighboring points
	point3D<indextype> current(x, y, z);		//start with the current point
	function->SetPixel(current.x, current.y, current.z, target_value);
	point3D<indextype> next;
	Q.push(current);
	indextype u, v, w;


	while(!Q.empty())							//continue until the queue is empty
	{
		current = Q.front();					//get the first element from the queue
		Q.pop();							
		
		if(current.x != function->DimY() - 1)
			if(function->GetPixel(current.x + 1, current.y, current.z) == old_value)
			{
				function->SetPixel(current.x + 1, current.y, current.z, target_value);
				result.push_back(point3D<unsigned int>(current.x + 1, current.y, current.z));
				Q.push(point3D<indextype>(current.x + 1, current.y, current.z));
			}
		if(current.x != 0)
			if(function->GetPixel(current.x - 1, current.y, current.z) == old_value)
			{
				function->SetPixel(current.x - 1, current.y, current.z, target_value);
				result.push_back(point3D<unsigned int>(current.x - 1, current.y, current.z));
				Q.push(point3D<indextype>(current.x - 1, current.y, current.z));
			}
		if(current.y != function->DimY() - 1)
			if(function->GetPixel(current.x, current.y +1, current.z) == old_value)
			{
				function->SetPixel(current.x, current.y+1, current.z, target_value);
				result.push_back(point3D<unsigned int>(current.x, current.y+1, current.z));
				Q.push(point3D<indextype>(current.x, current.y+1, current.z));
			}
		if(current.y != 0)
			if(function->GetPixel(current.x, current.y-1, current.z) == old_value)
			{
				function->SetPixel(current.x, current.y-1, current.z, target_value);
				result.push_back(point3D<unsigned int>(current.x, current.y-1, current.z));
				Q.push(point3D<indextype>(current.x, current.y-1, current.z));
			}
		if(current.z != function->DimZ() - 1)
			if(function->GetPixel(current.x, current.y, current.z+1) == old_value)
			{
				function->SetPixel(current.x, current.y, current.z+1, target_value);
				result.push_back(point3D<unsigned int>(current.x, current.y, current.z+1));
				Q.push(point3D<indextype>(current.x, current.y, current.z+1));
			}
		if(current.z != 0)
			if(function->GetPixel(current.x, current.y, current.z-1) == old_value)
			{
				function->SetPixel(current.x, current.y, current.z-1, target_value);
				result.push_back(point3D<unsigned int>(current.x, current.y, current.z-1));
				Q.push(point3D<indextype>(current.x, current.y, current.z-1));
			}

	}

	return result;

}

void rtsSkeletonizer::FloodFill26(rts_itkVolume<unsigned char>* function, int x, int y, int z, unsigned char target_value)
{
	unsigned char old_value = function->GetPixel(x, y, z);
	if(target_value == old_value)
		return;

	queue<point3D<indextype>> Q;
	point3D<indextype> current(x, y, z);
	point3D<indextype> next;
	Q.push(current);
	indextype u, v, w;
	while(!Q.empty())
	{
		current = Q.front();
		if(function->GetPixel(current.x, current.y, current.z) == old_value)
			function->SetPixel(current.x, current.y, current.z, target_value);
		Q.pop();

		unsigned int res_x = function->DimX();
		unsigned int res_y = function->DimY();
		unsigned int res_z = function->DimZ();
		for(u=-1; u<=1; u++)
		for(v=-1; v<=1; v++)
		for(w=-1; w<=1; w++)
		{
			next.x = current.x + u;
			next.y = current.y + v;
			next.z = current.z + w;

			if(next.x >= 0 && next.x < res_x &&
				next.y >= 0 && next.y < res_y &&
				next.z >= 0 && next.z < res_z &&
				function->GetPixel(next.x, next.y, next.z) == old_value)
				{
					function->SetPixel(next.x, next.y, next.z, target_value);
					Q.push(next);
				}
		}
	}

				
}







#endif