program test

#include "petsc/finclude/petscvec.h"
#include "petsc/finclude/petscdmplex.h"
#include "petsc/finclude/petscdmlabel.h"
 use petscvec
 use petscdmplex
 use petscsys

 implicit none

 DM                              :: dm_dumm, dm_geom, dm_geom_cell
 DM                              :: dm_dumm_g
 Vec                             :: coord
 Vec                             :: vglo, vloc, vscale
 Vec                             :: vcell, vcellg
 Vec                             :: vnode, vnodeg
 Mat                             :: Map
 PetscMPIInt                     :: myid, nproc
 PetscErrorCode                  :: ierr
 PetscSection                    :: section
 PetscInt                        :: dim, size, row, col
 PetscFE                         :: fe_0, fe_1
 PetscSF                         :: sf, sf_cell
 PetscReal, target, dimension(1) :: xx, yy, zz
 PetscReal, pointer              :: pxx(:) !=> xx
 PetscReal, pointer              :: pyy(:) !=> yy
 PetscReal, pointer              :: pzz(:) !=> zz
 PetscReal, target, dimension(2) :: cent_2D, norm_2D ! number 2 means x and y coordinates
 PetscReal, target, dimension(1) :: sol_tmp
 PetscReal, pointer              :: pSol(:) !=> sol_tmp
 PetscReal, pointer              :: pCent_2D(:)      !=> cent_2D
 PetscReal, pointer              :: pNorm_2D(:)      !=> norm_2D
 DMLabel, pointer                :: nolabel(:) => NULL()
 PetscInt                        :: numFields = 3 
 PetscInt                        :: numBC     = 1
 PetscInt, target, dimension(3)  :: numComp 
 PetscInt, target, dimension(12) :: numDof  
 PetscInt, pointer               :: pNumComp(:)
 PetscInt, pointer               :: pNumDof(:)
 PetscInt, target, dimension(1)  :: bcField
 PetscInt, pointer               :: pBcField(:)
 IS, target, dimension(1)        :: bcCompIS
 IS, target, dimension(1)        :: bcPointIS
 IS, pointer                     :: pBcCompIS(:)
 IS, pointer                     :: pBcPointIS(:)
 PetscViewer                     :: viewer
 integer                         :: i
 integer                         :: pst, pend
 integer                         :: vst, vend
 integer                         :: cst, cend
 integer                         :: nnode, ncell
 integer                         :: inode, ng
 double precision                :: area
 double precision, allocatable   :: sol(:), sol_cell(:) ! temporary solution vectors
 double precision, allocatable   :: xnd(:), ynd(:), znd(:)
 double precision, allocatable   :: xc(:), yc(:), zc(:)


 call PetscInitialize(PETSC_NULL_CHARACTER, ierr)
 if(ierr.ne.0) then
   write(*,*) 'Fail to initialize PETSc ...'
   stop
 endif
 call MPI_COMM_RANK(MPI_COMM_WORLD, myid, ierr)
 call MPI_COMM_SIZE(MPI_COMM_WORLD, nproc, ierr)
 write(*,*) 'MPI information ... ', myid, nproc

 call DMCreate(PETSC_COMM_WORLD, dm_geom, ierr);CHKERRQ(ierr)
 call DMSetType(dm_geom, DMPLEX, ierr);CHKERRQ(ierr)
 call DMSetFromOptions(dm_geom, ierr);CHKERRQ(ierr)
 call DMViewFromOptions(dm_geom, PETSC_NULL_DM, '-dm_view', ierr);CHKERRQ(ierr)

 call DMGetDimension(dm_geom, dim, ierr);CHKERRA(ierr)
 call DMClone(dm_geom, dm_geom_cell, ierr);CHKERRA(ierr)
 ! dm_geom      : solution on vertex
 ! dm_geom_cell : solution on cell
 ! mapping between them needed

 ! setup sections
 ! (1) define section with 1-dof on cell for "cell-dm"
 call PetscSectionCreate(PETSC_COMM_WORLD, section, ierr);CHKERRA(ierr)
 call PetscSectionSetNumFields(section, 1, ierr);CHKERRA(ierr)
 call PetscSectionSetFieldComponents(section, 0, 1, ierr);CHKERRA(ierr)
 call DMPlexGetHeightStratum(dm_geom_cell, 0, cst, cend, ierr);CHKERRA(ierr) ! cells
 call PetscSectionSetChart(section, cst, cend, ierr);CHKERRA(ierr)
 do i=cst,cend-1
   call PetscSectionSetDof(section, i, 1, ierr);CHKERRA(ierr)
   call PetscSectionSetFieldDof(section, i, 0, 1, ierr);CHKERRA(ierr)
 enddo
 call PetscSectionSetUp(section, ierr);CHKERRA(ierr)
 call DMSetLocalSection(dm_geom_cell, section, ierr);CHKERRA(ierr)
 call PetscSectionDestroy(section, ierr);CHKERRA(ierr)
 ! setup star forest using that section
 call DMGetSectionSF(dm_geom_cell, sf_cell, ierr);CHKERRA(ierr)
 call PetscSFSetUp(sf_cell, ierr);CHKERRA(ierr)

 ! (2) define section with 1-dof on vertex for "vertex-dm"
 call PetscSectionCreate(PETSC_COMM_WORLD, section, ierr);CHKERRA(ierr)
 call PetscSectionSetNumFields(section, 1, ierr);CHKERRA(ierr)
 call PetscSectionSetFieldComponents(section, 0, 1, ierr);CHKERRA(ierr)
 call DMPlexGetDepthStratum(dm_geom, 0, vst, vend, ierr);CHKERRA(ierr) ! vertex
 call PetscSectionSetChart(section, vst, vend, ierr);CHKERRA(ierr)
 do i=vst,vend-1
   call PetscSectionSetDof(section, i, 1, ierr);CHKERRA(ierr)
   call PetscSectionSetFieldDof(section, i, 0, 1, ierr);CHKERRA(ierr)
 enddo
 call PetscSectionSetUp(section, ierr);CHKERRA(ierr)
 call DMSetLocalSection(dm_geom, section, ierr);CHKERRA(ierr)
 call PetscSectionDestroy(section, ierr);CHKERRA(ierr)
 ! setup star forest using that section
 call DMGetSectionSF(dm_geom, sf, ierr);CHKERRA(ierr)
 call PetscSFSetUp(sf, ierr);CHKERRA(ierr)

 ! FE object for P1(or Q1) for dm_geom (i.e., vertex dm)
 call PetscFECreateLagrange(PETSC_COMM_WORLD, dim, 1, PETSC_FALSE, 1, PETSC_DETERMINE, fe_1, ierr);CHKERRA(ierr)
 call DMSetField(dm_geom, 0, PETSC_NULL_DMLABEL, fe_1, ierr);CHKERRA(ierr)
 call DMCreateDS(dm_geom, ierr);CHKERRA(ierr)
 call PetscFEDestroy(fe_1, ierr);CHKERRA(ierr)

 ! FE object for P0(or Q0) for dm_geom_cell (i.e., cell dm)
 call PetscFECreateLagrange(PETSC_COMM_WORLD, dim, 1, PETSC_FALSE, 0, PETSC_DETERMINE, fe_0, ierr);CHKERRA(ierr)
 call DMSetField(dm_geom_cell, 0, PETSC_NULL_DMLABEL, fe_0, ierr);CHKERRA(ierr)
 call DMCreateDS(dm_geom_cell, ierr);CHKERRA(ierr)
 call PetscFEDestroy(fe_0, ierr);CHKERRA(ierr)

 ! create interpolation 
 call DMCreateInterpolation(dm_geom, dm_geom_cell, Map, vscale, ierr);CHKERRA(ierr)
 call MatViewFromOptions(Map, PETSC_NULL_MAT, "-map_view", ierr);CHKERRA(ierr)
 call VecViewFromOptions(vscale, PETSC_NULL_VEC, "-map_view", ierr);CHKERRA(ierr)

 ! create local vectors
 call DMCreateLocalVector(dm_geom_cell, vcell, ierr);CHKERRA(ierr)
 call DMCreateLocalVector(dm_geom, vnode, ierr);CHKERRA(ierr)

 ! create global vectors
 call DMCreateGlobalVector(dm_geom_cell, vcellg, ierr);CHKERRA(ierr)
 call DMCreateGlobalVector(dm_geom, vnodeg, ierr);CHKERRA(ierr)

 ! get coordinates over vertex and cell-centroids
 call DMPlexGetDepthStratum(dm_geom, 0, vst, vend, ierr);CHKERRA(ierr)
 call DMPlexGetHeightStratum(dm_geom, 0, cst, cend, ierr);CHKERRA(ierr)
 nnode = vend - vst ! number of vertex each proc has
 ncell = cend - cst ! number of cell each proc has

 ! define vertex and cell-centroid coordinates
 allocate(xnd(nnode), ynd(nnode), znd(nnode))
 allocate(xc(ncell), yc(ncell), zc(ncell))
 xnd = 0.d0 ; ynd = 0.d0 ; znd = 0.d0
 xc  = 0.d0 ; yc  = 0.d0 ; zc  = 0.d0

 ! define coordinate of vertex
 call DMGetCoordinatesLocal(dm_geom, coord, ierr);CHKERRA(ierr)
 call VecGetSize(coord, size, ierr);CHKERRA(ierr)
 inode=1                 
 pxx => xx  ; pyy => yy  ; pzz => zz 
 do i=0,size-1 
   ! initialize
   xx  =  0.0 ; yy  =  0.0 ; zz = 0.d0
   if(    mod(i,dim).eq.0) then ! x coordinate
     call VecGetValues(coord, 1, i, pxx, ierr);CHKERRA(ierr)
     xnd(inode) = pxx(1) 
   elseif(mod(i,dim).eq.1) then ! y coordinate
     call VecGetValues(coord, 1, i, pyy, ierr);CHKERRA(ierr)
     ynd(inode) = pyy(1)
   endif
   ! update node index
   if(mod(i,dim).eq.dim-1) then
     if(i.ne.size-1) then
       inode = inode + 1
     endif
   endif
 enddo
 ! check local node counting
 if(inode.ne.nnode) then
   write(*,*) 'node counting error in xnd ynd ...'
   stop
 endif

 ! define coordinate of cell centroid
 cent_2D = 0.0 ; norm_2D = 0.0
 pCent_2D => cent_2D 
 pNorm_2D => norm_2D ! Note : two different pointers should be pointing two different target arrays
 do i=cst,cend-1 ! loop over cell
   call DMPlexComputeCellGeometryFVM(dm_geom, i, area, pCent_2D, pNorm_2D, ierr);CHKERRA(ierr)
   xc(i-cst+1) = pCent_2D(1)
   yc(i-cst+1) = pCent_2D(2)
 enddo

 ! fill the solution vector with cell-center's x-coordinate value for test!
 ! this will be mapped into vertex field and compared with xnd
 allocate(sol_cell(ncell))
 do i=1,ncell
   sol_cell(i) = xc(i)
 enddo

 ! fill local cell vector
 pSol => sol_tmp
 do i=0,ncell-1
   pSol = sol_cell(i+1)
   call VecSetValues(vcell, 1, i, pSol, INSERT_VALUES, ierr);CHKERRA(ierr)
 enddo

 ! fill global cell vector
 call DMLocalToGlobalBegin(dm_geom_cell, vcell, INSERT_VALUES, vcellg, ierr);CHKERRA(ierr)
 call DMLocalToGlobalEnd(dm_geom_cell, vcell, INSERT_VALUES, vcellg, ierr);CHKERRA(ierr)

 ! mapping global cell field to global vertex field
 call VecZeroEntries(vnodeg, ierr);CHKERRA(ierr)
 call MatRestrict(Map, vcellg, vnodeg, ierr);CHKERRA(ierr)
 call VecPointwiseMult(vnodeg, vnodeg, vscale, ierr);CHKERRQ(ierr)
 call VecViewFromOptions(vcellg, PETSC_NULL_VEC, "-cell_view", ierr);CHKERRA(ierr)
 call VecViewFromOptions(vnodeg, PETSC_NULL_VEC, "-node_view", ierr);CHKERRA(ierr)

 ! broadcast to local vertex field
 call DMGlobalToLocalBegin(dm_geom, vnodeg, INSERT_VALUES, vnode, ierr);CHKERRA(ierr)
 call DMGlobalToLocalEnd(dm_geom, vnodeg, INSERT_VALUES, vnode, ierr);CHKERRA(ierr)
 call VecViewFromOptions(vcell, PETSC_NULL_VEC, "-cell_view", ierr);CHKERRA(ierr)
 call VecViewFromOptions(vnode, PETSC_NULL_VEC, "-node_view", ierr);CHKERRA(ierr)

 ! create viewer
 call PetscViewerCreate(PETSC_COMM_WORLD, viewer, ierr);CHKERRA(ierr)
 
 ! set viewer type
 call PetscViewerSetType(viewer, PETSCVIEWERVTK, ierr);CHKERRA(ierr)
 
 ! set output file name for viewer
 ! fname = 'VertexField_Map_from_Cell.vtu'
 call PetscViewerFileSetName(viewer, 'VertexField_Map_from_Cell.vtu', ierr);CHKERRA(ierr) 
 
 ! print dm object to the file
 call VecView(vnode, viewer, ierr);CHKERRA(ierr)
 
 ! destroy and free memory
 call PetscViewerDestroy(viewer, ierr);CHKERRA(ierr)
 call MatDestroy(Map, ierr);CHKERRA(ierr)
 call VecDestroy(vscale, ierr);CHKERRA(ierr)
 call VecDestroy(vcell, ierr);CHKERRA(ierr)
 call VecDestroy(vnode, ierr);CHKERRA(ierr)
 call VecDestroy(vcellg, ierr);CHKERRA(ierr)
 call VecDestroy(vnodeg, ierr);CHKERRA(ierr)
 call DMDestroy(dm_geom, ierr);CHKERRA(ierr)
 call DMDestroy(dm_geom_cell, ierr);CHKERRA(ierr)
 deallocate(xnd, ynd, znd, xc, yc, zc)
 call PetscFinalize(ierr)

 stop
end program test