static const char help[] = "Test of FEA DMPlex w/CAD functionality";

#include <petsc.h>

/* User-defined Work Context  */
typedef struct {
    PetscInt  dummy;
    char filename[PETSC_MAX_PATH_LEN];    // context containing CAD filename from command line
} AppCtx;


void f0_function(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[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f0[])
{
  f0[0] = 0.0;
}

void f1_function(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[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f1[])
{
  for(PetscInt d = 0; d < dim; ++d) f1[d] = u_x[d];
}

void g3_uu_function(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[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g3[])
{
  for (PetscInt d=0; d < dim; ++d) g3[d*dim + d] = 1.0;
}

//void bc_inlet(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar bcval[])
static PetscErrorCode bc_inlet(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, void *ctx)
{
  u[0] = 1400.0;
  return PETSC_SUCCESS;
  //bcval[0] = 1400.0;
}

//void bc_outlet(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar bcval[])
static PetscErrorCode bc_outlet(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, void *ctx)
{
  u[0] = 100.0;
  return PETSC_SUCCESS;
  //bcval[0] = 100.0;
}

/* Procees Options - This should be removed once creation bug is fixed by Prof. Knepley */
static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options)
{
  PetscFunctionBeginUser;
  options->filename[0] = '\0';

  PetscOptionsBegin(comm, "", "FEA DMPlex w/CAD Options", "DMPlex w/CAD");
  PetscOptionsString("-filename", "The CAD/Geometry file", "ex18.c", options->filename, options->filename, PETSC_MAX_PATH_LEN, NULL);
  PetscOptionsEnd();
  PetscFunctionReturn(0);
}

static PetscErrorCode CADCreateLabels(DM dmSurface, DM dmVolume)
{
  DMLabel  inletLabel, outletLabel;
  PetscInt nStart, nEnd, eStart, eEnd, faceID;
  
  PetscFunctionBegin;
  PetscCall(DMCreateLabel(dmVolume, "inlet"));
  PetscCall(DMCreateLabel(dmVolume, "outlet"));
  PetscCall(DMGetLabel(dmVolume, "inlet", &inletLabel));
  PetscCall(DMGetLabel(dmVolume, "outlet", &outletLabel));
  PetscCall(DMPlexGetDepthStratum(dmSurface, 0, &nStart, &nEnd));
  PetscCall(DMPlexGetDepthStratum(dmSurface, 1, &eStart, &eEnd));
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [nStart = %d  ||  nEnd = %d] \n", nStart, nEnd));
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [eStart = %d  ||  eEnd = %d] \n", eStart, eEnd));
  for (PetscInt ii = nStart; ii < nEnd; ++ii) {
    PetscCall(DMGetLabelValue(dmSurface, "EGADS Face ID", ii, &faceID));
    
    PetscPrintf(PETSC_COMM_SELF, "Vertex %d FaceID %d\n", ii, faceID);
    if (faceID == 14) {PetscCall(DMLabelSetValue(inletLabel, ii, faceID));}
    if (faceID == 7)  {PetscCall(DMLabelSetValue(outletLabel, ii, faceID));}
  }
  for (PetscInt ii = eStart; ii < eEnd; ++ii) {
    PetscCall(DMGetLabelValue(dmSurface, "EGADS Face ID", ii, &faceID));
    
    PetscPrintf(PETSC_COMM_SELF, "Edge %d FaceID %d\n", ii, faceID);
    if (faceID == 14) {PetscCall(DMLabelSetValue(inletLabel, ii, faceID));}
    if (faceID == 7)  {PetscCall(DMLabelSetValue(outletLabel, ii, faceID));}
  }
  PetscFunctionReturn(PETSC_SUCCESS);
}

/* Main Function */
int main(int argc, char **argv)
{
  // Define PETSc Variables
  DM            dm, dmSurface;    // Unstructured Grid
  SNES          snes;             // Nonlinear Solver
  Vec           temp;             // Solutions
  AppCtx        ctx;      // User-defined Work Context
  PetscInt      dim;      // DM Geometric Dimension
  PetscBool     simplex;  // Is DMPlex Simplex type?
  PetscFE       fe;       // PETSc Finite Element Object
  PetscDS       ds;       // PETSc Discretization System Object
  PetscMPIInt   rank;     // PETSc MPI Processor ID
  
  PetscFunctionBeginUser;
  PetscCall(PetscInitialize(&argc, &argv, NULL, help));   // Initialize PETSc
  PetscCall(ProcessOptions(PETSC_COMM_WORLD, &ctx));      // Process Options :: Here we get the filename from the command line options

  // Create DMPlex from CAD file
  PetscCall(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));      // Get Rank of Current Processor
  PetscCall(DMPlexCreateFromFile(PETSC_COMM_WORLD, ctx.filename, "EGADS", PETSC_TRUE, &dmSurface)); // Return dm created from CAD file
  
  // Generate Volumetric Mesh
  PetscCall(DMPlexGenerate(dmSurface, "tetgen", PETSC_TRUE, &dm));   // Generate Volumetric Mesh
  PetscCall(CADCreateLabels(dmSurface, dm)); // Setup DMLabels for BC
  PetscCall(DMDestroy(&dmSurface));      // Destroy DM dmSurface no longer needed
  PetscCall(DMSetApplicationContext(dm, &ctx));            // Link context to dm
  PetscCall(DMViewFromOptions(dm, NULL, "-dm_view"));     // Write DM to file (hdf5)
  PetscCall(DMView(dm, PETSC_VIEWER_STDOUT_SELF));
  
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [CREATED MESH] \n"));
  
  // Setup SNES and link to DM
  PetscCall(SNESCreate(PETSC_COMM_WORLD, &snes));         // Create SNES object
  PetscCall(SNESSetDM(snes, dm));                         // Link SNES object to DM
  
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [SETUP SNES] \n"));
  
  // Setup Discretization
  PetscCall(DMGetDimension(dm, &dim));        // Get Geometric Dimension of the DM (2D or 3D)
  PetscCall(DMPlexIsSimplex(dm, &simplex));   // Check if DMPlex is Simplex Type
  PetscCall(PetscFECreateDefault(PetscObjectComm((PetscObject)dm), dim, 1, simplex, NULL, -1, &fe));    // Define and Create PetscFE Object
  PetscCall(DMAddField(dm, NULL, (PetscObject)fe));        // Add Field to Discretization object. In this case, the temperature field.
  PetscCall(PetscFEDestroy(&fe));
  PetscCall(DMCreateDS(dm));                  // Create The Discretization System (DS) based on the field(s) added to the DM
  
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [SETUP DISCRETIZATION] \n"));
  
  // Setup Residuals & (Optional) Jacobian
  PetscInt    testFieldID = 0;                // This is an assumption. How do we find out which field ID goes with each added field?
  PetscCall(DMGetDS(dm, &ds));                // Get DS associated with DM
  PetscCall(PetscDSSetResidual(ds, testFieldID, (void (*))f0_function, (void (*))f1_function));            // Set Residual Function
  PetscCall(PetscDSSetJacobian(ds, testFieldID, testFieldID, NULL, NULL, NULL, (void (*))g3_uu_function));    // Set Jacobian Function
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [SETUP RESIDUAL & JACOBIAN] \n"));
  
  // Apply inlet face conditions - Dirichlet    
  DMLabel        inletLabel, outletLabel;
  const PetscInt idA = 14, idB = 7;
  PetscCall(DMGetLabel(dm, "inlet", &inletLabel));
  PetscCall(DMGetLabel(dm, "outlet", &outletLabel));
  PetscCall(DMAddBoundary(dm, DM_BC_ESSENTIAL, "inTemp", inletLabel, 1, &idA, 0, 0, NULL, (void (*)(void))bc_inlet, NULL, &ctx, NULL));
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [SETUP BOUNDARY CONDITION inTemp] \n"));
  PetscCall(DMAddBoundary(dm, DM_BC_ESSENTIAL, "outTemp", outletLabel, 1, &idB, 0, 0, NULL, (void (*)(void))bc_outlet, NULL, &ctx, NULL));

  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [SETUP BOUNDARY CONDITIONS] \n"));

  // Create Global Vector, Initialize Values & Name
  PetscCall(DMCreateGlobalVector(dm, &temp));
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [AFTER GLOBAL VECTOR] \n"));
  PetscCall(VecSet(temp, 100.0));
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [AFTER temp set to 100] \n"));
  PetscCall(VecViewFromOptions(temp, NULL, "-init_sol_view"));
  PetscCall(PetscObjectSetName((PetscObject)temp, "temperature"));
  
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [AFTER SET OBJECT NAME] \n"));
  
  PetscCall(DMPlexSetSNESLocalFEM(dm, &ctx, &ctx, &ctx));
  
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [AFTER SNESLocalFEM] \n"));
  
  //Set SNES from Options
  PetscCall(SNESSetFromOptions(snes));
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [AFTER SNESSetFromOptions] \n"));
  
  // Solve System of Equations
  PetscCall(SNESSolve(snes, NULL, temp));
  
  PetscCall(PetscPrintf(PETSC_COMM_SELF, "   [SOLVED PROBLEM] \n"));
  
  //Get Solution and View
  PetscCall(SNESGetSolution(snes, &temp));
  PetscCall(VecViewFromOptions(temp, NULL, "-sol_view"));
  
  // Cleanup
  PetscCall(VecDestroy(&temp));
  PetscCall(SNESDestroy(&snes));
  PetscCall(DMDestroy(&dm));

  PetscCall(PetscFinalize());  
  return 0;
}