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

/*Model processing*/
int  EnthalpyAnalysis::DofsPerNode(int** doflist,int meshtype,int approximation){/*{{{*/
	return 1;
}/*}}}*/
void EnthalpyAnalysis::UpdateParameters(Parameters* parameters,IoModel* iomodel,int solution_enum,int analysis_enum){/*{{{*/

	int     numoutputs;
	char**  requestedoutputs = NULL;

	parameters->AddObject(iomodel->CopyConstantObject(ThermalStabilizationEnum));
	parameters->AddObject(iomodel->CopyConstantObject(ThermalIsenthalpyEnum));
	parameters->AddObject(iomodel->CopyConstantObject(ThermalIsdynamicbasalspcEnum));

	iomodel->FetchData(&requestedoutputs,&numoutputs,ThermalRequestedOutputsEnum);
	parameters->AddObject(new IntParam(ThermalNumRequestedOutputsEnum,numoutputs));
	if(numoutputs)parameters->AddObject(new StringArrayParam(ThermalRequestedOutputsEnum,requestedoutputs,numoutputs));
	iomodel->DeleteData(&requestedoutputs,numoutputs,ThermalRequestedOutputsEnum);
}/*}}}*/
void EnthalpyAnalysis::UpdateElements(Elements* elements,IoModel* iomodel,int analysis_counter,int analysis_type){/*{{{*/

	bool dakota_analysis;
	bool isenthalpy;

	/*Now, is the model 3d? otherwise, do nothing: */
	if(iomodel->meshtype==Mesh2DhorizontalEnum)return;

	/*Is enthalpy requested?*/
	iomodel->Constant(&isenthalpy,ThermalIsenthalpyEnum);
	if(!isenthalpy) return;

	/*Fetch data needed: */
	iomodel->FetchData(3,TemperatureEnum,WaterfractionEnum,PressureEnum);

	/*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,P1Enum);
			counter++;
		}
	}

	iomodel->Constant(&dakota_analysis,QmuIsdakotaEnum);

	iomodel->FetchDataToInput(elements,ThicknessEnum);
	iomodel->FetchDataToInput(elements,SurfaceEnum);
	iomodel->FetchDataToInput(elements,BedEnum);
	iomodel->FetchDataToInput(elements,FrictionCoefficientEnum);
	iomodel->FetchDataToInput(elements,FrictionPEnum);
	iomodel->FetchDataToInput(elements,FrictionQEnum);
	iomodel->FetchDataToInput(elements,MaskIceLevelsetEnum);
	iomodel->FetchDataToInput(elements,MaskGroundediceLevelsetEnum);
	iomodel->FetchDataToInput(elements,MeshElementonbedEnum);
	iomodel->FetchDataToInput(elements,MeshElementonsurfaceEnum);
	iomodel->FetchDataToInput(elements,FlowequationElementEquationEnum);
	iomodel->FetchDataToInput(elements,MaterialsRheologyBEnum);
	iomodel->FetchDataToInput(elements,MaterialsRheologyNEnum);
	iomodel->FetchDataToInput(elements,PressureEnum);
	iomodel->FetchDataToInput(elements,TemperatureEnum);
	iomodel->FetchDataToInput(elements,WaterfractionEnum);
	iomodel->FetchDataToInput(elements,EnthalpyEnum);
	iomodel->FetchDataToInput(elements,BasalforcingsGeothermalfluxEnum);
	iomodel->FetchDataToInput(elements,WatercolumnEnum);
	iomodel->FetchDataToInput(elements,BasalforcingsMeltingRateEnum);
	iomodel->FetchDataToInput(elements,VxEnum);
	iomodel->FetchDataToInput(elements,VyEnum);
	iomodel->FetchDataToInput(elements,VzEnum);
	InputUpdateFromConstantx(elements,0.,VxMeshEnum);
	InputUpdateFromConstantx(elements,0.,VyMeshEnum);
	InputUpdateFromConstantx(elements,0.,VzMeshEnum);
	if(dakota_analysis){
		elements->InputDuplicate(TemperatureEnum,QmuTemperatureEnum);
		elements->InputDuplicate(BasalforcingsMeltingRateEnum,QmuMeltingEnum);
		elements->InputDuplicate(VxMeshEnum,QmuVxMeshEnum);
		elements->InputDuplicate(VxMeshEnum,QmuVyMeshEnum);
		elements->InputDuplicate(VxMeshEnum,QmuVzMeshEnum);
	}

	/*Free data: */
	iomodel->DeleteData(3,TemperatureEnum,WaterfractionEnum,PressureEnum);
}/*}}}*/
void EnthalpyAnalysis::CreateNodes(Nodes* nodes,IoModel* iomodel){/*{{{*/

	if(iomodel->meshtype==Mesh3DEnum) iomodel->FetchData(2,MeshVertexonbedEnum,MeshVertexonsurfaceEnum);
	::CreateNodes(nodes,iomodel,EnthalpyAnalysisEnum,P1Enum);
	iomodel->DeleteData(2,MeshVertexonbedEnum,MeshVertexonsurfaceEnum);
}/*}}}*/
void EnthalpyAnalysis::CreateConstraints(Constraints* constraints,IoModel* iomodel){/*{{{*/

	/*Intermediary*/
	int        count;
	int        M,N;
	bool       spcpresent = false;
	IssmDouble heatcapacity;
	IssmDouble referencetemperature;

	/*Output*/
	IssmDouble *spcvector  = NULL;
	IssmDouble* times=NULL;
	IssmDouble* values=NULL;

	/*Fetch parameters: */
	iomodel->Constant(&heatcapacity,MaterialsHeatcapacityEnum);
	iomodel->Constant(&referencetemperature,ConstantsReferencetemperatureEnum);

	/*return if 2d mesh*/
	if(iomodel->meshtype==Mesh2DhorizontalEnum) return;

	/*Fetch data: */
	iomodel->FetchData(&spcvector,&M,&N,ThermalSpctemperatureEnum);

	//FIX ME: SHOULD USE IOMODELCREATECONSTRAINTS 
	/*Transient or static?:*/
	if(M==iomodel->numberofvertices){
		/*static: just create Constraints objects*/
		count=0;

		for(int i=0;i<iomodel->numberofvertices;i++){
			/*keep only this partition's nodes:*/
			if((iomodel->my_vertices[i])){

				if (!xIsNan<IssmDouble>(spcvector[i])){

					constraints->AddObject(new SpcStatic(iomodel->constraintcounter+count+1,iomodel->nodecounter+i+1,1,heatcapacity*(spcvector[i]-referencetemperature),EnthalpyAnalysisEnum));
					count++;

				}
			}
		}
	}
	else if (M==(iomodel->numberofvertices+1)){
		/*transient: create transient SpcTransient objects. Same logic, except we need to retrieve 
		 * various times and values to initialize an SpcTransient object: */
		count=0;

		/*figure out times: */
		times=xNew<IssmDouble>(N);
		for(int j=0;j<N;j++){
			times[j]=spcvector[(M-1)*N+j];
		}

		/*Create constraints from x,y,z: */
		for(int i=0;i<iomodel->numberofvertices;i++){

			/*keep only this partition's nodes:*/
			if((iomodel->my_vertices[i])){

				/*figure out times and values: */
				values=xNew<IssmDouble>(N);
				spcpresent=false;
				for(int j=0;j<N;j++){
					values[j]=heatcapacity*(spcvector[i*N+j]-referencetemperature);
					if(!xIsNan<IssmDouble>(values[j]))spcpresent=true; //NaN means no spc by default
				}

				if(spcpresent){
					constraints->AddObject(new SpcTransient(iomodel->constraintcounter+count+1,iomodel->nodecounter+i+1,1,N,times,values,EnthalpyAnalysisEnum));
					count++;
				}
				xDelete<IssmDouble>(values);
			}
		}
	}
	else{
		_error_("Size of field " << EnumToStringx(ThermalSpctemperatureEnum) << " not supported");
	}

	/*Free ressources:*/
	iomodel->DeleteData(spcvector,ThermalSpctemperatureEnum);
	xDelete<IssmDouble>(times);
	xDelete<IssmDouble>(values);
}/*}}}*/
void EnthalpyAnalysis::CreateLoads(Loads* loads, IoModel* iomodel){/*{{{*/

	/*No loads */
}/*}}}*/

/*Finite Element Analysis*/
ElementMatrix* EnthalpyAnalysis::CreateKMatrix(Element* element){/*{{{*/
	_error_("not implemented yet");
}/*}}}*/
ElementVector* EnthalpyAnalysis::CreatePVector(Element* element){/*{{{*/

	/*compute all load vectors for this element*/
	ElementVector* pe1=CreatePVectorVolume(element);
	ElementVector* pe2=CreatePVectorSheet(element);
	ElementVector* pe3=CreatePVectorShelf(element);
	ElementVector* pe =new ElementVector(pe1,pe2,pe3);

	/*clean-up and return*/
	delete pe1;
	delete pe2;
	delete pe3;
	return pe;
}/*}}}*/
ElementVector* EnthalpyAnalysis::CreatePVectorVolume(Element* element){/*{{{*/

	/*Intermediaries*/
	int         stabilization;
	IssmDouble  Jdet,phi,dt;
	IssmDouble  enthalpy;
	IssmDouble  kappa,tau_parameter,diameter;
	IssmDouble  u,v,w;
	IssmDouble  scalar_def,scalar_transient;
	IssmDouble* xyz_list = NULL;

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

	/*Initialize Element vector*/
	ElementVector* pe             = element->NewElementVector();
	IssmDouble*    basis          = xNew<IssmDouble>(numnodes);
	IssmDouble*    dbasis         = xNew<IssmDouble>(3*numnodes);
	IssmDouble*    pressure       = xNew<IssmDouble>(numvertices);
	IssmDouble*    enthalpypicard = xNew<IssmDouble>(numvertices);

	/*Retrieve all inputs and parameters*/
	element->GetVerticesCoordinates(&xyz_list);
	IssmDouble  rho_ice             = element->GetMaterialParameter(MaterialsRhoIceEnum);
	IssmDouble  thermalconductivity = element->GetMaterialParameter(MaterialsThermalconductivityEnum);
	element->FindParam(&dt,TimesteppingTimeStepEnum);
	element->FindParam(&stabilization,ThermalStabilizationEnum);
	Input* vx_input=element->GetInput(VxEnum); _assert_(vx_input);
	Input* vy_input=element->GetInput(VyEnum); _assert_(vy_input);
	Input* vz_input=element->GetInput(VzEnum); _assert_(vz_input);
	Input* enthalpy_input = NULL;
	if(reCast<bool,IssmDouble>(dt)){enthalpy_input = element->GetInput(EnthalpyEnum); _assert_(enthalpy_input);}
	if(stabilization==2){
		element->GetInputListOnVertices(enthalpypicard,EnthalpyPicardEnum);
		element->GetInputListOnVertices(pressure,PressureEnum);
		diameter=element->MinEdgeLength(xyz_list);
	}

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

		element->JacobianDeterminant(&Jdet,xyz_list,gauss);
		element->NodalFunctions(basis,gauss);
		element->ViscousHeating(&phi,xyz_list,gauss,vx_input,vy_input,vz_input);

		scalar_def=phi/rho_ice*Jdet*gauss->weight;
		if(reCast<bool,IssmDouble>(dt)) scalar_def=scalar_def*dt;

		/*TODO: add -beta*laplace T_m(p)*/
		for(int i=0;i<numnodes;i++) pe->values[i]+=scalar_def*basis[i];

		/* Build transient now */
		if(reCast<bool,IssmDouble>(dt)){
			enthalpy_input->GetInputValue(&enthalpy, gauss);
			scalar_transient=enthalpy*Jdet*gauss->weight;
			for(int i=0;i<numnodes;i++) pe->values[i]+=scalar_transient*basis[i];
		}

		if(stabilization==2){
			element->NodalFunctionsDerivatives(dbasis,xyz_list,gauss);

			vx_input->GetInputValue(&u,gauss);
			vy_input->GetInputValue(&v,gauss);
			vz_input->GetInputValue(&w,gauss);
			kappa          = element->EnthalpyDiffusionParameterVolume(numvertices,enthalpypicard,pressure) / rho_ice;
			tau_parameter  = element->StabilizationParameter(u,v,w,diameter,kappa);

			for(int i=0;i<numnodes;i++) pe->values[i]+=tau_parameter*scalar_def*(u*dbasis[0*3+i]+v*dbasis[1*3+i]+w*dbasis[2*3+i]);
			if(reCast<bool,IssmDouble>(dt)){
				for(int i=0;i<numnodes;i++) pe->values[i]+=tau_parameter*scalar_transient*(u*dbasis[0*3+i]+v*dbasis[1*3+i]+w*dbasis[2*3+i]);
			}
		}
	}

	/*Clean up and return*/
	xDelete<IssmDouble>(basis);
	xDelete<IssmDouble>(dbasis);
	xDelete<IssmDouble>(pressure);
	xDelete<IssmDouble>(enthalpypicard);
	xDelete<IssmDouble>(xyz_list);
	delete gauss;
	return pe;

}/*}}}*/
ElementVector* EnthalpyAnalysis::CreatePVectorSheet(Element* element){/*{{{*/
	return NULL;
}/*}}}*/
ElementVector* EnthalpyAnalysis::CreatePVectorShelf(Element* element){/*{{{*/

	IssmDouble  h_pmp,dt,Jdet,scalar_ocean,pressure;
	IssmDouble *xyz_list_base = NULL;

	/*Get basal element*/
	if(!element->IsOnBed() || !element->IsFloating()) return NULL;

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

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

	/*Retrieve all inputs and parameters*/
	element->GetVerticesCoordinatesBase(&xyz_list_base);
	element->FindParam(&dt,TimesteppingTimeStepEnum);
	Input*      pressure_input=element->GetInput(PressureEnum); _assert_(pressure_input);
	IssmDouble  gravity             = element->GetMaterialParameter(ConstantsGEnum);
	IssmDouble  rho_water           = element->GetMaterialParameter(MaterialsRhoWaterEnum);
	IssmDouble  rho_ice             = element->GetMaterialParameter(MaterialsRhoIceEnum);
	IssmDouble  heatcapacity        = element->GetMaterialParameter(MaterialsHeatcapacityEnum);
	IssmDouble  mixed_layer_capacity= element->GetMaterialParameter(MaterialsMixedLayerCapacityEnum);
	IssmDouble  thermal_exchange_vel= element->GetMaterialParameter(MaterialsThermalExchangeVelocityEnum);

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

		element->JacobianDeterminantBase(&Jdet,xyz_list_base,gauss);
		element->NodalFunctions(basis,gauss);

		pressure_input->GetInputValue(&pressure,gauss);
		h_pmp=element->PureIceEnthalpy(pressure);

		scalar_ocean=gauss->weight*Jdet*rho_water*mixed_layer_capacity*thermal_exchange_vel*h_pmp/(heatcapacity*rho_ice);
		if(reCast<bool,IssmDouble>(dt)) scalar_ocean=dt*scalar_ocean;

		for(int i=0;i<numnodes;i++) pe->values[i]+=scalar_ocean*basis[i];
	}

	/*Clean up and return*/
	delete gauss;
	xDelete<IssmDouble>(basis);
	xDelete<IssmDouble>(xyz_list_base);
	return pe;
}/*}}}*/
void EnthalpyAnalysis::GetSolutionFromInputs(Vector<IssmDouble>* solution,Element* element){/*{{{*/
	element->GetSolutionFromInputsOneDof(solution,EnthalpyEnum);
}/*}}}*/
void EnthalpyAnalysis::InputUpdateFromSolution(IssmDouble* solution,Element* element){/*{{{*/

	bool        converged;
	int         i,rheology_law;
	IssmDouble  B_average,s_average,T_average=0.,P_average=0.;
	int        *doflist   = NULL;
	IssmDouble *xyz_list  = NULL;

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

	/*Fetch dof list and allocate solution vector*/
	element->GetDofList(&doflist,NoneApproximationEnum,GsetEnum);
	IssmDouble* values        = xNew<IssmDouble>(numnodes);
	IssmDouble* pressure      = xNew<IssmDouble>(numnodes);
	IssmDouble* surface       = xNew<IssmDouble>(numnodes);
	IssmDouble* B             = xNew<IssmDouble>(numnodes);
	IssmDouble* temperature   = xNew<IssmDouble>(numnodes);
	IssmDouble* waterfraction = xNew<IssmDouble>(numnodes);

	/*Use the dof list to index into the solution vector: */
	for(i=0;i<numnodes;i++){
		values[i]=solution[doflist[i]];

		/*Check solution*/
		if(xIsNan<IssmDouble>(values[i])) _error_("NaN found in solution vector");
	}

	/*Get all inputs and parameters*/
	element->GetInputValue(&converged,ConvergedEnum);
	element->GetInputListOnNodes(&pressure[0],PressureEnum);
	if(converged){
		for(i=0;i<numnodes;i++){
			element->EnthalpyToThermal(&temperature[i],&waterfraction[i],values[i],pressure[i]);
			if(waterfraction[i]<0.) _error_("Negative water fraction found in solution vector");
			if(waterfraction[i]>1.) _error_("Water fraction >1 found in solution vector");
		}
		element->AddInput(EnthalpyEnum,values,P1Enum);
		element->AddInput(WaterfractionEnum,waterfraction,P1Enum);
		element->AddInput(TemperatureEnum,temperature,P1Enum);

		/*Update Rheology only if converged (we must make sure that the temperature is below melting point
		 * otherwise the rheology could be negative*/
		element->FindParam(&rheology_law,MaterialsRheologyLawEnum);
		element->GetInputListOnNodes(&surface[0],SurfaceEnum);
		switch(rheology_law){
			case NoneEnum:
				/*Do nothing: B is not temperature dependent*/
				break;
			case PatersonEnum:
				for(i=0;i<numnodes;i++) B[i]=Paterson(temperature[i]);
				element->AddMaterialInput(MaterialsRheologyBEnum,&B[0],P1Enum);
				break;
			case ArrheniusEnum:
				element->GetVerticesCoordinates(&xyz_list);
				for(i=0;i<numnodes;i++) B[i]=Arrhenius(temperature[i],surface[i]-xyz_list[i*3+2],element->GetMaterialParameter(MaterialsRheologyNEnum));
				element->AddMaterialInput(MaterialsRheologyBEnum,&B[0],P1Enum);
				break;
			case LliboutryDuvalEnum:
				for(i=0;i<numnodes;i++) B[i]=LliboutryDuval(values[i],pressure[i],element->GetMaterialParameter(MaterialsRheologyNEnum),element->GetMaterialParameter(MaterialsBetaEnum),element->GetMaterialParameter(ConstantsReferencetemperatureEnum),element->GetMaterialParameter(MaterialsHeatcapacityEnum),element->GetMaterialParameter(MaterialsLatentheatEnum)); 
				element->AddMaterialInput(MaterialsRheologyBEnum,&B[0],P1Enum); 
				break; 
			default: _error_("Rheology law " << EnumToStringx(rheology_law) << " not supported yet");
		}
	}
	else{
		element->AddInput(EnthalpyPicardEnum,values,P1Enum);
	}

	/*Free ressources:*/
	xDelete<IssmDouble>(values);
	xDelete<IssmDouble>(pressure);
	xDelete<IssmDouble>(surface);
	xDelete<IssmDouble>(B);
	xDelete<IssmDouble>(temperature);
	xDelete<IssmDouble>(waterfraction);
	xDelete<IssmDouble>(xyz_list);
	xDelete<int>(doflist);
}/*}}}*/