#include <petscdmplex.h>
#include <petscts.h>
#include <petscsnes.h>
#include <petscds.h>
#include <petscviewerhdf5.h>

typedef struct {
  int debug;
  int dim;
  PetscBool interpolate;
  PetscReal refinementLimit;
  PetscErrorCode (**exactFuncs)(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *ctx);
} AppCtx;

PetscErrorCode CreateMesh(MPI_Comm, DM*, AppCtx*);
PetscErrorCode SetupDiscretization(DM, AppCtx*);

PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *ctx)
{
  u[0] = 0.0;
  return 0;
}

PetscErrorCode analytic_temp_2d(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *ctx)
{
  *u = 2.0*time + x[0]*x[0] + x[1]*x[1];
  /*
    if (x[0] == 0) {
    *u = time * 1.0;
    } else {
    *u = 0.0;
    }
  */
  return 0;
}

void f0_temp(PetscInt dim, PetscInt Nf, PetscInt NfAux,
          const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[],
          const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[],
          PetscReal t, const PetscReal x[], PetscScalar f0[])
{
  f0[0] = u_t[0] - 2.0;
}

void f1_temp(PetscInt dim, PetscInt Nf, PetscInt NfAux,
	     const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[],
	     const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[],
	     PetscReal t, const PetscReal x[], PetscScalar f1[])
{
  PetscInt d;
  for (d = 0; d < dim; ++d) {
    f1[d] = u_x[d];
  }
}

void g3_temp(PetscInt dim, PetscInt Nf, PetscInt NfAux,
	     const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[],
	     const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[],
	     PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscScalar g3[])
{
  PetscInt d;
  for (d = 0; d < dim; ++d) {
    g3[d*dim+d] = 1.0;
  }
}

void g0_temp(PetscInt dim, PetscInt Nf, PetscInt NfAux,
	     const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[],
	     const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[],
	     PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscScalar g0[])
{
  g0[0] = 1.0;
}

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc, char **argv)
{
  PetscErrorCode ierr;
  AppCtx ctx;
  DM dm;
  TS ts;
  Vec u, r;
  PetscReal t = 0.0;
  PetscReal l2error = 0.0; /* L_2 error in the solution */
  
  ierr = PetscInitialize(&argc, &argv, NULL, NULL);
  if (ierr) {
    return ierr;
  }

  /* Default context initialization */
  ctx.debug = 0;
  ctx.dim = 2;
  ctx.interpolate = PETSC_TRUE;
  ctx.refinementLimit = 0.0;

  ierr = CreateMesh(PETSC_COMM_WORLD, &dm, &ctx);CHKERRQ(ierr);
  ierr = DMSetApplicationContext(dm, &ctx);CHKERRQ(ierr);
  ierr = PetscMalloc(1, &ctx.exactFuncs);CHKERRQ(ierr);
  ierr = SetupDiscretization(dm, &ctx);CHKERRQ(ierr);

  ierr = DMCreateGlobalVector(dm, &u);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) u, "temperature");CHKERRQ(ierr);
  ierr = VecDuplicate(u, &r);CHKERRQ(ierr);

  ierr = TSCreate(PETSC_COMM_WORLD, &ts);CHKERRQ(ierr);
  ierr = TSSetDM(ts, dm);CHKERRQ(ierr);
  ierr = DMTSSetIFunctionLocal(dm, DMPlexTSComputeIFunctionFEM, &ctx);CHKERRQ(ierr);
  ierr = DMTSSetIJacobianLocal(dm, DMPlexTSComputeIJacobianFEM, &ctx);CHKERRQ(ierr);
  ierr = TSSetType(ts, TSBEULER);CHKERRQ(ierr);
  ierr = TSSetDuration(ts, 1, 1.0);CHKERRQ(ierr);
  ierr = TSSetInitialTimeStep(ts, 0.0, 0.001);CHKERRQ(ierr);
  ierr = TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER);CHKERRQ(ierr);
  
  ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
  ierr = TSSolve(ts, u);CHKERRQ(ierr);

  ierr = TSGetTime(ts, &t);CHKERRQ(ierr);
  ierr = DMComputeL2Diff(dm, t, ctx.exactFuncs, NULL, u, &l2error);CHKERRQ(ierr);
  if (l2error < 1.0e-11) {
    ierr = PetscPrintf(PETSC_COMM_WORLD, "L_2 Error: < 1.0e-11\n");CHKERRQ(ierr);
  }
  else {
    ierr = PetscPrintf(PETSC_COMM_WORLD, "L_2 Error: %g\n", l2error);CHKERRQ(ierr);
  }
  
  // ierr = DMProjectFunction(dm, t, ctx.exactFuncs, NULL, INSERT_ALL_VALUES, u);CHKERRQ(ierr);
  ierr = VecViewFromOptions(u, NULL, "-sol_view");CHKERRQ(ierr);
  
  ierr = VecDestroy(&u);CHKERRQ(ierr);
  ierr = VecDestroy(&r);CHKERRQ(ierr);
  ierr = TSDestroy(&ts);CHKERRQ(ierr);
  ierr = DMDestroy(&dm);CHKERRQ(ierr);
  ierr = PetscFree(ctx.exactFuncs);CHKERRQ(ierr);
  ierr = PetscFinalize();
  return ierr;
}

#undef __FUNCT__
#define __FUNCT__ "SetupDiscretization"
PetscErrorCode SetupDiscretization(DM dm, AppCtx* ctx)
{
  PetscErrorCode ierr;
  const PetscInt id = 1;
  const PetscInt dim = ctx->dim;
  PetscFE temp_fe;
  PetscDS temp_ds;
    
  PetscFunctionBeginUser;
  ierr = PetscFECreateDefault(dm, dim, 1, PETSC_FALSE, "temp_", -1, &temp_fe);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) temp_fe, "temperature");CHKERRQ(ierr);
  ierr = DMGetDS(dm, &temp_ds);CHKERRQ(ierr);
  ierr = PetscDSSetDiscretization(temp_ds, 0, (PetscObject) temp_fe);CHKERRQ(ierr);

  ierr = PetscDSSetResidual(temp_ds, 0, f0_temp, f1_temp);CHKERRQ(ierr);
  ierr = PetscDSSetJacobian(temp_ds, 0, 0, g0_temp, NULL, NULL, g3_temp);CHKERRQ(ierr);

  switch (ctx->dim) {
  case 2:
    ctx->exactFuncs[0] = analytic_temp_2d;
    break;
  default:
    SETERRQ1(PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Invalid dimension %d", ctx->dim);
  }
  ierr = PetscDSAddBoundary(temp_ds, DM_BC_ESSENTIAL, "wall", "marker", 0, 0, NULL, (void (*)()) ctx->exactFuncs[0], 1, &id, ctx);CHKERRQ(ierr);
  
  ierr = PetscFEDestroy(&temp_fe);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "CreateBCLabel"
PetscErrorCode CreateBCLabel(DM dm, const char name[])
{
  DMLabel label;
  PetscErrorCode ierr;

  PetscFunctionBeginUser;
  ierr = DMCreateLabel(dm, name);CHKERRQ(ierr);
  ierr = DMGetLabel(dm, name, &label);CHKERRQ(ierr);
  ierr = DMPlexMarkBoundaryFaces(dm, label);CHKERRQ(ierr);
  ierr = DMPlexLabelComplete(dm, label);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

#undef __FUNCT__
#define __FUNCT__ "CreateMesh"
PetscErrorCode CreateMesh(MPI_Comm comm, DM *dm, AppCtx *ctx)
{
  PetscErrorCode ierr;
  const int dim = ctx->dim;
  const PetscReal refinementLimit = ctx->refinementLimit;
  PetscInt cells[3] = {1, 1, 1}; /* coarse mesh is one cell; refine from there */

  PetscFunctionBeginUser;
  ierr = DMPlexCreateHexBoxMesh(comm, dim, cells, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, dm);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) *dm, "Mesh");CHKERRQ(ierr);

  PetscBool hasLabel;
  ierr = DMHasLabel(*dm,"marker",&hasLabel);CHKERRQ(ierr);
  if (!hasLabel) {
    ierr = CreateBCLabel(*dm, "marker");CHKERRQ(ierr);
  }
  {
    DM refinedMesh = NULL;
    DM distributedMesh = NULL;

    /* Refine mesh using a volume constraint */
    ierr = DMPlexSetRefinementLimit(*dm, refinementLimit);CHKERRQ(ierr);
    ierr = DMRefine(*dm, comm, &refinedMesh);CHKERRQ(ierr);
    if (refinedMesh) {
      const char *name;

      ierr = PetscObjectGetName((PetscObject) *dm, &name);CHKERRQ(ierr);
      ierr = PetscObjectSetName((PetscObject) refinedMesh, name);CHKERRQ(ierr);
      ierr = DMDestroy(dm);CHKERRQ(ierr);
      *dm = refinedMesh;
    }
    /* Distribute mesh over processes */
    ierr = DMPlexDistribute(*dm, 0, NULL, &distributedMesh);CHKERRQ(ierr);
    if (distributedMesh) {
      ierr = DMDestroy(dm);CHKERRQ(ierr);
      *dm = distributedMesh;
    }
  }
  ierr = DMSetFromOptions(*dm);CHKERRQ(ierr);
  ierr = DMViewFromOptions(*dm, NULL, "-dm_view");CHKERRQ(ierr);
  PetscFunctionReturn(0);
}