/*
  
  Steady-state non-negative diffusion
  
  Written by: Justin Chang

*/

static char help[] = "Solves a diffusion problem with concentration condition\n\
3D cube with hole - requires input mesh and boundary file.";

#include <petscdmplex.h>
#include <petscsnes.h>
#include <petscds.h>
#include <petscmath.h>
#include <petsctime.h>
#include <petsctao.h>
#include <petscviewerhdf5.h>

PetscInt    dim = 3;	             /* Dimension of problem */
PetscScalar D[9];						       /* Diffusivity tensor */
PetscScalar bc_values[2];          /* Concentration BCs */
PetscScalar alpha_t = 0.1;         /* Tensorial dispersion coefficients */
PetscScalar alpha_l = 10;          /* Tensorial dispersion coefficients */

/* Timings */
PetscLogDouble      v_create,v_create_distribute,v_create_total,v_post,v_solve,v_total;

typedef struct {
  PetscLogEvent createMeshEvent;
	char          meshData[2048];		  /* Input mesh file */
	char          boundaryData[2048]; /* Input boundary file */
	char          output[2048];		 	  /* Output VTU file name */
	PetscBool			flg_out;					  /* Flag for outputting concentration */
	PetscBool			flg_input;				  /* Flag for input parameters */
	PetscBool			nonneg;				  /* Flag for input parameters */
	PetscViewer		viewer;						  /* Multi-purpose viewer */
	PetscSF				sf;								  /* Mapping for distributed mesh */
	Mat A,J;												  /* Jacobian/Hessian matrices */
	Vec u,r,b;											 /* Solution vectors */
	Vec f,lb,ub;											 /* Gradient and lower bounds vectors */
	
	/* Input datafile parameters */
	PetscInt			*numPoints,*coneSize,*connectivity,*orientations,*bc_nodes;
	PetscScalar		*coord;
	PetscInt			Nele,Nnodes,npe,Nbc,Naux;
	
} AppCtx;

/* Initial guess */
void initial(const PetscReal x[], PetscScalar *u, void *ctx)
{
  *u = bc_values[1];
}

void maxSol(const PetscReal x[], PetscScalar *u, void *ctx)
{
  *u = bc_values[0];
}

/* Calculate diffusivity tensor */
void calculate_D()
{
	const PetscScalar u_x[3] = {1,1,1};
	PetscScalar coeff1,coeff2,vnorm;
	vnorm = sqrt(pow(u_x[0],2)+pow(u_x[1],2)+pow(u_x[2],2));
	coeff1 = alpha_t*vnorm;
	if (vnorm == 0)
		coeff2 = 0;
	else
		coeff2 = (alpha_l-alpha_t)/vnorm;
	
	/* tensorial dispersion */
	D[0] = coeff1+coeff2*u_x[0]*u_x[0];
	D[1] = coeff2*u_x[0]*u_x[1];
	D[2] = coeff2*u_x[0]*u_x[2];
	D[3] = coeff2*u_x[1]*u_x[0];
	D[4] = coeff1+coeff2*u_x[1]*u_x[1];
	D[5] = coeff2*u_x[1]*u_x[2];
	D[6] = coeff2*u_x[2]*u_x[0];
	D[7] = coeff2*u_x[2]*u_x[1];
	D[8] = coeff1+coeff2*u_x[2]*u_x[2];
}

/* Boundary conditions */
void bcs(const PetscReal x[], PetscScalar *u, void *ctx)
{
  if (x[0] >= 0.45 && x[0] <= 0.55 && x[1] >= 0.45 && x[1] <= 0.55 && x[2] >= 0.45 && x[2] <= 0.55) {
    *u = bc_values[0];
  }else {
    *u = bc_values[1];
    }
}

/* Residual -<w,f(x)>  */
void f0(const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], const PetscReal x[], PetscScalar f0[])
{
	f0[0] = 0.0;  
}

/* Residual <grad[w],D(x)grad[u]> */
void f1(const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], const PetscReal x[], PetscScalar f1[])
{
	PetscInt d,d2;
	calculate_D();
	for (d = 0; d < dim; d++) {
		f1[d] = 0.0;
		for (d2 = 0; d2 < dim; d2++)
			f1[d] += D[d*dim+d2]*u_x[d2];
	}
}

/* Jacobian <grad[w],D(x)grad[u]> */
void g3(const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], const PetscReal x[], PetscScalar g3[])
{
	PetscInt d,d2;
	calculate_D();
	for (d = 0; d < dim; d++) {
		for (d2 = 0; d2 < dim; d2++)
			g3[d*dim+d2] = D[d*dim+d2];
	}
}

#undef __FUNCT__
#define __FUNCT__ "ProcessOptions"
PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
{
  PetscErrorCode ierr;
  PetscFunctionBeginUser;
	options->meshData[0]        = '\0';
	options->boundaryData[0]    = '\0';
	options->output[0]          = '\0';
	options->flg_out					= PETSC_FALSE;
	options->nonneg					= PETSC_FALSE;
	options->flg_input					= PETSC_FALSE;
	PetscInt numCond;
  ierr = PetscOptionsBegin(comm, "", "Flow problem Options", "DMPLEX");CHKERRQ(ierr);
  
  /* Nonnegative optimization */
  ierr = PetscOptionsBool("-nonnegative", "Solve with the non-negative constraint", "trans.c", options->nonneg, &options->nonneg, NULL);CHKERRQ(ierr);
  
  /* Input mesh file */
  ierr = PetscOptionsString("-f_mesh", "Required mesh file", "trans.c", options->meshData, options->meshData, sizeof(options->output), &options->flg_input);CHKERRQ(ierr);
  if (!options->flg_input) {
  	SETERRQ(comm, PETSC_ERR_ARG_WRONG, "Input mesh file is required: -f_mesh <name>");
  }
	
	/* Input boundary file */
	ierr = PetscOptionsString("-f_boundary", "Required boundary file", "trans.c", options->boundaryData, options->boundaryData, sizeof(options->output), &options->flg_input);CHKERRQ(ierr);
	if (!options->flg_input) {
		SETERRQ(comm, PETSC_ERR_ARG_WRONG, "Input boundary file is required: -f_boundary <name>");
	}
	
  /* Output result - optional */
  ierr = PetscOptionsBool("-f_out", "Output result", "flow.c", options->flg_out, &options->flg_out, NULL);CHKERRQ(ierr);
  
	/* Output VTU file names */
  ierr = PetscOptionsString("-f_outname", "Optional output VTU file name", "trans.c", options->output, options->output, sizeof(options->output), &options->flg_input);CHKERRQ(ierr);
  if (options->flg_out && !options->flg_input) {
  	SETERRQ(comm, PETSC_ERR_ARG_WRONG, "Output VTU file names required: -f_outname <name>");
  }
	
	/* Read concentration */
	numCond = 2;
	ierr = PetscOptionsGetScalarArray(NULL,"-concentration", bc_values, &numCond, &options->flg_input);CHKERRQ(ierr);
	if (!options->flg_input) {
		SETERRQ(comm, PETSC_ERR_ARG_WRONG, "Input concentration boundary condition value is required: -concentration <input>");
	}
  
  /* Read alpha_l */
  ierr = PetscOptionsGetScalar(NULL,"-f_al", &alpha_l, &options->flg_input);CHKERRQ(ierr);
  if (!options->flg_input) {
  	SETERRQ(comm, PETSC_ERR_ARG_WRONG, "Input longitudinal alpha is required: -f_al <input>");
  }
  
  /* Read alpha_t */
  ierr = PetscOptionsGetScalar(NULL,"-f_at", &alpha_t, &options->flg_input);CHKERRQ(ierr);
  if (!options->flg_input) {
  	SETERRQ(comm, PETSC_ERR_ARG_WRONG, "Input transverse alpha is required: -f_at <input>");
  }

  ierr = PetscOptionsEnd();
  ierr = PetscLogEventRegister("CreateMesh", DM_CLASSID, &options->createMeshEvent);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "CreateMesh"
PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm)
{
  PetscErrorCode ierr;
  PetscMPIInt    rank,size;
  PetscInt			 d=1,i,id = 1,fd;
	const char     *meshData = user->meshData;
	const char     *boundaryData = user->boundaryData;

  PetscFunctionBeginUser;
  ierr = MPI_Comm_rank(comm, &rank);CHKERRQ(ierr);
  ierr = MPI_Comm_size(comm, &size);CHKERRQ(ierr);
  ierr = PetscLogEventBegin(user->createMeshEvent,0,0,0,0);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,"Creating DMPlex... ");CHKERRQ(ierr);
	
	/* Read input files */
	if (!rank) {
		
		/* Open meshData */
		ierr = PetscViewerBinaryOpen(PETSC_COMM_SELF,meshData,FILE_MODE_READ,&user->viewer);CHKERRQ(ierr);
		ierr = PetscViewerBinaryGetDescriptor(user->viewer,&fd);CHKERRQ(ierr);
		/* Read number of elements */
		ierr = PetscBinaryRead(fd,&user->Nele,1,PETSC_INT);CHKERRQ(ierr);
		/* Read number of vertices */
		ierr = PetscBinaryRead(fd,&user->Nnodes,1,PETSC_INT);CHKERRQ(ierr);
		/* Read number of nodes per element (assume same for all) */
		ierr = PetscBinaryRead(fd,&user->npe,1,PETSC_INT);CHKERRQ(ierr);
		/* Read spatial dimension */
		ierr = PetscBinaryRead(fd,&dim,1,PETSC_INT);CHKERRQ(ierr);
		ierr = PetscCalloc5(2,&user->numPoints,user->Nele+user->Nnodes,&user->coneSize,user->Nele*user->npe,&user->connectivity,user->Nele*user->npe,&user->orientations,user->Nnodes*dim,&user->coord);CHKERRQ(ierr);
		/* Read element connectivity */
		ierr = PetscBinaryRead(fd,user->connectivity,user->Nele*user->npe,PETSC_INT);CHKERRQ(ierr);
		/* Read vertex coordinates */
		ierr = PetscBinaryRead(fd,user->coord,user->Nnodes*dim,PETSC_SCALAR);CHKERRQ(ierr);
		ierr = PetscBinaryClose(fd);CHKERRQ(ierr);
		ierr = PetscViewerDestroy(&user->viewer);CHKERRQ(ierr);
		
		/* Open boundaryData */
		ierr = PetscViewerBinaryOpen(PETSC_COMM_SELF,boundaryData,FILE_MODE_READ,&user->viewer);CHKERRQ(ierr);
		ierr = PetscViewerBinaryGetDescriptor(user->viewer,&fd);CHKERRQ(ierr);
		/* Read number of boundary nodes */
		ierr = PetscBinaryRead(fd,&user->Nbc,1,PETSC_INT);CHKERRQ(ierr);
		ierr = PetscCalloc1(user->Nbc,&user->bc_nodes);CHKERRQ(ierr);
		/* Read list of boundary nodes */
		ierr = PetscBinaryRead(fd,user->bc_nodes,user->Nbc,PETSC_INT);CHKERRQ(ierr);
		ierr = PetscBinaryClose(fd);CHKERRQ(ierr);
		ierr = PetscViewerDestroy(&user->viewer);CHKERRQ(ierr);
		
		/* Data for creating DMPlex from DAG */
		user->numPoints[0] = user->Nnodes;
		user->numPoints[1] = user->Nele;
		for (i=0;i<user->Nele;i++)
			user->coneSize[i] = user->npe;
	}
	
	/* Create DMPlex */
	MPI_Bcast(&dim,1,MPI_INT,0,comm);
	ierr = DMCreate(comm,dm);CHKERRQ(ierr);
	ierr = DMSetType(*dm, DMPLEX);CHKERRQ(ierr);
	ierr = DMSetDimension(*dm, dim);CHKERRQ(ierr);
	if (!rank) {
		ierr = DMPlexCreateFromDAG(*dm,d,user->numPoints,user->coneSize,user->connectivity,user->orientations,user->coord);CHKERRQ(ierr);
		for (i=0;i<user->Nbc;i++) {
			ierr = DMPlexSetLabelValue(*dm,"marker",user->bc_nodes[i],id);CHKERRQ(ierr);
		}
		ierr = PetscFree6(user->numPoints,user->coneSize,user->connectivity,user->orientations,user->coord,user->bc_nodes);CHKERRQ(ierr);
		
	} else {
		PetscInt d2[2] = {0, 0};
		ierr = DMPlexCreateFromDAG(*dm, d, d2, NULL, NULL, NULL, NULL);CHKERRQ(ierr);
	}
	ierr = PetscPrintf(PETSC_COMM_WORLD,"done\n");CHKERRQ(ierr);
	
	/* Interpolate */
	ierr = PetscPrintf(PETSC_COMM_WORLD,"\nDM before interpolating:\n");CHKERRQ(ierr);
	ierr = DMView(*dm,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
	DM idm = NULL;
	DMLabel label;
	ierr = DMPlexInterpolate(*dm, &idm);CHKERRQ(ierr);
  ierr = DMPlexCopyCoordinates(*dm, idm);CHKERRQ(ierr);
  ierr = DMPlexCopyLabels(*dm, idm);CHKERRQ(ierr);
  ierr = DMPlexGetLabel(idm, "marker", &label);CHKERRQ(ierr);
  ierr = DMPlexMarkBoundaryFaces(idm,label);CHKERRQ(ierr);
  ierr = DMPlexLabelComplete(idm,label);CHKERRQ(ierr);
  ierr = DMDestroy(dm);CHKERRQ(ierr);
  *dm = idm;
	
	ierr = PetscPrintf(PETSC_COMM_WORLD,"\nDM after interpolation:\n");CHKERRQ(ierr);
	ierr = DMView(*dm,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
	
	/* Distribute mesh if necessary */
	ierr = PetscTime(&v_create_distribute);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD,"\nDistributing mesh... ");CHKERRQ(ierr);
	if (size > 1) {
		DM distributedMesh = NULL;
		ierr = DMPlexDistribute(*dm,0,&user->sf,&distributedMesh);CHKERRQ(ierr);
		if (distributedMesh) {
			ierr = DMDestroy(dm);CHKERRQ(ierr);
			*dm  = distributedMesh;
		}
	}
	ierr = PetscPrintf(PETSC_COMM_WORLD,"done\n");CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD,"\nRefining mesh... ");CHKERRQ(ierr);
	ierr = DMPlexSetRefinementUniform(*dm, PETSC_TRUE);CHKERRQ(ierr);
  ierr = DMSetFromOptions(*dm);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD,"done\n");CHKERRQ(ierr);
  
  /* Uninterpolate
  ierr = DMView(*dm,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
	ierr = DMPlexUninterpolate(*dm, &idm);CHKERRQ(ierr);
  ierr = DMPlexCopyCoordinates(*dm, idm);CHKERRQ(ierr);
  ierr = DMPlexCopyLabels(*dm, idm);CHKERRQ(ierr);
  ierr = DMPlexGetLabel(idm, "marker", &label);CHKERRQ(ierr);
  ierr = DMPlexMarkBoundaryFaces(idm,label);CHKERRQ(ierr);
  ierr = DMDestroy(dm);CHKERRQ(ierr);
  *dm = idm;
	PetscInt markerPts,vend;
	IS markerIS;
	const PetscInt *localIS;
	ierr = DMPlexGetDepthStratum(*dm,0,NULL,&vend);CHKERRQ(ierr);
	ierr = DMPlexGetStratumIS(*dm,"marker",id,&markerIS);CHKERRQ(ierr);
	ierr = ISGetLocalSize(markerIS,&markerPts);CHKERRQ(ierr);
	ierr = ISGetIndices(markerIS,&localIS);CHKERRQ(ierr);
	for (i=0;i<markerPts;i++) {
		if (localIS[i] >= vend) {
			ierr = DMPlexClearLabelValue(*dm,"marker",localIS[i],id);CHKERRQ(ierr);
		}
	}
	ISRestoreIndices(markerIS,&localIS);CHKERRQ(ierr);
	ISDestroy(&markerIS);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD,"\nDM after uninterpolation:\n");CHKERRQ(ierr);
	ierr = DMView(*dm,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
  */
  
  ierr = DMViewFromOptions(*dm, NULL, "-dm_view");CHKERRQ(ierr);
  ierr = PetscLogEventEnd(user->createMeshEvent,0,0,0,0);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "SetupProblem"
PetscErrorCode SetupProblem(DM dm, AppCtx *user)
{
  DM             cdm   = dm;
  const PetscInt id    = 1;
  PetscFE        fe;
	PetscDS        prob;
	PetscMPIInt	 	 rank,size;
	PetscErrorCode ierr;
	
	PetscFunctionBeginUser;
	/* Create finite element for diffusion */
	ierr = MPI_Comm_rank(PETSC_COMM_WORLD, &rank);CHKERRQ(ierr);
	ierr = MPI_Comm_size(PETSC_COMM_WORLD, &size);CHKERRQ(ierr);
  ierr = PetscFECreateDefault(dm, dim, 1, PETSC_TRUE, NULL, -1, &fe);CHKERRQ(ierr);
	ierr = PetscObjectSetName((PetscObject) fe, "Concentration");CHKERRQ(ierr);
	
  /* Set discretization and boundary conditions for each mesh */
  while (cdm) {
    ierr = DMGetDS(cdm, &prob);CHKERRQ(ierr);
    ierr = PetscDSSetDiscretization(prob, 0, (PetscObject) fe);CHKERRQ(ierr);
		
		/* Set discretization for diffusion */
		ierr = PetscDSSetResidual(prob, 0, f0, f1);CHKERRQ(ierr);
		ierr = PetscDSSetJacobian(prob, 0, 0, NULL, NULL, NULL, g3);CHKERRQ(ierr);
		
		/* Set boundary conditions */
    ierr = DMPlexAddBoundary(cdm, PETSC_TRUE, "wall", "marker", 0, (void (*)())bcs, 1, &id, user);CHKERRQ(ierr);
    ierr = DMPlexGetCoarseDM(cdm, &cdm);CHKERRQ(ierr);
  }
  ierr = PetscFEDestroy(&fe);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "FormFunctionGradient"
PetscErrorCode FormFunctionGradient(Tao tao, Vec x, PetscReal *f, Vec g, void *ctx)
{
  AppCtx *user = (AppCtx*)ctx;
  PetscScalar       xtHx;
  PetscErrorCode    ierr;

  PetscFunctionBegin;
  ierr = MatMult(user->A,x,g);CHKERRQ(ierr);
  ierr = VecDot(x,g,&xtHx);CHKERRQ(ierr);
  ierr = VecDot(x,user->f,f);CHKERRQ(ierr);
  *f += 0.5*xtHx;
  ierr = VecAXPY(g,1.0,user->f);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "FormHessian"
PetscErrorCode FormHessian(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx)
{
  PetscFunctionBegin;
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc, char **argv)
{
  DM             			dm;          /* Problem specification */
  KSP           			ksp;         /* nonlinear solver */
  SNES           			snes;        /* nonlinear solver */
  Tao						 			tao;				 /* Optimization solver */
  TaoConvergedReason 	reason;			 /* Convergence reason */
  AppCtx         			user;        /* user-defined work context */
  PetscInt       			its;         /* iterations for convergence */
	PetscMPIInt	 	      size;
  PetscErrorCode 			ierr;
  
  /* Begin PETSc program */
	ierr = PetscInitialize(&argc, &argv, NULL, help);CHKERRQ(ierr);
	ierr = MPI_Comm_size(PETSC_COMM_WORLD, &size);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "\n========================\n");CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "  %d processors:",size);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "\n========================\n");CHKERRQ(ierr);
	ierr = ProcessOptions(PETSC_COMM_WORLD, &user);CHKERRQ(ierr);
	ierr = PetscTime(&v_create);CHKERRQ(ierr);
  ierr = CreateMesh(PETSC_COMM_WORLD, &user, &dm);CHKERRQ(ierr);
	ierr = PetscTime(&v_create_total);CHKERRQ(ierr);
  ierr = SNESCreate(PETSC_COMM_WORLD, &snes);CHKERRQ(ierr);
  ierr = SNESSetDM(snes, dm);CHKERRQ(ierr);
  ierr = DMSetApplicationContext(dm, &user);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "\nSetting up and initializing problem... ");CHKERRQ(ierr);
  ierr = SetupProblem(dm, &user);CHKERRQ(ierr);
	
	/* Create global matrices and vectors */
  ierr = DMCreateGlobalVector(dm, &user.u);CHKERRQ(ierr);
	ierr = VecDuplicate(user.u, &user.r);CHKERRQ(ierr);
	ierr = VecDuplicate(user.u, &user.f);CHKERRQ(ierr);
	ierr = VecDuplicate(user.u, &user.b);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) user.u, "solution");CHKERRQ(ierr);
  ierr = DMSetMatType(dm,MATAIJ);CHKERRQ(ierr);
  ierr = DMCreateMatrix(dm, &user.J);CHKERRQ(ierr);
  user.A = user.J;

	/* Set residual and jacobian evaluation functions */
  ierr = DMSNESSetFunctionLocal(dm,(PetscErrorCode(*)(DM,Vec,Vec,void*))DMPlexSNESComputeResidualFEM, &user);CHKERRQ(ierr);
  ierr = DMSNESSetJacobianLocal(dm,(PetscErrorCode (*)(DM,Vec,Mat,Mat,void*))DMPlexSNESComputeJacobianFEM, &user);CHKERRQ(ierr);
  ierr = SNESSetJacobian(snes, user.A, user.J, NULL, NULL);CHKERRQ(ierr);
  ierr = SNESSetFromOptions(snes);CHKERRQ(ierr);

	/* Initial guess */
  void (*initialGuess[1])(const PetscReal x[], PetscScalar *, void *ctx) = {initial};
  void (*maximumSol[1])(const PetscReal x[], PetscScalar *, void *ctx) = {maxSol};
	ierr = DMPlexProjectFunction(dm, initialGuess, NULL, INSERT_VALUES, user.u);CHKERRQ(ierr);
	
	/* Set up nonnegative solver if needed */
  if (user.nonneg) {
    ierr = TaoCreate(PETSC_COMM_WORLD,&tao);CHKERRQ(ierr);
		ierr = VecDuplicate(user.u, &user.lb);CHKERRQ(ierr);
		ierr = VecDuplicate(user.u, &user.ub);CHKERRQ(ierr);
		ierr = DMPlexProjectFunction(dm, initialGuess, NULL, INSERT_ALL_VALUES, user.lb);CHKERRQ(ierr);
		ierr = DMPlexProjectFunction(dm, maximumSol, NULL, INSERT_ALL_VALUES, user.ub);CHKERRQ(ierr);
	}
	ierr = PetscPrintf(PETSC_COMM_WORLD, "done\n");CHKERRQ(ierr);
	ierr = PetscTime(&v_solve);CHKERRQ(ierr);
	
	/* Option 1 - classical Galerkin */
	if (!user.nonneg) {
		ierr = PetscPrintf(PETSC_COMM_WORLD, "\n========================\n");CHKERRQ(ierr);
		ierr = PetscPrintf(PETSC_COMM_WORLD, "  Solving Galerkin...");CHKERRQ(ierr);
		ierr = PetscPrintf(PETSC_COMM_WORLD, "\n========================\n");CHKERRQ(ierr);
		
	  ierr = SNESGetKSP(snes,&ksp);CHKERRQ(ierr);
		ierr = SNESComputeFunction(snes,user.u,user.r);CHKERRQ(ierr);
		ierr = SNESComputeJacobian(snes,user.u,user.A,user.A);CHKERRQ(ierr);
		ierr = KSPSolve(ksp,user.r,user.b);CHKERRQ(ierr);
		ierr = VecChop(user.b,1.0e-10);CHKERRQ(ierr);
		ierr = VecWAXPY(user.f,-1.0,user.b,user.u);CHKERRQ(ierr);
		ierr = VecCopy(user.f,user.u);CHKERRQ(ierr);
		ierr = KSPGetIterationNumber(ksp,&its);CHKERRQ(ierr);
	  ierr = PetscPrintf(PETSC_COMM_WORLD, "\nNumber of KSP iterations = %D\n", its);CHKERRQ(ierr);
	
	/* Option 2 - nonnegative methodology */
	} else {
		ierr = PetscPrintf(PETSC_COMM_WORLD, "\n========================\n");CHKERRQ(ierr);
		ierr = PetscPrintf(PETSC_COMM_WORLD, "  Solving non-negative...");CHKERRQ(ierr);
		ierr = PetscPrintf(PETSC_COMM_WORLD, "\n========================\n");CHKERRQ(ierr);
		
		ierr = TaoGetKSP(tao,&ksp);CHKERRQ(ierr);
		ierr = SNESComputeFunction(snes,user.u,user.r);CHKERRQ(ierr);
		ierr = SNESComputeJacobian(snes,user.u,user.A,user.A);CHKERRQ(ierr);
		ierr = MatMult(user.A,user.u,user.b);CHKERRQ(ierr);
		ierr = VecChop(user.r,1.0e-10);CHKERRQ(ierr);
		ierr = VecChop(user.b,1.0e-10);CHKERRQ(ierr);
		ierr = VecWAXPY(user.f,-1.0,user.b,user.r);CHKERRQ(ierr);
		
		/* Set up Tao solver */
  	ierr = TaoSetInitialVector(tao,user.u);CHKERRQ(ierr);
  	ierr = TaoSetVariableBounds(tao,user.lb,user.ub);CHKERRQ(ierr);
  	ierr = TaoSetObjectiveAndGradientRoutine(tao,FormFunctionGradient,(void*)&user);CHKERRQ(ierr);
  	ierr = TaoSetHessianRoutine(tao,user.A,user.A,FormHessian,(void*)&user);CHKERRQ(ierr);
    ierr = TaoSetTolerances(tao,1e-9,1e-12,0,0,0);CHKERRQ(ierr);
    ierr = TaoSetFromOptions(tao);CHKERRQ(ierr);
    
    /* Solve */
    ierr = TaoSolve(tao);CHKERRQ(ierr);
    ierr = TaoGetConvergedReason(tao,&reason);CHKERRQ(ierr);
    if (reason < 0) {
      ierr = PetscPrintf(MPI_COMM_WORLD, "\nTAO failed to converge.\n");CHKERRQ(ierr);
		} else {
			//ierr = KSPGetIterationNumber(ksp,&its);CHKERRQ(ierr);
      ierr = PetscPrintf(MPI_COMM_WORLD, "\nTao converged with reason %D\n", reason);CHKERRQ(ierr);
    }
	}
	ierr = PetscTime(&v_post);CHKERRQ(ierr);
  
  /* Output concentration to VTK */
  if (user.flg_out) {
		ierr = PetscViewerVTKOpen(PETSC_COMM_WORLD,user.output,FILE_MODE_WRITE,&user.viewer);CHKERRQ(ierr);
		ierr = VecView(user.u,user.viewer);CHKERRQ(ierr);
		ierr = PetscViewerDestroy(&user.viewer);CHKERRQ(ierr);
	}
	
  ierr = PetscTime(&v_total);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD, "\n========================\n");CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "  Time summary:");CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "\n========================\n\n");CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "Create DMPlex: %gs\n",v_create_distribute-v_create);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "Distribute and refine DMPlex: %gs\n",v_create_total-v_create_distribute);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "Set up problem: %gs\n",v_solve-v_create_total);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "Solver: %gs\n",v_post-v_solve);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD, "Postprocess into VTU: %gs\n\n",v_total-v_post);CHKERRQ(ierr);
	
	/* Clean up */
	if (user.A != user.J) {ierr = MatDestroy(&user.A);CHKERRQ(ierr);}
	ierr = VecDestroy(&user.r);CHKERRQ(ierr);
	ierr = VecDestroy(&user.b);CHKERRQ(ierr);
  ierr = VecDestroy(&user.f);CHKERRQ(ierr);
	if (user.nonneg) {
		ierr = VecDestroy(&user.lb);CHKERRQ(ierr);
		ierr = VecDestroy(&user.ub);CHKERRQ(ierr);
	  ierr = TaoDestroy(&tao);CHKERRQ(ierr);
	}
  ierr = MatDestroy(&user.J);CHKERRQ(ierr);
  ierr = VecDestroy(&user.u);CHKERRQ(ierr);
  ierr = SNESDestroy(&snes);CHKERRQ(ierr);
  ierr = DMDestroy(&dm);CHKERRQ(ierr);
  ierr = PetscFinalize();
  return 0;
}