/*!\file: control_core.cpp
 * \brief: core of the control solution 
 */ 
#include <config.h>
#include "../toolkits/toolkits.h"
#include "../objects/objects.h"
#include "../shared/shared.h"
#include "../io/io.h"
#include "../EnumDefinitions/EnumDefinitions.h"
#include "./solutions.h"
#include "../modules/modules.h"
#include "../include/include.h"
#include "../solvers/solvers.h"

#if defined (_HAVE_TAO_) && (_PETSC_MAJOR_ == 3 && _PETSC_MINOR_ == 2)
#include <tao.h>

/*Local prototype*/
int FormFunctionGradient(TaoSolver,Vec,double*,Vec,void*);
int IssmMonitor(TaoSolver,void*);
typedef struct {
	FemModel* femmodel;
} AppCtx;

void controltao_core(FemModel* femmodel){

	/*TAO*/
	int        ierr;
	int        num_controls,solution_type;
	int        nsteps,maxiter;
	AppCtx     user;
	TaoSolver  tao;
	double    *dummy          = NULL;
	int       *control_list   = NULL;
	Vec        X              = NULL;
	Vec        XL             = NULL;
	Vec        XU             = NULL;

	/*Initialize TAO*/
	int argc; char **args=NULL;
	PetscGetArgs(&argc,&args);
	ierr = TaoInitialize(&argc,&args,(char*)0,"");
	if(ierr) _error_("Could not initialize Tao");

	/*Recover some parameters*/
	femmodel->parameters->FindParam(&solution_type,SolutionTypeEnum);
	femmodel->parameters->FindParam(&num_controls,InversionNumControlParametersEnum);
	femmodel->parameters->FindParam(&control_list,NULL,InversionControlParametersEnum);
	femmodel->parameters->FindParam(&nsteps,InversionNstepsEnum);
	femmodel->parameters->FindParam(&dummy,NULL,NULL,InversionMaxiterPerStepEnum);
	maxiter=nsteps*(int)dummy[0]; xfree((void**)&dummy);

	/*Initialize TAO*/
	_printf_(VerboseControl(),"%s\n","   Initializing the Toolkit for Advanced Optimization (TAO)");
	TaoCreate(PETSC_COMM_WORLD,&tao);
	TaoSetFromOptions(tao);
	TaoSetType(tao,"tao_blmvm");
	//TaoSetType(tao,"tao_cg");
	//TaoSetType(tao,"tao_lmvm");

	/*Prepare all TAO parameters*/
	TaoSetMonitor(tao,IssmMonitor,&user,NULL);
	TaoSetMaximumFunctionEvaluations(tao,maxiter);
	TaoSetMaximumIterations(tao,nsteps);
	TaoSetTolerances(tao,0.,0.,0.,0.,0.);

	GetVectorFromControlInputsx(&X, femmodel->elements,femmodel->nodes,femmodel->vertices,femmodel->loads,femmodel->materials,femmodel->parameters,"value");
	GetVectorFromControlInputsx(&XL,femmodel->elements,femmodel->nodes,femmodel->vertices,femmodel->loads,femmodel->materials,femmodel->parameters,"lowerbound");
	GetVectorFromControlInputsx(&XU,femmodel->elements,femmodel->nodes,femmodel->vertices,femmodel->loads,femmodel->materials,femmodel->parameters,"upperbound");
	TaoSetInitialVector(tao,X);
	TaoSetVariableBounds(tao,XL,XU);
	VecFree(&XL);
	VecFree(&XU);

	user.femmodel=femmodel;
	TaoSetObjectiveAndGradientRoutine(tao,FormFunctionGradient,(void*)&user); 

	/*Solver optimization problem*/
	_printf_(VerboseControl(),"%s\n","   Starting optimization");
	TaoSolve(tao);
	TaoView(tao,PETSC_VIEWER_STDOUT_WORLD);
	TaoGetSolutionVector(tao,&X);
	SetControlInputsFromVectorx(femmodel->elements,femmodel->nodes,femmodel->vertices,femmodel->loads,femmodel->materials,femmodel->parameters,X);
	for(int i=0;i<num_controls;i++){
		InputToResultx(femmodel->elements,femmodel->nodes,femmodel->vertices,femmodel->loads,femmodel->materials,femmodel->parameters,control_list[i]);
	}

	/*Finalize*/
	_printf_(VerboseControl(),"%s\n","   preparing final solution");
	femmodel->parameters->SetParam(false,InversionIscontrolEnum); //needed to turn control result output in solutioncore
	void (*solutioncore)(FemModel*)=NULL;
	CorePointerFromSolutionEnum(&solutioncore,femmodel->parameters,solution_type);
	solutioncore(femmodel);

	/*Clean up and return*/
	xfree((void**)&control_list);
	VecFree(&X);
	TaoDestroy(&tao);
	TaoFinalize();
}
int FormFunctionGradient(TaoSolver tao, Vec X, double *fcn,Vec G,void *userCtx){

	/*Retreive arguments*/
	int       solution_type,num_cost_functions;
	AppCtx   *user           = (AppCtx *)userCtx;
	FemModel *femmodel       = user->femmodel;
	int      *cost_functions = NULL;
	double   *cost_functionsd= NULL;
	Vec       gradient       = NULL;

	/*Set new variable*/
	//VecView(X,PETSC_VIEWER_STDOUT_WORLD);
	SetControlInputsFromVectorx(femmodel->elements,femmodel->nodes,femmodel->vertices,femmodel->loads,femmodel->materials,femmodel->parameters,X);

	/*Recover some parameters*/
	femmodel->parameters->FindParam(&solution_type,SolutionTypeEnum);
	femmodel->parameters->FindParam(&num_cost_functions,InversionNumCostFunctionsEnum);
	femmodel->parameters->FindParam(&cost_functionsd,NULL,NULL,InversionCostFunctionsEnum);

	/*Prepare objective function*/
	cost_functions=(int*)xmalloc(num_cost_functions*sizeof(int));
	for(int i=0;i<num_cost_functions;i++) cost_functions[i]=(int)cost_functionsd[i]; //FIXME
	femmodel->parameters->SetParam(cost_functions,1,num_cost_functions,StepResponsesEnum);

	/*Compute solution and adjoint*/
	void (*solutioncore)(FemModel*)=NULL;
	void (*adjointcore)(FemModel*)=NULL;
	CorePointerFromSolutionEnum(&solutioncore,femmodel->parameters,solution_type);
	AdjointCorePointerFromSolutionEnum(&adjointcore,solution_type);
	solutioncore(femmodel);
	adjointcore(femmodel);

	/*Compute objective function*/
	CostFunctionx(fcn,femmodel->elements,femmodel->nodes,femmodel->vertices,femmodel->loads,femmodel->materials,femmodel->parameters);

	/*Compute gradient*/
	Gradjx(&gradient,NULL,femmodel->elements,femmodel->nodes,femmodel->vertices,femmodel->loads,femmodel->materials,femmodel->parameters);
	VecCopy(gradient,G); VecFree(&gradient);
	VecScale(G,-1.);

	/*Clean-up and return*/
	xfree((void**)&cost_functions);
	xfree((void**)&cost_functionsd);
	return 0;
}
int IssmMonitor(TaoSolver tao, void *userCtx){

	int       i,its,num_responses;
	double    f,gnorm,cnorm,xdiff;
	AppCtx   *user      = (AppCtx *)userCtx;
	FemModel *femmodel  = user->femmodel;
	Element  *element   = NULL;
	int      *responses = NULL;

	femmodel->parameters->FindParam(&num_responses,InversionNumCostFunctionsEnum);
	femmodel->parameters->FindParam(&responses,NULL,NULL,StepResponsesEnum);

	TaoGetSolutionStatus(tao, &its, &f, &gnorm, &cnorm, &xdiff, NULL);
	if(its==0) _printf_(true,"Iter       Function      Residual  |  List of contributions\n");
	if(its==0) _printf_(true,"-----------------------------------+-----------------------\n");
	_printf_(true,"%4i   %12.7g  %12.7g  | ",its,f,gnorm);

	/*Retrieve objective functions independently*/
	for(i=0;i<num_responses;i++){
		Responsex(&f,femmodel->elements,femmodel->nodes,femmodel->vertices,femmodel->loads,femmodel->materials,femmodel->parameters,EnumToStringx(responses[i]),false,i);
		_printf_(true," %12.7g ","",f);
	}
	_printf_(true,"\n");

	/*Clean-up and return*/
	xfree((void**)&responses);
	return 0;
}

#else
void controltao_core(FemModel* femmodel){
	_error_("TAO not installed or PETSc version not supported");
}
#endif //_HAVE_TAO_ 
