#include <petscdmplex.h>
#include <petscsf.h>
#include <petscviewerhdf5.h>

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

  PetscCall( PetscInitialize(&argc, &argv, NULL,"") );
  PetscMPIInt rank;
  PetscCallMPI( MPI_Comm_rank(PETSC_COMM_WORLD,&rank) );
 
  PetscScalar lower[2]={0.,0.};
  PetscScalar upper[2]={1.,1.};
  DM dmMesh;
  PetscScalar dx=0.1;
  PetscInt faces[2];
  DMBoundaryType periodicity[2]={DM_BOUNDARY_NONE,DM_BOUNDARY_NONE};
  PetscBool simplex=PETSC_FALSE;

  PetscCall( PetscOptionsGetBool(NULL,NULL,"-simplex",NULL,&simplex) );
  const PetscBool interpolate=PETSC_TRUE;
  PetscCall( PetscOptionsGetScalar(NULL,NULL,"-dx",&dx,NULL) );
  faces[0] = std::ceil( (upper[0] - lower[0]) / dx );
  faces[1] = std::ceil( (upper[1] - lower[1]) / dx );
  PetscCall( PetscPrintf( PETSC_COMM_WORLD, "Creating mesh with (%d,%d) faces  \n", faces[0],faces[1] ) );
  PetscCall( DMPlexCreateBoxMesh(MPI_COMM_WORLD, 2, simplex, faces, lower, upper, periodicity, interpolate, &(dmMesh)) );
  PetscCall( PetscObjectSetName((PetscObject) dmMesh, "Mesh") );
  PetscCall( DMSetFromOptions(dmMesh) );

  //PARTITION THE MESH
  PetscPartitioner partitioner;
  PetscCall( DMPlexGetPartitioner(dmMesh, &partitioner) );
  //PetscCall(  PetscPartitionerSetType(partitioner, PETSCPARTITIONERPARMETIS) );
  PetscCall(  PetscPartitionerSetFromOptions(partitioner) );
  DM dmDist=NULL;
  PetscCall( DMSetBasicAdjacency(dmMesh, PETSC_TRUE, PETSC_TRUE) );
  PetscCall( DMPlexDistribute(dmMesh, 1, NULL, &dmDist) );
  if (dmDist) {PetscCall( DMDestroy(&(dmMesh)) ); dmMesh = dmDist; }

  PetscInt     cStart, cEnd;
  PetscCall( DMPlexGetHeightStratum(dmMesh, 0, &cStart, &cEnd) );

  PetscSF sfPoint;
  PetscCall( DMGetPointSF(dmMesh, &sfPoint) );
  //PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD,PETSC_VIEWER_ASCII_INFO_DETAIL);
  //PetscSFView(sfPoint,PETSC_VIEWER_STDOUT_WORLD);
  //PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD);

  PetscInt numLeaves;
  const PetscInt    *localPoints;
  PetscCall( PetscSFGetGraph(sfPoint, NULL, &numLeaves, &localPoints, NULL) );
  DMLabel inghostLabel;
  PetscCall( DMCreateLabel(dmMesh, "inghosts") );
  PetscCall( DMGetLabel(dmMesh, "inghosts", &inghostLabel) );
  for (PetscInt i=0; i<numLeaves; i++){
    //PetscPrintf(PETSC_COMM_SELF,"[%d] %D is inner ghost\n",user.rank,localPoints[i]);
    PetscCall( DMLabelSetValue(inghostLabel, localPoints[i], 1) );
  }

  PetscInt numRoots;
  const PetscSFNode *remotePoints;
  PetscCall( PetscSFGetGraph(sfPoint, &numRoots, &numLeaves, &localPoints, &remotePoints) );
  for (PetscInt i=0; i<numLeaves; i++)
    if (localPoints[i]<cEnd)
      PetscCall( PetscPrintf(PETSC_COMM_SELF,"[on rank %d] %d <- (%d,%d)\n",rank,localPoints[i],remotePoints[i].rank,remotePoints[i].index) );

  /* Setup the section for cell averages*/
  PetscSection sUavg;
  {
    const PetscInt numFields  = 3;
    const PetscInt numComp[numFields] = {1,2,1};
    const PetscInt numDof[numFields*3] = { 0,0,1 , 0,0,2, 0,0,1};
    const PetscInt numBC = 0;
    PetscCall(DMSetNumFields(dmMesh, numFields));
    PetscCall(DMPlexCreateSection(dmMesh, NULL, numComp, numDof, numBC, NULL, NULL, NULL, NULL, &sUavg));
    /* Name the Field variables */
    PetscCall(PetscSectionSetFieldName(sUavg, 0, "rho"));
    PetscCall(PetscSectionSetFieldName(sUavg, 1, "M"));
    PetscCall(PetscSectionSetFieldName(sUavg, 2, "E"));
    //PetscCall(PetscSectionView(section, PETSC_VIEWER_STDOUT_WORLD));
  }

//PetscSF      sectionSF;
//PetscInt    *remoteOffsets;
///** Construct the communication pattern for halo exchange between local vectors */
///* Get the point SF: an object that says which copies of mesh points (cells,
 //* vertices, faces, edges) are copies of points on other processes */
////PetscCall(DMGetPointSF(dm, &pointSF));
///* Relate the locations of ghost degrees of freedom on this process
 //* to their locations of the non-ghost copies on a different process */
//PetscCall(PetscSFCreateRemoteOffsets(sfPoint, sUavg, sUavg, &remoteOffsets));
///* Use that information to construct a star forest for halo exchange
 //* for data described by the local section */
//PetscCall(PetscSFCreateSectionSF(sfPoint, sUavg, remoteOffsets, sUavg, &sectionSF));
//PetscCall(PetscFree(remoteOffsets));

  //PetscCall( DMSetSectionSF(dmMesh, sectionSF) );

  PetscCall( DMSetLocalSection(dmMesh, sUavg) );
  //maybe we need this? Alas, it does not make a difference...
  PetscCall( PetscSectionDestroy(&sUavg) );

  Vec solGlob, solLoc;  
  PetscCall( DMCreateGlobalVector(dmMesh, &solGlob) );
  PetscCall( PetscObjectSetName((PetscObject) solGlob , "Uavg") );
  PetscCall( DMCreateLocalVector(dmMesh, &solLoc) );

  //Fill in local values and then loc2glob
  PetscScalar *values, *cellVal;
  PetscCall( VecGetArray(solLoc, &values) );
  for (PetscInt c=cStart; c<cEnd; ++c){
    PetscCall( DMPlexPointLocalRef(dmMesh, c, values, &cellVal) );
    for (int m=0; m<4; ++m) cellVal[m] = rank + m/10.;
  }
  PetscCall( VecRestoreArray(solLoc, &values) );
  PetscCall( DMLocalToGlobal(dmMesh, solLoc, INSERT_VALUES, solGlob) );

  //output the global solution
  PetscViewer h5Viewer;
  PetscCall(PetscViewerCreate(PETSC_COMM_WORLD, &h5Viewer));
  PetscCall(PetscViewerHDF5Open(PETSC_COMM_WORLD, "output.h5", FILE_MODE_WRITE, &h5Viewer));
  PetscCall(VecView(solGlob, h5Viewer));
  PetscCall(PetscViewerDestroy(&h5Viewer));

  PetscCall( VecDestroy( &solGlob) );
  PetscCall( VecDestroy( &solLoc) );
  PetscCall( DMDestroy( &dmMesh) );
  PetscCall( PetscFinalize() );
}