#ifndef STIM_TENSOR3_H
#define STIM_TENSOR3_H

#include "matrix_sym.h"
#include <stim/math/constants.h>

namespace stim {

	/*This class represents a symmetric rank-2 2D tensor, useful for structure tensors
	*/

	//Matrix ID cheat sheet
	//	| 0  1  2 |
	//	| 1  3  4 |
	//	| 2  4  5 |
	template<typename T>
	class tensor3 : public matrix_sym<T, 3> {

	protected:

	public:

		//calculates the determinant of the tensor
		CUDA_CALLABLE T det() {
			return M[0] * M[3] * M[5] + 2 * (M[1] * M[4] * M[2]) - M[2] * M[3] * M[2] - M[1] * M[1] * M[5] - M[0] * M[4] * M[4];
		}

		//calculate the eigenvalues for the tensor
		//adapted from https://en.wikipedia.org/wiki/Eigenvalue_algorithm

		CUDA_CALLABLE stim::vec3<T> lambda() {
			stim::vec3<T> lam;
			T p1 = M[1] * M[1] + M[2] * M[2] + M[4] * M[4];		//calculate the sum of the squared off-diagonal values
			if (p1 == 0) {										//if this value is zero, the matrix is diagonal
				lam[0] = M[0];									//the eigenvalues are the diagonal values
				lam[1] = M[3];
				lam[2] = M[5];
				return lam;										//return the eigenvalue vector
			}

			T tr = matrix_sym<T, 3>::trace();					//calculate the trace of the matrix
			T q = tr / 3;
			T p2 = (M[0] - q) * (M[0] - q) + (M[3] - q) * (M[3] - q) + (M[5] - q) * (M[5] - q) + 2 * p1;
			T p = sqrt(p2 / 6);
			tensor3<T> Q;										//allocate space for Q (q along the diagonals)
			Q = (T)0;											//initialize Q to zeros
			Q(0, 0) = Q(1, 1) = Q(2, 2) = q;					//set the diagonal values to q
			tensor3<T> B = *this;								// B1 = A
			B.M[0] = (B.M[0] - q);
			B.M[3] = (B.M[3] - q);
			B.M[5] = (B.M[5] - q);						
			matrix_sym<T, 3>::operator_product(B, 1/p);				// B = (1/p) * (A - q*I)
			//B.M[0] = B.M[0] * 1/p;
			//B.M[1] = B.M[1] * 1/p;
			//B.M[2] = B.M[2] * 1/p;
			//B.M[3] = B.M[3] * 1/p;
			//B.M[4] = B.M[4] * 1/p;
			//B.M[5] = B.M[5] * 1/p;
			T r = B.det() / 2;									//calculate det(B) / 2

			// In exact arithmetic for a symmetric matrix - 1 <= r <= 1
			// but computation error can leave it slightly outside this range.
			T phi;
			if (r <= -1) phi = stim::PI / 3;
			else if (r >= 1) phi = 0;
			else phi = acos(r) / 3;
			
			// the eigenvalues satisfy eig3 >= eig2 >= eig1
			lam[2] = q + 2 * p * cos(phi);
			lam[0] = q + 2 * p * cos(phi + (2 * stim::PI / 3));
			lam[1] = 3 * q - (lam[2] + lam[0]);

			return lam;
		}

		CUDA_CALLABLE stim::matrix<T> eig(stim::vec3<T>& lambda = stim::vec3<T>()) {
			stim::matrix<T> V;
		
			stim::matrix<T> M1 = matrix_sym<T, 3>::mat();
			stim::matrix<T> M2 = matrix_sym<T, 3>::mat();
			stim::matrix<T> M3 = matrix_sym<T, 3>::mat();	// fill a tensor with symmetric values

			M1 = M1 - lambda[0];					// M1 = A - lambda[0] * I
	
			M2 = M2 - lambda[1];					// M2 = A - lambda[1] * I

			M3 = M3 - lambda[2];					// M3 = A - lambda[2] * I

			T Mod = 0;											// module of one column

			T tmp1[9] = {0};
			for(int i = 0; i < 9; i++) {
				for(int j = 0; j < 3; j++){
					tmp1[i] += M2(i%3, j) * M3(j, i/3); 
				}
			}
			if(tmp1[0] * tmp1[1] * tmp1[2] != 0) {				// test whether it is zero column
				Mod = sqrt(pow(tmp1[0],2) + pow(tmp1[1],2) + pow(tmp1[2],2));
				V(0, 0) = tmp1[0]/Mod;
				V(1, 0) = tmp1[1]/Mod;
				V(2, 0) = tmp1[2]/Mod;
			}
			else {
				Mod = sqrt(pow(tmp1[3],2) + pow(tmp1[4],2) + pow(tmp1[5],2));
				V(0, 0) = tmp1[3]/Mod;
				V(1, 0) = tmp1[4]/Mod;
				V(2, 0) = tmp1[5]/Mod;
			}

			T tmp2[9] = {0};
			for(int i = 0; i < 9; i++) {
				for(int j = 0; j < 3; j++){
					tmp2[i] += M1(i%3, j) * M3(j, i/3); 
				}
			}
			if(tmp2[0] * tmp2[1] * tmp2[2] != 0) {
				Mod = sqrt(pow(tmp2[0],2) + pow(tmp2[1],2) + pow(tmp2[2],2));
				V(0, 1) = tmp2[0]/Mod;
				V(1, 1) = tmp2[1]/Mod;
				V(2, 1) = tmp2[2]/Mod;
			}
			else {
				Mod = sqrt(pow(tmp2[3],2) + pow(tmp2[4],2) + pow(tmp2[5],2));
				V(0, 1) = tmp2[3]/Mod;
				V(1, 1) = tmp2[4]/Mod;
				V(2, 1) = tmp2[5]/Mod;
			}

			T tmp3[9] = {0};
			for(int i = 0; i < 9; i++) {
				for(int j = 0; j < 3; j++){
					tmp3[i] += M1(i%3, j) * M2(j, i/3); 
				}
			}
			if(tmp3[0] * tmp3[1] * tmp3[2] != 0) {
				Mod = sqrt(pow(tmp3[0],2) + pow(tmp3[1],2) + pow(tmp3[2],2));
				V(0, 2) = tmp3[0]/Mod;
				V(1, 2) = tmp3[1]/Mod;
				V(2, 2) = tmp3[2]/Mod;
			}
			else {
				Mod = sqrt(pow(tmp3[3],2) + pow(tmp3[4],2) + pow(tmp3[5],2));
				V(0, 2) = tmp3[3]/Mod;
				V(1, 2) = tmp3[4]/Mod;
				V(2, 2) = tmp3[5]/Mod;
			}
			return V;					//return the eigenvector matrix
		}
		// return one specific eigenvector
		CUDA_CALLABLE stim::vec3<T> eig(int n, stim::vec3<T>& lambda = stim::vec3<T>()) {
			stim::matrix<T, 3> V = eig(lambda);
			stim::vec3<T> v;
			for(int i = 0; i < 3; i++)
				v[i] = V(i, n);
			return v;
		}


		CUDA_CALLABLE T linear(stim::vec3<T>& lambda = stim::vec3<T>()) {
			T cl = (lambda[2] - lambda[1]) / (lambda[0] + lambda[1] + lambda[2]);
			return cl;
		}

		CUDA_CALLABLE T Planar(stim::vec3<T>& lambda = stim::vec3<T>()) {
			T cp = 2 * (lambda[1] - lambda[0]) / (lambda[0] + lambda[1] + lambda[2]);
			return cp;
		}

		CUDA_CALLABLE T spherical(stim::vec3<T>& lambda = stim::vec3<T>()) {
			T cs = 3 * lambda[0] / (lambda[0] + lambda[1] + lambda[2]);
			return cs;
		}

		CUDA_CALLABLE T fa(stim::vec3<T>& lambda = stim::vec3<T>()) {
			T fa = sqrt(1/2) * sqrt(pow(lambda[2] - lambda[1], 2) + pow(lambda[1] - lambda[0], 2) + pow(lambda[0] - lambda[2], 2)) / sqrt(pow(lambda[2], 2) + pow(lambda[1], 2) + pow(lambda[0], 2));
		}
		//JACK 2: write functions to calculate anisotropy
		//ex: fa(), linear(), planar(), spherical()


		//calculate the eigenvectors and eigenvalues of the tensor
		//CUDA_CALLABLE void eig(stim::matrix<T, 3>& v, stim::matrix<T, 3>& lambda){

		//}
		CUDA_CALLABLE tensor3<T> operator=(T rhs) {
			stim::matrix_sym<T, 3>::operator=(rhs);
			return *this;
		}

		CUDA_CALLABLE tensor3<T> operator=(stim::matrix_sym<T, 3> rhs) {
			stim::matrix_sym<T, 3>::operator=(rhs);
			return *this;
		}
	};


}	//end namespace stim


#endif