#include "petsc.h"
#include "petscdmplex.h"
#include <petscdmforest.h>

static char help[] = "Minimal working example using DMPLEX/DMFOREST";

PetscErrorCode CreateAdaptivityLabel(DM dmForest, DMLabel *adaptLabel)
{
  DMLabel identLabel;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = DMLabelCreate(PETSC_COMM_SELF, "adapt", adaptLabel); CHKERRQ(ierr);
  ierr = DMLabelSetDefaultValue(*adaptLabel, DM_ADAPT_COARSEN); CHKERRQ(ierr);
  ierr = DMGetLabel(dmForest, "identity", &identLabel); CHKERRQ(ierr);
  /* Just adapt cell 0 */
  ierr = DMLabelSetValue(*adaptLabel, 0, DM_ADAPT_REFINE); CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

PetscErrorCode AddIdentityLabel(DM dm)
{
  PetscInt pStart, pEnd, p;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = DMCreateLabel(dm, "identity"); CHKERRQ(ierr);
  ierr = DMPlexGetChart(dm, &pStart, &pEnd); CHKERRQ(ierr);
  for (p = pStart; p < pEnd; p++)
  {
    ierr = DMSetLabelValue(dm, "identity", p, p);
    CHKERRQ(ierr);
  }
  PetscFunctionReturn(0);
}

int main(int argc, char **argv)
{

  PetscErrorCode ierr;
  DM dm, dmForest, dmPostForest;
  const PetscInt dim = 3;
  const PetscBool cellSimplex = PETSC_FALSE;
  const DMBoundaryType Periodic[] = { DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE };
  const PetscInt iSize[] = { 5, 5, 5 };
  const PetscReal Xmin[] = { 0.0, 0.0, 0.0 };
  const PetscReal Xmax[] = { 1.0, 1.0, 1.0 };
  PetscFE fe;
  const PetscInt numFields = 1;
  PetscInt *numComp, *numDof;
  PetscInt i;
  Vec GlobalVector;
  Vec postForestCCVector;
  PetscSection sectionCellCenter;
  DMLabel adaptLabel;
  PetscSection NewPreSection, NewPostSection;
  DM dmForestCloned, dmPostForestCloned;
  Vec faceCenteredForestVector, faceCenteredPostForestVector;


  PetscFunctionBegin;
  /* Set-up Mesh and DM */
  ierr = PetscInitialize(&argc, &argv, (char*) 0, help); CHKERRQ(ierr);
  ierr = DMPlexCreateBoxMesh(PETSC_COMM_WORLD, dim, cellSimplex, iSize, Xmin, Xmax, Periodic, PETSC_TRUE, &dm); CHKERRQ(ierr);

  /* Add label to DM */
  ierr = AddIdentityLabel(dm); CHKERRQ(ierr);
  /* Creating an FE object, just to use P8EST  */
  ierr = PetscFECreateDefault(PETSC_COMM_WORLD, 3, 1, PETSC_FALSE, NULL, PETSC_DEFAULT, &fe); CHKERRQ(ierr);
  ierr = DMSetField(dm, 0, NULL, (PetscObject) fe); CHKERRQ(ierr);
  ierr = PetscFEDestroy(&fe); CHKERRQ(ierr);
  ierr = DMCreateDS(dm); CHKERRQ(ierr);

  /* Data saved at cell centers, create section to do so: */
  ierr = PetscCalloc2(numFields, &numComp, numFields * (dim + 1), &numDof); CHKERRQ(ierr);
  for (i = 0; i < numFields; i++)
  {
    numComp[i] = 1;
    numDof[i * (dim + 1) + dim] = 1;
  }
  ierr = DMSetNumFields(dm, numFields); CHKERRQ(ierr);
  ierr = DMPlexCreateSection(dm, NULL, numComp, numDof, 0, NULL, NULL, NULL, NULL, &sectionCellCenter); CHKERRQ(ierr);
  ierr = PetscSectionSetFieldName(sectionCellCenter, 0, "Uvelocity"); CHKERRQ(ierr);

  /* Set the section onto the standard DM */
  ierr = DMSetLocalSection(dm, sectionCellCenter); CHKERRQ(ierr);
  ierr = PetscFree2(numComp, numDof); CHKERRQ(ierr);
  ierr = PetscSectionDestroy(&sectionCellCenter); CHKERRQ(ierr);

  /* Create a vector on the dm */
  ierr = DMCreateGlobalVector(dm, &GlobalVector); CHKERRQ(ierr);
  ierr = VecSet(GlobalVector, 1.0); CHKERRQ(ierr);


  /* Create a dmForest, which is the starting P4EST dm */
  ierr = DMCreate(PETSC_COMM_WORLD, &dmForest); CHKERRQ(ierr);
  ierr = DMSetType(dmForest, DMP8EST); CHKERRQ(ierr);
  ierr = DMCopyDisc(dm, dmForest); CHKERRQ(ierr);
  ierr = DMForestSetBaseDM(dmForest, dm); CHKERRQ(ierr);
  ierr = DMForestSetMinimumRefinement(dmForest, 0); CHKERRQ(ierr);
  ierr = DMForestSetInitialRefinement(dmForest, 0); CHKERRQ(ierr);
  ierr = DMForestSetMaximumRefinement(dmForest, 2); CHKERRQ(ierr);
  ierr = DMSetFromOptions(dmForest); CHKERRQ(ierr);
  ierr = DMSetUp(dmForest); CHKERRQ(ierr);

  /* Now refine the dmForest to dmPostForest */
   ierr = CreateAdaptivityLabel(dmForest, &adaptLabel); CHKERRQ(ierr);
   ierr = DMForestTemplate(dmForest, PETSC_COMM_WORLD, &dmPostForest); CHKERRQ(ierr);
   ierr = DMForestSetAdaptivityLabel(dmPostForest, adaptLabel); CHKERRQ(ierr);
   ierr = DMLabelDestroy(&adaptLabel); CHKERRQ(ierr);
   ierr = DMSetUp(dmPostForest); CHKERRQ(ierr);

   /* First, create a "standard" cell face vector and translate from one forest to another */
   ierr = DMCreateGlobalVector(dmPostForest, &postForestCCVector); CHKERRQ(ierr);
   ierr = DMForestTransferVec(dmForest, GlobalVector, dmPostForest, postForestCCVector, PETSC_FALSE, 0.0); CHKERRQ(ierr);
   ierr = VecDestroy(&GlobalVector); CHKERRQ(ierr);


   /*****************************************************************************/
   /* Create a new section and clone DM's to translate a different vector.
    * Clone the dmForest pre and post,  set a new section onto them and then
    * use them to Transfer a vector                                             */
   /*****************************************************************************/
   /* Clone pre and post forests */
   ierr = DMClone(dmForest, &dmForestCloned); CHKERRQ(ierr);
   ierr = DMClone(dmPostForest, &dmPostForestCloned); CHKERRQ(ierr);

   ierr = PetscCalloc2(numFields, &numComp, numFields * (dim + 1), &numDof); CHKERRQ(ierr);
   /* Data saved at face centers: */
   for (i = 0; i < numFields; i++)
   {
     numComp[i] = 1;
     numDof[i * (dim + 1) + 3] = 1;
   }

   ierr = DMSetNumFields(dmForestCloned, numFields); CHKERRQ(ierr);
   ierr = DMSetNumFields(dmPostForestCloned, numFields); CHKERRQ(ierr);

   /* Create and Set a new Section to the cloned Forest DMs using a face centered layout */
   {
     DM plexCloned, plexPostCloned;

     ierr = DMConvert(dmForestCloned, DMPLEX, &plexCloned);CHKERRQ(ierr);
     ierr = DMConvert(dmPostForestCloned, DMPLEX, &plexPostCloned);CHKERRQ(ierr);
     ierr = DMPlexCreateSection(plexCloned, NULL, numComp, numDof, 0, NULL, NULL, NULL, NULL, &NewPreSection); CHKERRQ(ierr);
     ierr = DMPlexCreateSection(plexPostCloned, NULL, numComp, numDof, 0, NULL, NULL, NULL, NULL, &NewPostSection); CHKERRQ(ierr);
     ierr = DMDestroy(&plexCloned);CHKERRQ(ierr);
     ierr = DMDestroy(&plexPostCloned);CHKERRQ(ierr);
   }
   ierr = DMSetLocalSection(dmForestCloned, NewPreSection); CHKERRQ(ierr);
   ierr = DMSetLocalSection(dmPostForestCloned, NewPostSection); CHKERRQ(ierr);
   ierr = PetscFree2(numComp, numDof); CHKERRQ(ierr);
   ierr = PetscSectionDestroy(&NewPreSection); CHKERRQ(ierr);
   ierr = PetscSectionDestroy(&NewPostSection); CHKERRQ(ierr);

   /* With the new section, create a vector and translate from one forest to another */
    ierr = DMCreateGlobalVector(dmForestCloned, &faceCenteredForestVector); CHKERRQ(ierr);
    ierr = DMCreateGlobalVector(dmPostForestCloned, &faceCenteredPostForestVector); CHKERRQ(ierr);
    ierr = VecSet(faceCenteredForestVector, 1.0); CHKERRQ(ierr);
    ierr = DMForestTransferVec(dmForestCloned, faceCenteredForestVector, dmPostForestCloned, faceCenteredPostForestVector, PETSC_FALSE, 0.0); CHKERRQ(ierr);


   /* Clean up "face-center" vectors and cloned DMs*/
   ierr = VecDestroy(&faceCenteredForestVector); CHKERRQ(ierr);
   ierr = VecDestroy(&faceCenteredPostForestVector); CHKERRQ(ierr);
   ierr = DMDestroy(&dmForestCloned); CHKERRQ(ierr);
   ierr = DMDestroy(&dmPostForestCloned); CHKERRQ(ierr);

  /* Finalize */
  ierr = VecDestroy(&GlobalVector); CHKERRQ(ierr);
  ierr = DMDestroy(&dm); CHKERRQ(ierr);
  ierr = DMDestroy(&dmForest); CHKERRQ(ierr);
  ierr = PetscFinalize(); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}