#include <petsc.h>
#include <petscdmforest.h>
#include <petscviewerhdf5.h>

int main (int argc, char **argv)
{
  DM             base, forest, plex;
  PetscSection   s;
  Vec            g, l;
  PetscViewer    viewer;
  PetscInt       vStart, vEnd, v;
  PetscBool      writeEarly = PETSC_FALSE, nog2l = PETSC_FALSE;
  PetscErrorCode ierr;

  ierr = PetscInitialize(&argc, &argv, NULL, NULL);if (ierr) return ierr;
  ierr = PetscOptionsGetBool(NULL, NULL, "-write_early", writeEarly, &writeEarly);CHKERRQ(ierr);
  ierr = PetscOptionsGetBool(NULL, NULL, "-no_g2l", nog2l, &nog2l);CHKERRQ(ierr);

  /* Create a base DMPlex mesh */
  ierr = DMPlexCreateBoxMesh(PETSC_COMM_WORLD, 2, PETSC_FALSE, NULL, NULL, NULL, NULL, PETSC_TRUE, &base);CHKERRQ(ierr);
  ierr = DMSetFromOptions(base);CHKERRQ(ierr);

  /* Covert Plex mesh to Forest and destroy base */
  ierr = DMCreate(PETSC_COMM_WORLD, &forest);CHKERRQ(ierr);
  ierr = DMSetType(forest, DMP4EST);CHKERRQ(ierr);
  ierr = DMForestSetBaseDM(forest, base);CHKERRQ(ierr);
  ierr = DMSetUp(forest);CHKERRQ(ierr);
  ierr = DMDestroy(&base);CHKERRQ(ierr);
  ierr = DMViewFromOptions(forest, NULL, "-dm_view");CHKERRQ(ierr);

  /*  Setup the section */
  ierr = DMConvert(forest, DMPLEX, &plex);CHKERRQ(ierr);
  ierr = DMPlexGetDepthStratum(plex, 0, &vStart, &vEnd);CHKERRQ(ierr);
  ierr = PetscSectionCreate(PetscObjectComm((PetscObject) forest), &s);CHKERRQ(ierr);
  ierr = PetscSectionSetChart(s, vStart, vEnd);CHKERRQ(ierr);
  for(v = vStart; v < vEnd; ++v) {ierr = PetscSectionSetDof(s, v, 1);CHKERRQ(ierr);}
  ierr = PetscSectionSetUp(s);CHKERRQ(ierr);
  ierr = DMSetLocalSection(forest, s);CHKERRQ(ierr);
  ierr = PetscSectionDestroy(&s);CHKERRQ(ierr);
  ierr = DMDestroy(&plex);CHKERRQ(ierr);

  /*  Save a vector*/
  ierr = DMCreateGlobalVector(forest, &g);CHKERRQ(ierr);
  ierr = PetscObjectSetName((PetscObject) g, "g");CHKERRQ(ierr);
  ierr = VecSet(g, 1.0);CHKERRQ(ierr);

  ierr = PetscViewerHDF5Open(PETSC_COMM_WORLD, "forestHDF.h5", FILE_MODE_WRITE, &viewer);
  if (writeEarly) {ierr = VecView(g, viewer);CHKERRQ(ierr);}

  ierr = DMCreateLocalVector(forest, &l);CHKERRQ(ierr);
  if (!nog2l) {ierr = DMGlobalToLocal(forest, g, INSERT_VALUES, l);CHKERRQ(ierr);}

  ierr = VecView(g, viewer);CHKERRQ(ierr);
  ierr = PetscViewerDestroy(&viewer);CHKERRQ(ierr);

  ierr = VecDestroy(&g);CHKERRQ(ierr);
  ierr = VecDestroy(&l);CHKERRQ(ierr);
  ierr = DMDestroy(&forest);CHKERRQ(ierr);
  ierr = PetscFinalize();
  return ierr;
}