/*
 _____ ___  ____   ____ 
|_   _/ _ \|  _ \ / ___|
  | || | | | |_) | |    
  | || |_| |  _ <| |___ 
  |_| \___/|_| \_\\____|
                        
*/

#include "UserData.h"
#include "solution.h"
#include "residue.h"
#include "macros.h"
#include "timing.h"

#include <ida/ida.h>
#include <ida/ida_direct.h>
#include <sunmatrix/sunmatrix_band.h>
#include <sunlinsol/sunlinsol_lapackband.h>
//#include <ida/ida_band.h>

static int check_flag(void *flagvalue, 
		      const char *funcname, 
		      int opt);

void freeAtLast(void* mem, N_Vector *y, 
		N_Vector *ydot, 
		N_Vector *res, 
		N_Vector *id, 
		N_Vector *atolv,
		N_Vector *constraints,UserData data);

int main(){

	// Read input file specifying the details of the case and store them
  	FILE *input;input=fopen("input.dat","r");
  	UserData data;data=NULL;data=allocateUserData(input);
	fclose(input);
	data->clockStart=get_wall_time();

	if(data==NULL){
		printf("check input file!\n");
  		freeUserData(data);
		return(-1);
	}

	// Allocate solution variables
	long int ier,mu,ml,count,netf,ncfn,njevals,nrevals;
	realtype tNow,*atolvdata,*constraintsdata,finalTime,dtMax,tolsfac;


  	N_Vector y,ydot,id,res,atolv,constraints;
  	y=ydot=id=res=atolv=constraints=NULL;
 	ier=allocateSolution(data->neq,data->nThreads,&y,&ydot,&id,&res,&atolv,&constraints);
	ier=setAlgebraicVariables(&id,data);
	ier=setInitialCondition(&y,&ydot,data);
	if(ier==-1)return(-1);
	tNow=data->tNow;
	finalTime=data->finalTime;
	//TODO: dtMax should be a user input
	dtMax = 1e-4;

	double* ydata;
	double* ydotdata;
  	ydata    = N_VGetArrayPointer_OpenMP(y);
  	ydotdata = N_VGetArrayPointer_OpenMP(ydot);
	////////// DEBUG ///////////////////
	//double* resdata;
	//double* iddata;
	//resdata = N_VGetArrayPointer_OpenMP(res);
	//iddata = N_VGetArrayPointer_OpenMP(id);
	///////////////////////////////////

	void *mem;mem=NULL;mem = IDACreate();
  	ier = IDASetUserData(mem, data);
  	ier = IDASetId(mem, id);
  	ier = IDAInit(mem, residue, tNow, y, ydot);

        // Atol array
        atolvdata = N_VGetArrayPointer_OpenMP(atolv);
	for (size_t i = 1; i <=data->npts; i++) {
             atolT(i) = data->temperatureTolerance;
             for (size_t k = 1; k <=data->nsp; k++) {
		     	if(k!=data->k_bath){
                		atolY(i,k) = data->massFractionTolerance;
			}
			else{
                		atolY(i,k) = data->bathGasTolerance;
			}

             }
             atolR(i) = data->radiusTolerance;
             atolP(i) = data->pressureTolerance;
			 atolMdot(i) = data->MdotTolerance;
	}
  	ier = IDASVtolerances(mem, data->relativeTolerance, atolv);

  	mu = 2*data->nvar; ml = mu;
	SUNMatrix A; A=NULL;
	A=SUNBandMatrix(data->neq,mu,ml,mu+ml);
	SUNLinearSolver LS; LS=NULL;
	//LS=SUNBandLinearSolver(y,A);
	LS=SUNLapackBand(y,A);
	ier=IDADlsSetLinearSolver(mem,LS,A);
  	//ier = IDABand(mem, data->neq, mu, ml);

        constraintsdata = N_VGetArrayPointer_OpenMP(constraints);
	if(data->setConstraints){
		for (size_t i = 1; i <=data->npts; i++) {
             		for (size_t k = 1; k <=data->nsp; k++) {
             		   constraintsY(i,k) = ONE;
             		}
		}
  		ier=IDASetConstraints(mem, constraints);
	}
	
	if(!data->quasiSteady){
		constraintsR(1) = ONE;
		ier=IDASetConstraints(mem, constraints);
	}

	ier = IDASetSuppressAlg(mem, data->suppressAlg);
  	if(check_flag(&ier, "IDASetSuppressAlg", 1)) return(1);

	//ier= IDASetMaxNumStepsIC(mem, 1);
	//ier= IDASetMaxNumJacsIC(mem,8);
	//ier= IDASetMaxNumItersIC(mem,100);
	//ier= IDASetMaxBacksIC(mem,2000);
	//ier = IDASetLineSearchOffIC(mem,SUNTRUE);

	//////// DEBUG ///////////
	//if(data->dryRun){
    //	ier = residue(tNow,y, ydot, res, data);
	//	for(size_t k=0; k < data->nvar*data->npts; k++){
	//		printf("%i: %15.6e\n",k,resdata[k]);
	//	}
	//	for(size_t k=0; k < data->nvar*data->npts; k++){
	//		if(iddata[k] == 1){
	//			ydotdata[k] = -resdata[k];
	//		}
	//	}
    //	ier = residue(tNow,y, ydot, res, data);
	//	for(size_t k=0; k < data->nvar*data->npts; k++){
	//		printf("%i: %15.6e\n",k,resdata[k]);
	//	}
	//}
	//for(size_t k=0; k < data->neq; k++){
	//	if(iddata[k] == 1){
	//		ydotdata[k] = -resdata[k];
	//	}
	//}
	//////////////////////////
	
	if(!data->dryRun){
		printf("Calculating Initial Conditions:\n");
		printf("Cross your fingers...\n");
  		ier = IDACalcIC(mem, IDA_YA_YDP_INIT, 1e-05*finalTime);
		//If at first IDACalcIC doesn't succeed, try, try, try again:
		for (int i = 0; i < 10; i++) {
  			ier = IDACalcIC(mem, IDA_YA_YDP_INIT, (1e-01+pow(10,i)*finalTime));
			if(ier==0){
				break;
			}
		}

		//...or else cry again :(
  		if(check_flag(&ier, "IDACalcIC", 1)){
			freeAtLast(mem,&y,&ydot,&res,&id,&atolv,&constraints,data);
		       	return(-1);
		}else{
			printf("Initial (Consistent) Conditions Calculated!\n");
		}
		ier = IDASetInitStep(mem,1e-12);
	}

	printSpaceTimeHeader(data);
	printGlobalHeader(data);
	printSpaceTimeOutput(tNow, &y, data->output, data);
	printSpaceTimeOutput(tNow, &y, data->gridOutput, data);

	if(!data->dryRun){
		count=0;
		double dt=1e-08;
		double t1=0.0e0;
		double xOld=0.0e0;
		double x=0.0e0;
		double dx=0.0e0;
		double dxMin=1.0e0;
		double dxRatio=dx/dxMin;
		int move=0;
		int kcur=0;
		if(data->adaptiveGrid){
			dxMin=data->grid->leastMove;
			//xOld=maxCurvPosition(ydata, data->nt, data->nvar,
			//		     data->grid->x, data->npts);
			xOld=isothermPosition(ydata, data->isotherm, data->nt, 
	                		   data->nvar, data->grid->x, data->npts);
		}
		while (tNow<=finalTime && R(1)>100e-9) {

			t1=tNow;

			if(data->quasiSteady){
  				ier = IDASolve(mem, finalTime, &tNow, y, ydot, IDA_ONE_STEP);
			}else{
				/*This prevents the solver from taking a step so large that
				 *the droplet radius becomes negative.
				 *TODO:Try adding the constraint that the droplet radius must
				 *     be a positive number*/
  				ier = IDASolve(mem, tNow+dtMax, &tNow, y, ydot, IDA_ONE_STEP);
  				//ier = IDASolve(mem, tNow+dtMax, &tNow, y, ydot, IDA_NORMAL);
			}

  			if(check_flag(&ier, "IDASolve", 1)){
				freeAtLast(mem,&y,&ydot,&res,&id,&atolv,&constraints,data);
		       		return(-1);
			}
			dt=tNow-t1;
  			ier = IDAGetCurrentOrder(mem, &kcur);

			if(data->adaptiveGrid==1 && data->moveGrid==1){
				x=maxCurvPosition(ydata, data->nt, data->nvar,
						  data->grid->x, data->npts);

				//x=isothermPosition(ydata, data->isotherm, data->nt, 
	                	//		   data->nvar, data->grid->x, data->npts);

				//x=maxGradPosition(ydata, data->nt, data->nvar,
				//		  data->grid->x, data->npts);
				dx=x-xOld;

				if(dx*dxMin>0.0e0){
					move=1;
				}else{
					move=-1;
				}

				//if(fabs(dx)>=dxMin && x+(double)(move)*0.5e0*dxMin<=1.0e0){
				dxRatio=fabs(dx/dxMin);
				if(dxRatio>=1.0e0 && dxRatio<=2.0e0){
					printf("Regridding!\n");

					data->regrid=1;
					printSpaceTimeOutput(tNow, &y, data->gridOutput, data);

					ier=reGrid(data->grid, x+(double)(move)*0.5e0*dxMin);
					if(ier==-1){
						freeAtLast(mem,&y,&ydot,&res,&id,&atolv,&constraints,data);
						return(-1);
					}

					updateSolution(ydata, ydotdata, data->nvar,
					               data->grid->xOld,data->grid->x,data->npts);
					storeGrid(data->grid->x,data->grid->xOld,data->npts);
					xOld=x;

					printf("Regrid Complete! Restarting Problem at %15.6e s\n",tNow);
					ier = IDAReInit(mem, tNow, y, ydot);
  					if(check_flag(&ier, "IDAReInit", 1)){
						freeAtLast(mem,&y,&ydot,&res,&id,&atolv,&constraints,data);
		       				return(-1);
					}
					ier = IDASetInitStep(mem,1e-01*dt);
					printf("Reinitialized! Calculating Initial Conditions:\n");
					printf("Cross your fingers...\n");
  					ier = IDACalcIC(mem, IDA_YA_YDP_INIT, tNow+1e-01*dt);
  					if(check_flag(&ier, "IDACalcIC", 1)){
  						ier = IDACalcIC(mem, IDA_YA_YDP_INIT, tNow+1e+01*dt);
					}
					//Every once in a while, for reasons
					//that befuddle this mathematically
					//lesser author, IDACalcIC fails. Here,
					//I desperately try to make it converge
					//again by sampling increasingly larger
					//time-steps:
					for (int i = 0; i < 10; i++) {
  						ier = IDACalcIC(mem, IDA_YA_YDP_INIT, tNow+(1e-01+pow(10,i)*dt));
						if(ier==0){
							break;
						}
					}
					//Failure :( Back to the drawing board:
  					if(check_flag(&ier, "IDACalcIC", 1)){
						freeAtLast(mem,&y,&ydot,&res,&id,&atolv,&constraints,data);
		       				return(-1);
					}
					printf("Initial (Consistent) Conditions Calculated!\n");
					printf("------------------------------------------\n\n");
					if(data->writeEveryRegrid){
						printSpaceTimeOutput(tNow, &y, data->output, data);
						FILE* fp;
						fp=fopen("restart.bin","w");
						writeRestart(tNow, &y, &ydot, fp, data);
						fclose(fp);
					}

				}
			}

			if(count%data->nSaves==0 && !data->writeEveryRegrid){
				printSpaceTimeOutput(tNow, &y, data->output, data);
				//if(data->writeRates){
				//	printSpaceTimeRates(tNow, ydot, data);
				//}
			}

			/*Print global variables only if time-step is of high order!*/
			if(data->nTimeSteps==0){
				data->flamePosition[0]=0.0e0;
				data->flamePosition[1]=0.0e0;
				data->flameTime[0]=tNow;
				data->flameTime[1]=tNow;
			}

  			ier = IDAGetNumErrTestFails(mem, &netf);
  			ier = IDAGetNumNonlinSolvConvFails(mem, &ncfn);
  			ier = IDADlsGetNumJacEvals(mem, &njevals);
  			ier = IDADlsGetNumResEvals(mem, &nrevals);
  			printf("etf = %ld ,"
			       "nlsf= %ld ," 
			       "J=%ld ,"
			       "R=%ld ,"
			       "o=%d ,",netf, ncfn, njevals, nrevals, kcur);
			printf("Time=%15.6e s,",tNow);
			printf("frac: %15.6e\n",dxRatio);
			count++;
			data->nTimeSteps=count;
		}
	}

	SUNLinSolFree(LS);
	SUNMatDestroy(A);
	freeAtLast(mem,&y,&ydot,&res,&id,&atolv,&constraints,data);

	return(0);
}

void freeAtLast(void* mem, 
		N_Vector *y, 
		N_Vector *ydot, 
		N_Vector *res, 
		N_Vector *id, 
		N_Vector *atolv,
		N_Vector *constraints,UserData data){

  	IDAFree(&mem);
	freeSolution(y,ydot,res,id,atolv,constraints);
  	freeUserData(data);
}

static int check_flag(void *flagvalue, const char *funcname, int opt)
{
  int *errflag;

  /* Check if SUNDIALS function returned NULL pointer - no memory allocated */
  if (opt == 0 && flagvalue == NULL) {
    fprintf(stderr, 
            "\nSUNDIALS_ERROR: %s() failed - returned NULL pointer\n\n",
	    funcname);
    return(1);
  }

  /* Check if flag < 0 */
  else if (opt == 1) {
    errflag = (int *) flagvalue;
    if (*errflag < 0) {
      fprintf(stderr,
              "\nSUNDIALS_ERROR: %s() failed with flag = %d\n\n",
	      funcname, *errflag);
      return(1); 
    }
  }

  /* Check if function returned NULL pointer - no memory allocated */
  else if (opt == 2 && flagvalue == NULL) {
    fprintf(stderr,
            "\nMEMORY_ERROR: %s() failed - returned NULL pointer\n\n",
	    funcname);
    return(1);
  }

  return(0);
}