#include "./BalancethicknessAnalysis.h"
#include "../toolkits/toolkits.h"
#include "../classes/classes.h"
#include "../shared/shared.h"
#include "../modules/modules.h"

/*Model processing*/
int  BalancethicknessAnalysis::DofsPerNode(int** doflist,int meshtype,int approximation){/*{{{*/
	return 1;
}/*}}}*/
void BalancethicknessAnalysis::UpdateParameters(Parameters* parameters,IoModel* iomodel,int solution_enum,int analysis_enum){/*{{{*/
}/*}}}*/
void BalancethicknessAnalysis::UpdateElements(Elements* elements,IoModel* iomodel,int analysis_counter,int analysis_type){/*{{{*/

	int    stabilization,finiteelement;

	/*Fetch data needed: */
	iomodel->Constant(&stabilization,BalancethicknessStabilizationEnum);

	/*Finite element type*/
	finiteelement = P1Enum;
	if(stabilization==3){
		finiteelement = P1DGEnum;
	}

	/*Update elements: */
	int counter=0;
	for(int i=0;i<iomodel->numberofelements;i++){
		if(iomodel->my_elements[i]){
			Element* element=(Element*)elements->GetObjectByOffset(counter);
			element->Update(i,iomodel,analysis_counter,analysis_type,finiteelement);
			counter++;
		}
	}

	iomodel->FetchDataToInput(elements,ThicknessEnum);
	iomodel->FetchDataToInput(elements,SurfaceEnum);
	iomodel->FetchDataToInput(elements,BaseEnum);
	iomodel->FetchDataToInput(elements,MaskIceLevelsetEnum);
	iomodel->FetchDataToInput(elements,VxEnum);
	iomodel->FetchDataToInput(elements,VyEnum);
	iomodel->FetchDataToInput(elements,BasalforcingsMeltingRateEnum);
	iomodel->FetchDataToInput(elements,SurfaceforcingsMassBalanceEnum);
	iomodel->FetchDataToInput(elements,BalancethicknessThickeningRateEnum);

	if(iomodel->meshtype==Mesh3DEnum){
		iomodel->FetchDataToInput(elements,MeshElementonbedEnum);
		iomodel->FetchDataToInput(elements,MeshElementonsurfaceEnum);
	}
}/*}}}*/
void BalancethicknessAnalysis::CreateNodes(Nodes* nodes,IoModel* iomodel){/*{{{*/

	int  stabilization;
	iomodel->Constant(&stabilization,BalancethicknessStabilizationEnum);

	/*Check in 3d*/
	if(stabilization==3 && iomodel->meshtype==Mesh3DEnum) _error_("DG 3d not implemented yet");

	/*First fetch data: */
	if(iomodel->meshtype==Mesh3DEnum) iomodel->FetchData(2,MeshVertexonbedEnum,MeshVertexonsurfaceEnum);
	if(stabilization!=3){
		::CreateNodes(nodes,iomodel,BalancethicknessAnalysisEnum,P1Enum);
	}
	else{
		::CreateNodes(nodes,iomodel,BalancethicknessAnalysisEnum,P1DGEnum);
	}
	iomodel->DeleteData(2,MeshVertexonbedEnum,MeshVertexonsurfaceEnum);
}/*}}}*/
void BalancethicknessAnalysis::CreateConstraints(Constraints* constraints,IoModel* iomodel){/*{{{*/

	/*Fetch parameters: */
	int    stabilization;	
	iomodel->Constant(&stabilization,BalancethicknessStabilizationEnum);

	/*Do not add constraints in DG*/
	if(stabilization!=3){
		IoModelToConstraintsx(constraints,iomodel,BalancethicknessSpcthicknessEnum,BalancethicknessAnalysisEnum,P1Enum);
	}

}/*}}}*/
void BalancethicknessAnalysis::CreateLoads(Loads* loads, IoModel* iomodel){/*{{{*/

	/*Intermediary*/
	int element;
	int stabilization;

	/*Fetch parameters: */
	iomodel->Constant(&stabilization,BalancethicknessStabilizationEnum);

	/*Loads only in DG*/
	if (stabilization==3){

		/*Get faces and elements*/
		CreateFaces(iomodel);
		iomodel->FetchData(1,ThicknessEnum);

		/*First load data:*/
		for(int i=0;i<iomodel->numberoffaces;i++){

			/*Get left and right elements*/
			element=iomodel->faces[4*i+2]-1; //faces are [node1 node2 elem1 elem2]

			/*Now, if this element is not in the partition, pass: */
			if(!iomodel->my_elements[element]) continue;

			/* Add load */
			loads->AddObject(new Numericalflux(iomodel->loadcounter+i+1,i,i,iomodel,BalancethicknessAnalysisEnum));
		}

		/*Free data: */
		iomodel->DeleteData(1,ThicknessEnum);
	}
}/*}}}*/

/*Finite Element Analysis*/
void           BalancethicknessAnalysis::Core(FemModel* femmodel){/*{{{*/
	_error_("not implemented");
}/*}}}*/
ElementVector* BalancethicknessAnalysis::CreateDVector(Element* element){/*{{{*/
	/*Default, return NULL*/
	return NULL;
}/*}}}*/
ElementMatrix* BalancethicknessAnalysis::CreateJacobianMatrix(Element* element){/*{{{*/
_error_("Not implemented");
}/*}}}*/
ElementMatrix* BalancethicknessAnalysis::CreateKMatrix(Element* element){/*{{{*/

	if(!element->IsOnBed()) return NULL;
	Element* basalelement = element->SpawnBasalElement();

	ElementMatrix* Ke = NULL;
	switch(element->FiniteElement()){
		case P1Enum: case P2Enum:
			Ke = CreateKMatrixCG(basalelement);
			break;
		case P1DGEnum:
			Ke = CreateKMatrixDG(basalelement);
			break;
		default:
			_error_("Element type " << EnumToStringx(element->FiniteElement()) << " not supported yet");
	}

	int meshtype;
	element->FindParam(&meshtype,MeshTypeEnum);
	if(meshtype!=Mesh2DhorizontalEnum){basalelement->DeleteMaterials(); delete basalelement;};
	return Ke;
}/*}}}*/
ElementMatrix* BalancethicknessAnalysis::CreateKMatrixCG(Element* element){/*{{{*/

	/*Intermediaries */
	int        stabilization;
	int        meshtype;
	IssmDouble Jdet,D_scalar,h;
	IssmDouble vel,vx,vy,dvxdx,dvydy;
	IssmDouble dvx[2],dvy[2];
	IssmDouble* xyz_list = NULL;

	/*Fetch number of nodes and dof for this finite element*/
	int numnodes = element->GetNumberOfNodes();

	/*Initialize Element vector and other vectors*/
	ElementMatrix* Ke     = element->NewElementMatrix();
	IssmDouble*    B      = xNew<IssmDouble>(2*numnodes);
	IssmDouble*    Bprime = xNew<IssmDouble>(2*numnodes);
	IssmDouble     D[2][2];

	/*Retrieve all inputs and parameters*/
	element->GetVerticesCoordinates(&xyz_list);
	element->FindParam(&meshtype,MeshTypeEnum);
	element->FindParam(&stabilization,BalancethicknessStabilizationEnum);
	Input* vxaverage_input=NULL;
	Input* vyaverage_input=NULL;
	if(meshtype==Mesh2DhorizontalEnum){
		vxaverage_input=element->GetInput(VxEnum); _assert_(vxaverage_input);
		vyaverage_input=element->GetInput(VyEnum); _assert_(vyaverage_input);
	}
	else{
		vxaverage_input=element->GetInput(VxAverageEnum); _assert_(vxaverage_input);
		vyaverage_input=element->GetInput(VyAverageEnum); _assert_(vyaverage_input);
	}
	h = element->CharacteristicLength();

	/* Start  looping on the number of gaussian points: */
	Gauss* gauss=element->NewGauss(2);
	for(int ig=gauss->begin();ig<gauss->end();ig++){
		gauss->GaussPoint(ig);

		element->JacobianDeterminant(&Jdet,xyz_list,gauss);
		GetB(B,element,xyz_list,gauss);
		GetBprime(Bprime,element,xyz_list,gauss);

		vxaverage_input->GetInputValue(&vx,gauss);
		vyaverage_input->GetInputValue(&vy,gauss);
		vxaverage_input->GetInputDerivativeValue(&dvx[0],xyz_list,gauss);
		vyaverage_input->GetInputDerivativeValue(&dvy[0],xyz_list,gauss);
		dvxdx=dvx[0];
		dvydy=dvy[1];
		D_scalar=gauss->weight*Jdet;

		D[0][0]=D_scalar*dvxdx;
		D[0][1]=0.;
		D[1][0]=0.;
		D[1][1]=D_scalar*dvydy;
		TripleMultiply(B,2,numnodes,1,
					&D[0][0],2,2,0,
					B,2,numnodes,0,
					&Ke->values[0],1);

		D[0][0]=D_scalar*vx;
		D[1][1]=D_scalar*vy;
		TripleMultiply(B,2,numnodes,1,
					&D[0][0],2,2,0,
					Bprime,2,numnodes,0,
					&Ke->values[0],1);

		if(stabilization==1){
			/*Streamline upwinding*/
			vel=sqrt(vx*vx+vy*vy);
			D[0][0]=h/(2*vel)*vx*vx;
			D[1][0]=h/(2*vel)*vy*vx;
			D[0][1]=h/(2*vel)*vx*vy;
			D[1][1]=h/(2*vel)*vy*vy;
		}
		else if(stabilization==2){
			/*SSA*/
			vxaverage_input->GetInputAverage(&vx);
			vyaverage_input->GetInputAverage(&vy);
			D[0][0]=h/2.0*fabs(vx);
			D[0][1]=0.;
			D[1][0]=0.;
			D[1][1]=h/2.0*fabs(vy);
		}
		if(stabilization==1 || stabilization==2){
			D[0][0]=D_scalar*D[0][0];
			D[1][0]=D_scalar*D[1][0];
			D[0][1]=D_scalar*D[0][1];
			D[1][1]=D_scalar*D[1][1];
			TripleMultiply(Bprime,2,numnodes,1,
						&D[0][0],2,2,0,
						Bprime,2,numnodes,0,
						&Ke->values[0],1);
		}
	}

	/*Clean up and return*/
	xDelete<IssmDouble>(xyz_list);
	xDelete<IssmDouble>(B);
	xDelete<IssmDouble>(Bprime);
	delete gauss;
	return Ke;
}/*}}}*/
ElementMatrix* BalancethicknessAnalysis::CreateKMatrixDG(Element* element){/*{{{*/

	/*Intermediaries */
	int        meshtype;
	IssmDouble Jdet,D_scalar,vx,vy,dvxdx,dvydy,vel;
	IssmDouble dvx[2],dvy[2];
	IssmDouble* xyz_list = NULL;

	/*Fetch number of nodes and dof for this finite element*/
	int numnodes = element->GetNumberOfNodes();

	/*Initialize Element vector and other vectors*/
	ElementMatrix* Ke     = element->NewElementMatrix();
	IssmDouble*    B      = xNew<IssmDouble>(2*numnodes);
	IssmDouble*    Bprime = xNew<IssmDouble>(2*numnodes);
	IssmDouble     D[2][2];

	/*Retrieve all inputs and parameters*/
	element->GetVerticesCoordinates(&xyz_list);
	element->FindParam(&meshtype,MeshTypeEnum);
	Input* vxaverage_input=NULL;
	Input* vyaverage_input=NULL;
	if(meshtype==Mesh2DhorizontalEnum){
		vxaverage_input=element->GetInput(VxEnum); _assert_(vxaverage_input);
		vyaverage_input=element->GetInput(VyEnum); _assert_(vyaverage_input);
	}
	else{
		vxaverage_input=element->GetInput(VxAverageEnum); _assert_(vxaverage_input);
		vyaverage_input=element->GetInput(VyAverageEnum); _assert_(vyaverage_input);
	}

	/* Start  looping on the number of gaussian points: */
	Gauss* gauss=element->NewGauss(2);
	for(int ig=gauss->begin();ig<gauss->end();ig++){
		gauss->GaussPoint(ig);

		element->JacobianDeterminant(&Jdet,xyz_list,gauss);
		vxaverage_input->GetInputValue(&vx,gauss);
		vyaverage_input->GetInputValue(&vy,gauss);
		vxaverage_input->GetInputDerivativeValue(&dvx[0],xyz_list,gauss);
		vyaverage_input->GetInputDerivativeValue(&dvy[0],xyz_list,gauss);
		D_scalar=gauss->weight*Jdet;

		/*WARNING: B and Bprime are inverted compared to CG*/
		GetB(Bprime,element,xyz_list,gauss);
		GetBprime(B,element,xyz_list,gauss);

		D_scalar = - gauss->weight*Jdet;
		D[0][0]  = D_scalar*vx;
		D[0][1]  = 0.;
		D[1][0]  = 0.;
		D[1][1]  = D_scalar*vy;
		TripleMultiply(B,2,numnodes,1,
					&D[0][0],2,2,0,
					Bprime,2,numnodes,0,
					&Ke->values[0],1);

	}

	/*Clean up and return*/
	xDelete<IssmDouble>(xyz_list);
	xDelete<IssmDouble>(B);
	xDelete<IssmDouble>(Bprime);
	delete gauss;
	return Ke;
}/*}}}*/
ElementVector* BalancethicknessAnalysis::CreatePVector(Element* element){/*{{{*/

	if(!element->IsOnBed()) return NULL;
	Element* basalelement = element->SpawnBasalElement();

	ElementVector* pe = NULL;
	switch(element->FiniteElement()){
		case P1Enum: case P2Enum:
			pe = CreatePVectorCG(basalelement);
			break;
		case P1DGEnum:
			pe = CreatePVectorDG(basalelement);
			break;
		default:
			_error_("Element type " << EnumToStringx(element->FiniteElement()) << " not supported yet");
	}

	int meshtype;
	element->FindParam(&meshtype,MeshTypeEnum);
	if(meshtype!=Mesh2DhorizontalEnum){basalelement->DeleteMaterials(); delete basalelement;};
	return pe;
}/*}}}*/
ElementVector* BalancethicknessAnalysis::CreatePVectorCG(Element* element){/*{{{*/

	/*Intermediaries */
	IssmDouble  dhdt,mb,ms,Jdet;
	IssmDouble* xyz_list = NULL;

	/*Fetch number of nodes and dof for this finite element*/
	int numnodes = element->GetNumberOfNodes();

	/*Initialize Element vector and other vectors*/
	ElementVector* pe    = element->NewElementVector();
	IssmDouble*    basis = xNew<IssmDouble>(numnodes);

	/*Retrieve all inputs and parameters*/
	element->GetVerticesCoordinates(&xyz_list);
	Input* mb_input   = element->GetInput(BasalforcingsMeltingRateEnum);       _assert_(mb_input);
	Input* ms_input   = element->GetInput(SurfaceforcingsMassBalanceEnum);     _assert_(ms_input);
	Input* dhdt_input = element->GetInput(BalancethicknessThickeningRateEnum); _assert_(dhdt_input);

	/*Initialize mb_correction to 0, do not forget!:*/
	/* Start  looping on the number of gaussian points: */
	Gauss* gauss=element->NewGauss(2);
	for(int ig=gauss->begin();ig<gauss->end();ig++){
		gauss->GaussPoint(ig);

		element->JacobianDeterminant(&Jdet,xyz_list,gauss);
		element->NodalFunctions(basis,gauss);

		ms_input->GetInputValue(&ms,gauss);
		mb_input->GetInputValue(&mb,gauss);
		dhdt_input->GetInputValue(&dhdt,gauss);

		for(int i=0;i<numnodes;i++) pe->values[i]+=Jdet*gauss->weight*(ms-mb-dhdt)*basis[i];
	}

	/*Clean up and return*/
	xDelete<IssmDouble>(xyz_list);
	xDelete<IssmDouble>(basis);
	delete gauss;
	return pe;
}/*}}}*/
ElementVector* BalancethicknessAnalysis::CreatePVectorDG(Element* element){/*{{{*/

	/*Intermediaries */
	IssmDouble  dhdt,mb,ms,Jdet;
	IssmDouble* xyz_list = NULL;

	/*Fetch number of nodes and dof for this finite element*/
	int numnodes = element->GetNumberOfNodes();

	/*Initialize Element vector and other vectors*/
	ElementVector* pe    = element->NewElementVector();
	IssmDouble*    basis = xNew<IssmDouble>(numnodes);

	/*Retrieve all inputs and parameters*/
	element->GetVerticesCoordinates(&xyz_list);
	Input* mb_input   = element->GetInput(BasalforcingsMeltingRateEnum);       _assert_(mb_input);
	Input* ms_input   = element->GetInput(SurfaceforcingsMassBalanceEnum);     _assert_(ms_input);
	Input* dhdt_input = element->GetInput(BalancethicknessThickeningRateEnum); _assert_(dhdt_input);

	/*Initialize mb_correction to 0, do not forget!:*/
	/* Start  looping on the number of gaussian points: */
	Gauss* gauss=element->NewGauss(2);
	for(int ig=gauss->begin();ig<gauss->end();ig++){
		gauss->GaussPoint(ig);

		element->JacobianDeterminant(&Jdet,xyz_list,gauss);
		element->NodalFunctions(basis,gauss);

		ms_input->GetInputValue(&ms,gauss);
		mb_input->GetInputValue(&mb,gauss);
		dhdt_input->GetInputValue(&dhdt,gauss);

		for(int i=0;i<numnodes;i++) pe->values[i]+=Jdet*gauss->weight*(ms-mb-dhdt)*basis[i];
	}

	/*Clean up and return*/
	xDelete<IssmDouble>(xyz_list);
	xDelete<IssmDouble>(basis);
	delete gauss;
	return pe;
}/*}}}*/
void BalancethicknessAnalysis::GetB(IssmDouble* B,Element* element,IssmDouble* xyz_list,Gauss* gauss){/*{{{*/
	/*Compute B  matrix. B=[B1 B2 B3] where Bi is of size 3*NDOF2. 
	 * For node i, Bi can be expressed in the actual coordinate system
	 * by: 
	 *       Bi=[ N ]
	 *          [ N ]
	 * where N is the finiteelement function for node i.
	 *
	 * We assume B_prog has been allocated already, of size: 2x(NDOF1*numnodes)
	 */

	/*Fetch number of nodes for this finite element*/
	int numnodes = element->GetNumberOfNodes();

	/*Get nodal functions*/
	IssmDouble* basis=xNew<IssmDouble>(numnodes);
	element->NodalFunctions(basis,gauss);

	/*Build B: */
	for(int i=0;i<numnodes;i++){
		B[numnodes*0+i] = basis[i];
		B[numnodes*1+i] = basis[i];
	}

	/*Clean-up*/
	xDelete<IssmDouble>(basis);
}/*}}}*/
void BalancethicknessAnalysis::GetBprime(IssmDouble* Bprime,Element* element,IssmDouble* xyz_list,Gauss* gauss){/*{{{*/
	/*Compute B'  matrix. B'=[B1' B2' B3'] where Bi' is of size 3*NDOF2. 
	 * For node i, Bi' can be expressed in the actual coordinate system
	 * by: 
	 *       Bi_prime=[ dN/dx ]
	 *                [ dN/dy ]
	 * where N is the finiteelement function for node i.
	 *
	 * We assume B' has been allocated already, of size: 3x(NDOF2*numnodes)
	 */

	/*Fetch number of nodes for this finite element*/
	int numnodes = element->GetNumberOfNodes();

	/*Get nodal functions derivatives*/
	IssmDouble* dbasis=xNew<IssmDouble>(2*numnodes);
	element->NodalFunctionsDerivatives(dbasis,xyz_list,gauss);

	/*Build B': */
	for(int i=0;i<numnodes;i++){
		Bprime[numnodes*0+i] = dbasis[0*numnodes+i];
		Bprime[numnodes*1+i] = dbasis[1*numnodes+i];
	}

	/*Clean-up*/
	xDelete<IssmDouble>(dbasis);

}/*}}}*/
void BalancethicknessAnalysis::GetSolutionFromInputs(Vector<IssmDouble>* solution,Element* element){/*{{{*/
	   _error_("not implemented yet");
}/*}}}*/
void BalancethicknessAnalysis::InputUpdateFromSolution(IssmDouble* solution,Element* element){/*{{{*/

	int meshtype;
	element->FindParam(&meshtype,MeshTypeEnum);
	switch(meshtype){
		case Mesh2DhorizontalEnum:
			element->InputUpdateFromSolutionOneDof(solution,ThicknessEnum);
			break;
		case Mesh3DEnum:
			element->InputUpdateFromSolutionOneDofCollapsed(solution,ThicknessEnum);
			break;
		default: _error_("mesh "<<EnumToStringx(meshtype)<<" not supported yet");
	}
}/*}}}*/
void BalancethicknessAnalysis::UpdateConstraints(FemModel* femmodel){/*{{{*/
	/*Default, do nothing*/
	return;
}/*}}}*/
