// Amira Abdel-Rahman
// (c) Massachusetts Institute of Technology 2019

////////////3D bar-truss fea/////////////////////
function solveFea(){
	// # determine the global matrices
	var M,K,F;
	[M,K,F]=getMatrices(setup);
	// M.print();
	// K.print();
	// F.print();
	
	// M, K, F = get_matrices(properties)
	
	// [Q,R]=tf.linalg.qr(K);

	var displacements=lusolve(K.arraySync(), F.arraySync(),false);
	// console.log(displacements);

	updateDisplacement(displacements);
	var X= tf.tensor(displacements);//change later to more effecient solver
	console.log("Displacements:")
	X.print();

	// # determine the stresses in each element
	var stresses = get_stresses(setup, displacements);
	updateStresses(stresses.arraySync());

	console.log("Stresses:");
	stresses.print()

	// console.log(setup);
	//TODO CHECK SINGULARITY AND DOFSC

}

function getMatrices(setup){
	

	var M = tf.zeros([setup.ndofs,setup.ndofs]);
	var K = tf.zeros([setup.ndofs,setup.ndofs]);
	  
	for(var ii=0;ii<setup.edges.length;ii++){
		var element=setup.edges[ii];
		var fromPoint= toTf3D(setup.nodes[element.source].position);
		var toPoint= toTf3D(setup.nodes[element.target].position);
		var dofs=[];
		dofs.push(...setup.nodes[element.source].degrees_of_freedom);
		dofs.push(...setup.nodes[element.target].degrees_of_freedom);

		var  element_vector=tf.sub (toPoint,fromPoint);
		
		
		var length   = element_vector.norm();
		var rho = tf.scalar(element.density);
		var area = tf.scalar(element.area);
		var E = tf.scalar(element.stiffness);// youngs modulus
	
		var Cm = tf.div(tf.mul(rho,tf.mul(area,length)),tf.scalar(6.0));//rho * area * length /6.0
		var Ck= tf.div(tf.mul(E,area),length);//E * area / length 
		// Cm.print();

		var m = tf.tensor ([[2,1],[1,2]]) ;
		var k = tf.tensor([[1,-1],[-1,1]]);

		var x_axis=tf.tensor ([1,0,0]);
		var y_axis=tf.tensor ([0,1,0]);
		var z_axis=tf.tensor ([0,0,1]);

		// # find rotated mass and stifness matrices
		var tau = rotation_matrix(element_vector, x_axis,y_axis,z_axis);
		var m_r=tf.matMul(tau,tf.matMul(m,tau) ,true ,false) //tau.T.dot(m).dot(tau)
		var k_r=tf.matMul(tau,tf.matMul(k,tau) ,true ,false) //tau.T.dot(k).dot(tau)

		// # change from element to global coordinate
		var B=tf.zeros([6,setup.ndofs]);
		var B=[]
		for(var i=0;i<6;i++){
			B.push([]);
			for(var j=0;j<setup.ndofs;j++){
				B[i].push(0.0);
			}
		}
		for(var i=0;i<6;i++){
			B[i][dofs[i]]=1.0;
		}
		
		B=tf.tensor(B);
		// B.print()
		
		
		var M_rG=tf.matMul(B,tf.matMul(m_r,B) ,true ,false) //B.T.dot(m_r).dot(B)
		var K_rG=tf.matMul(B,tf.matMul(k_r,B) ,true ,false) //B.T.dot(k_r).dot(B)

		M=tf.add(M,tf.mul(Cm, M_rG));  //+= Cm * M_rG
		K=tf.add(K,tf.mul(Ck, K_rG));  //+= Ck * K_rG

	}
	


	// # construct the force vector
	var F=[];
	for(var i=0;i<setup.nodes.length;i++){
		F.push(setup.nodes[i].force.x);
		F.push(setup.nodes[i].force.y);
		F.push(setup.nodes[i].force.z);
	}
	F=tf.tensor(F);
	// # remove the restrained dofs
	var remove_indices=[];
	for(var i=0;i<setup.nodes.length;i++){
		for(var j=0;j<setup.nodes[i].degrees_of_freedom.length;j++){
			if(setup.nodes[i].restrained_degrees_of_freedom[j]){
				remove_indices.push(setup.nodes[i].degrees_of_freedom[j]);
			}
		}
	}
	// M.print();
	// K.print();
	// F.print();

	for(var i=0;i<2;i++){
		M=tf_delete(M,remove_indices,i);
		K=tf_delete(K,remove_indices,i);
	}
	

	F=tf_delete(F,remove_indices,0);

	// M.print()
	// K.print()
	// F.print()
	return [ M,K,F ]
}

function get_stresses(setup,X){
	//   # find the stresses in each member
	var stresses=[];
	for(var ii=0;ii<setup.edges.length;ii++){
		var element=setup.edges[ii];
		// # find the element geometry
		var fromPoint= toTf3D(setup.nodes[element.source].position);
		var toPoint= toTf3D(setup.nodes[element.target].position);
		var dofs=[];
		dofs.push(...setup.nodes[element.source].degrees_of_freedom);
		dofs.push(...setup.nodes[element.target].degrees_of_freedom);
		var  element_vector=tf.sub (toPoint,fromPoint);

		var x_axis=tf.tensor ([1,0,0]);
		var y_axis=tf.tensor ([0,1,0]);
		var z_axis=tf.tensor ([0,0,1]);

		// # find rotated mass and stifness matrices
		var tau = rotation_matrix(element_vector, x_axis,y_axis,z_axis);

		var global_displacements=tf.tensor([setup.nodes[element.source].displacement.x,setup.nodes[element.source].displacement.y,setup.nodes[element.source].displacement.z,
			setup.nodes[element.target].displacement.x,setup.nodes[element.target].displacement.y,setup.nodes[element.target].displacement.z]) //# todo change

		// # nodal displacement
		var q=tf.dot(tau,global_displacements);
		var q=tf.unstack(q);

		var strain=tf.div(tf.sub(q[1],q[0]),element_vector.norm()); //todo check smarter way to do it
		
		var E = tf.scalar(element.stiffness);
		var stress=tf.mul(E,strain);
		stresses.push(stress);
	} 

	return tf.stack(stresses,0);
}

///////utils////////
function rotation_matrix(element_vector, x_axis, y_axis, z_axis){
	// # find the direction cosines
	var x_proj = direction_cosine(element_vector, x_axis);
	var y_proj = direction_cosine(element_vector, y_axis);
	var z_proj = direction_cosine(element_vector, z_axis);
	var zero=tf.tensor([0]);
	var temp1=tf.concat([x_proj.expandDims(0), y_proj.expandDims(0), z_proj.expandDims(0), zero, zero, zero]);
	var temp2=tf.concat([ zero, zero, zero, x_proj.expandDims(0), y_proj.expandDims(0), z_proj.expandDims(0)]);
	return tf.stack([temp1,temp2]);

}

function direction_cosine(vec1, vec2){
	return tf.div(tf.dot(vec1,vec2) ,tf.mul( vec1.norm() , vec2.norm()));
}