#ifndef JACK_CENTERLINE_H
#define JACK_CENTERLINE_H

#include <stdlib.h>
#include <vector>
#include <stim/math/vec3.h>
#include <stim/structures/kdtree.cuh>

namespace stim {

	/// we always assume that one centerline has a flow direction even it actually does not have. Also, we allow loop centerline
	/// NOTE: centerline is not derived from std::vector<stim::vec3<T>> class!!!
	template<typename T>
	class centerline {

	private:
		size_t n;								// number of points on the centerline, can be zero if NULL

		// update length information at each point (distance from starting point) starting from index "start"
		void update_L(size_t start = 0) {
			L.resize(n);

			if (start == 0) {
				L[0] = 0.0f;
				start++;
			}

			stim::vec3<T> dir;					// temp direction vector for calculating length between two points
			for (size_t i = start; i < n; i++) {
				dir = C[i] - C[i - 1];			// calculate the certerline extending direction
				L[i] = L[i - 1] + dir.len();	// addition
			}
		}

	protected:
		std::vector<stim::vec3<T> > C;			// points on the centerline
		std::vector<T> L;						// stores the integrated length along the fiber

	public:
		/// constructors
		// empty constructor
		centerline() {
			n = 0;
		}

		// constructor that allocate memory
		centerline(size_t s) {
			n = s;
			C.resize(s);		// allocate memory for points
			L.resize(s);		// allocate memory for lengths

			update_L();
		}

		// constructor that constructs a centerline based on a list of points
		centerline(std::vector<stim::vec3<T> > rhs) {
			n = rhs.size();						// get the number of points
			C.resize(n);
			for (size_t i = 0; i < n; i++)
				C[i] = rhs[i];					// copy data
			update_L();
		}


		/// vector operations
		// add a new point to current centerline
		void push_back(stim::vec3<T> p) {
			C.push_back(p);
			n++;								// increase the number of points
			update_L(n - 1);
		}

		// insert a new point at specific location to current centerline
		void insert(typename std::vector<stim::vec3<T> >::iterator pos, stim::vec3<T> p) {
			C.insert(pos, p);									// insert a new point
			n++;
			size_t d = std::distance(C.begin(), pos);			// get the index
			update_L(d);
		}
		// insert a new point at C[idx]
		void insert(size_t idx, stim::vec3<T> p) {
			n++;
			C.resize(n);
			for (size_t i = n - 1; i > idx; i--)				// move point location
				C[i] = C[i - 1];
			C[idx] = p;
			update_L(idx);
		}

		// assign a point at specific location to current centerline
		void assign(size_t idx, stim::vec3<T> p) {
			C[idx] = p;
			update_L(idx);
		}

		// erase a point at specific location on current centerline
		void erase(typename std::vector<stim::vec3<T> >::iterator pos) {
			C.erase(pos);										// erase a point
			n--;
			size_t d = std::distance(C.begin(), pos);			// get the index
			update_L(d);
		}
		// erase a point at C[idx]
		void erase(size_t idx) {
			n--;
			for (size_t i = idx; i < n; i++)
				C[i] = C[i + 1];
			C.resize(n);
			update_L(idx);
		}

		// clear up all the points
		void clear() {
			C.clear();			// clear list
			L.clear();			// clear length information
			n = 0;				// set number to zero
		}

		// reverse current centerline in terms of points order
		centerline<T> reverse() {
			centerline<T> result = *this;

			std::reverse(result.C.begin(), result.C.end());
			result.update_L();

			return result;
		}

		/// functions for reading centerline information
		// return the number of points on current centerline
		size_t size() {
			return n;
		}

		// return the length
		T length() {
			return L.back();
		}

		// finds the index of the point closest to the length "l" on the lower bound
		size_t findIdx(T l) {
			for (size_t i = 1; i < L.size(); i++) {
				if (L[i] > l)
					return i - 1;
			}

			return L.size() - 1;
		}

		// get a position vector at the given length into the fiber (based on the pvalue), interpolate
		stim::vec3<T> p(T l, size_t idx) {
			T rate = (l - L[idx]) / (L[idx + 1] - L[idx]);
			stim::vec3<T> v1 = C[idx];
			stim::vec3<T> v2 = C[idx + 1];
		
			return (v1 + (v2 - v1) * rate);
		}
		// get a position vector at the given pvalue(pvalue[0.0f, 1.0f])
		stim::vec3<T> p(T pvalue) {
			// degenerated cases
			if (pvalue <= 0.0f) return C[0];
			if (pvalue >= 1.0f) return C.back();
		
			T l = pvalue * L.back();		// get the length based on the given pvalue
			size_t idx = findIdx(l);
			
			return p(l, idx);				// interpolation
		}

		// get the normalized direction vector at point idx (average of the incoming and outgoing directions)
		stim::vec3<T> d(size_t idx) {
			if (n <= 1) return stim::vec3<T>(0.0f, 0.0f, 0.0f);		// if there is insufficient information to calculate the direction, return null
			if (n == 2) return (C[1] - C[0]).norm();				// if there are only two points, the direction vector at both is the direction of the line segment
			
			// degenerate cases at two ends
			if (idx == 0) return (C[1] - C[0]).norm();				// the first direction vector is oriented towards the first line segment
			if (idx == n - 1) return (C[n - 1] - C[n - 2]).norm();	// the last direction vector is oriented towards the last line segment

			// all other direction vectors are the average direction of the two joined line segments
			stim::vec3<T> a = C[idx] - C[idx - 1];
			stim::vec3<T> b = C[idx + 1] - C[idx];
			stim::vec3<T> ab = a.norm() + b.norm();
			return ab.norm();
		}


		/// arithmetic operations
		// '=' operation
		centerline<T> & operator=(centerline<T> rhs) {
			n = rhs.n;
			C = rhs.C;
			L = rhs.L;

			return *this;
		}

		// "[]" operation
		stim::vec3<T> & operator[](size_t idx) {
			return C[idx];
		}

		// "==" operation
		bool operator==(centerline<T> rhs) const {

			if (n != rhs.size())
				return false;
			else {
				size_t num = rhs.size();
				stim::vec3<T> tmp;			// weird situation that I can only use tmp instead of C itself in comparison
				for (size_t i = 0; i < num; i++) {
					stim::vec3<T> tmp = C[i];
					if (tmp[0] != rhs[i][0] || tmp[1] != rhs[i][1] || tmp[2] != rhs[i][2])
						return false;
				}
				return true;
			}
		}

		// "+" operation
		centerline<T> operator+(stim::vec3<T> p) const {
			centerline<T> result(*this);
			result.C.push_back(p);
			result.n++;
			result.update_L(n - 1);
			
			return result;
		}
		centerline<T> operator+(centerline<T> c) const {
			centerline<T> result(*this);
			size_t num1 = result.size();
			size_t num2 = c.size();
			for (size_t i = 0; i < num2; i++)
				result.push_back(c[i]);
			result.update_L(num1);

			return result;
		}

		
		/// advanced operation
		// stitch two centerlines if possible (mutual-stitch and self-stitch)
		static std::vector<centerline<T> > stitch(centerline<T> c1, centerline<T> c2, size_t end = 0) {
			std::vector<centerline<T> > result;
			centerline<T> new_centerline;
			stim::vec3<T> new_vertex;
			// ********** for Pavel **********
			// ********** JACK thinks that ultimately we want it AUTOMATEDLY! **********

			// check stitch case
			if (c1 == c2) {		// self-stitch case
				// ***** don't know how it works *****
			}
			else {				// mutual-stitch case
				size_t num1 = c1.size();	// get the numbers of two centerlines
				size_t num2 = c2.size();

				T* reference = (T*)malloc(sizeof(T) * num1 * 3);	// c1 as reference set
				T* query = (T*)malloc(sizeof(T) * num2 * 3);		// c2 as query set
				for (size_t p = 0; p < num1; p++)					// read points
					for (size_t d = 0; d < 3; d++)
						reference[p * 3 + d] = c1[p][d];			// KDTREE is stilla close code, it has its own structure
				for (size_t p = 0; p < num2; p++)					// read points
					for (size_t d = 0; d < 3; d++)
						query[p * 3 + d] = c2[p][d];

				stim::kdtree<T, 3> kdt;
				kdt.create(reference, num1, 5);						// create a tree based on reference set, max_tree_level is set to 5
				
				std::vector<size_t> index(num2);					// stores index results
				std::vector<float> dist(num2);

				kdt.search(query, num2, &index[0], &dist[0]);		// find the nearest neighbor in c1 for c2

				// clear up
				free(reference);
				free(query);

				std::vector<float>::iterator sm = std::min_element(dist.begin(), dist.end());	// find smallest distance
				size_t id = std::distance(dist.begin(), sm);				// find the target index

				size_t id1 = index[id];
				size_t id2 = id;

				// for centerline c1
				bool flag = false;						// flag indicates whether it has already added the bifurcation point to corresponding new centerline
				if (id1 == 0 || id1 == num1 - 1) { 		// splitting bifurcation is on the end
					new_centerline = c1;
					new_vertex = c2[id2];
					if (id1 == 0) {
						new_centerline.insert(0, new_vertex);
						flag = true;
					}
					if (id1 == num1 - 1) {
						new_centerline.push_back(new_vertex);
						flag = true;
					}

					result.push_back(new_centerline);
				}
				else {									// splitting bifurcation is on the centerline
					std::vector<centerline<T> > tmp_centerline = c1.split(id1);
					result = tmp_centerline;
				}

				// for centerline c2
				if (id2 == 0 || id2 == num2 - 1) {		// splitting bidurcation is on the end
					new_centerline = c2;
					if (flag)
						result.push_back(new_centerline);
					else {			// add the bifurcation point to this centerline
						new_vertex = c1[id1];
						if (id2 == 0) {
							new_centerline.insert(0, new_vertex);
							flag = true;
						}
						if (id2 == num2 - 1) {
							new_centerline.push_back(new_vertex);
							flag = true;
						}

						result.push_back(new_centerline);
					}
				}
				else {									// splitting bifurcation is on the centerline
					std::vector<centerline<T> > tmp_centerline = c2.split(id2);
					result.push_back(tmp_centerline[0]);
					result.push_back(tmp_centerline[1]);
				}
			}

			return result;
		}

		// split current centerline at specific position
		std::vector<centerline<T> > split(size_t idx) {
			std::vector<centerline<T> > result;

			// won't split
			if (idx <= 0 || idx >= n - 1) {
				result.resize(1);
				result[0] = *this;				// return current centerline
			}
			// do split
			else {
				size_t n1 = idx + 1;			// vertex idx would appear twice
				size_t n2 = n - idx;			// in total n + 1 points

				centerline<T> tmp;				// temp centerline

				result.resize(2);

				for (size_t i = 0; i < n1; i++)	// first half
					tmp.push_back(C[i]);
				tmp.update_L();
				result[0] = tmp;
				tmp.clear();					// clear up for next computation

				for (size_t i = 0; i < n2; i++)	// second half
					tmp.push_back(C[i + idx]);
				tmp.update_L();
				result[1] = tmp;
			}

			return result;
		}

		// resample current centerline
		centerline<T> resample(T spacing) {
			
			stim::vec3<T> dir;				// direction vector
			stim::vec3<T> tmp;				// intermiate point to be added
			stim::vec3<T> p1;				// starting point
			stim::vec3<T> p2;				// ending point

			centerline<T> result;

			for (size_t i = 0; i < n - 1; i++) {
				p1 = C[i];
				p2 = C[i + 1];
				
				dir = p2 - p1;				// compute the direction of current segment
				T seg_len = dir.len();

				if (seg_len > spacing) {	// current segment can be sampled
					for (T step = 0.0f; step < seg_len; step += spacing) {
						tmp = p1 + dir * (step / seg_len);		// add new point
						result.push_back(tmp);
					}
				}
				else
					result.push_back(p1);	// push back starting point
			}
			result.push_back(p2);			// push back ending point

			return result;
		}
	};
}

#endif