PROGRAM ParAssembly
!
!      Mesh : Q1_4cells.msh
!
!
!      7----------10---------6
!      |          |          |
!      |          |          |
!      |     3    |    2     |
!      |          |          |
!      |          |          |
!     11----------12---------9
!      |          |          |
!      |          |          |
!      |     0    |    1     |
!      |          |          |
!      |          |          |
!      4----------8----------5
!
!  Unknowns in u_global = [u9,u11,u12]
!
implicit none
#include <petsc/finclude/petsc.h90>

  PetscInt, parameter :: ndime = 2, ndof = 1, nodes = 4

  PetscErrorCode      :: ierr
  PetscMPIInt         :: nranks,rank
  DM                  :: dm,dm_distrib
  Mat                 :: A
  !Vec                 :: u_global
  PetscBool           :: interpolate
  PetscSection        :: section

  PetscInt, pointer :: cone(:),closure(:),indices(:)
  PetscInt, target, dimension(1) :: numComp ! We define a scalar field
  PetscInt, pointer :: pNumComp(:)
  PetscInt, target, dimension(3) :: numDof  ! number of dof per point
  PetscInt, pointer :: pNumDof(:)
  PetscInt, target, dimension(1) :: bcField
  PetscInt, pointer :: pBcField(:)
  IS, target, dimension(1) :: bcCompIS
  IS, pointer :: pBcCompIS(:)
  IS, target, dimension(1) :: bcPointIS
  IS, pointer :: pBcPointIS(:)

  PetscInt , dimension(nodes*ndof)   :: local_ldof
  PetscReal, dimension(  ndof*nodes,  ndof*nodes) :: elstf

  PetscInt       :: nfields, numBC, dofs, dofsfix
  PetscInt       :: cStart, cEnd, vStart, eStart, eEnd, vEnd, pStart, pEnd
  PetscInt       :: idface
  PetscInt       :: size_cone,size_closure
  PetscInt       :: i,ic,ie,icell,idof,inod,ipoin,indx

  ISLocalToGlobalMapping :: ltog

  !*************************************
  !--- Beginning of the program
  !*************************************
  call PetscInitialize(PETSC_NULL_CHARACTER,ierr)

  call MPI_Comm_size(PETSC_COMM_WORLD,nranks,ierr)
  call MPI_Comm_rank(PETSC_COMM_WORLD,rank,ierr)


  !*************************************
  !--- Creating Mesh
  !*************************************

  interpolate = PETSC_TRUE
  call DMPlexCreateGmshFromFile(PETSC_COMM_WORLD,"Q1_4cells.msh",interpolate,dm,ierr)

  !--- Distributing the mesh over processes
  if (nranks > 1) then
     call DMPlexDistribute(dm,0,PETSC_NULL_OBJECT,dm_distrib,ierr);
     call DMDestroy(dm,ierr);
     dm = dm_distrib;
  endif

  !*************************************
  !--- Set Petsc Layout
  !*************************************

 !--- We create additional labels for vertices on labeled faces
  call DMCreateLabel(dm,"AllVertBC",ierr)

  !--- We get the range Id of Sieve Points which are faces/edges in 2D
  call DMPlexGetDepthStratum(dm,ndime-1,eStart,eEnd,ierr)
  do ie = eStart, eEnd-1
      !- This returns FaceNum if ifac is external face or -1 if not
      !- IMPORTANT : Physical Entities must be defined in .geo to have FaceNums
      call DMGetLabelValue(dm,"Face Sets",ie,idface,ierr)
      if (idface .gt. 0) then
         !-- For Dirichlet BCs on vertices we want to extract vertices
         call DMPlexGetConeSize(dm,ie,size_cone,ierr)
         call DMPlexGetCone(dm,ie,cone,ierr)
         do i=1,size_cone
            !-- Select only 2 faces for Bcs
            if (idface .eq. 1 .or. idface .eq. 3) then
               call DMSetLabelValue(dm,"AllVertBC",cone(i),1,ierr)
            endif
         enddo        
      endif
  enddo

  call DMView(dm,PETSC_VIEWER_STDOUT_WORLD, ierr)


  !---- Defining a Scalar Field on vertices
  nFields = 1    ! Num of fields
  numComp(1) = 1 ! Num of components
  pnumComp => numComp
  do i = 1, nFields*(ndime+1)
     numDof(i) = 0
  end do
  numDof(0*(ndime+1)+1) = 1  !1 DOFs on vertices
  pnumDof => numDof

  !--- Defining Boundary Conditions on faces 1 and 3
  numBC = 1
  bcField(1) = 0  ! Field 0 : u
  !--- Data structure defining contrained components of u
  call ISCreateStride(PETSC_COMM_WORLD, 1, 0, 1, bcCompIS(1), ierr);
  !--- Data structure with num of SievePoints with val=1 in Label 'Vert Sets'
  call DMGetStratumIS(dm, "AllVertBC", 1, bcPointIS(1),ierr);
  pBcField => bcField
  pBcCompIS => bcCompIS
  pBcPointIS => bcPointIS

  !--- Create a PetscSection with this data layout
  call DMPlexCreateSection(dm, ndime, nFields, pNumComp,        &
           pNumDof, numBC, pBcField, pBcCompIS, pBcPointIS,     &
           PETSC_NULL_OBJECT, section, ierr)

  !--- Name the Field variables
  call PetscSectionSetFieldName(section, 0, 'u', ierr)

  !--- Tell DM to use this section
  call DMSetDefaultSection(dm,section,ierr)

  !--- Visualizing Section
  call PetscPrintf(PETSC_COMM_WORLD,"\n Visualizing Section ... \n",ierr)
  call PetscSectionView(section, PETSC_VIEWER_STDOUT_WORLD, ierr)


  !*************************************
  !--- Matrix and RHS Assembly
  !*************************************

  !--- Create a Sparse Matrix with sparsity struct. based on dm
  call DMCreateMatrix(dm,A,ierr)
  !call DMCreateGlobalVector(dm,u_global,ierr)
  !--- Set Mapping Local to Global Indexes
  call DMGetLocalToGlobalMapping(dm,ltog,ierr)
  call MatSetLocalToGlobalMapping(A,ltog,ltog,ierr);
  call ISLocalToGlobalMappingDestroy(ltog,ierr);

  call MatZeroEntries(A,ierr)


  !--- Starting assembly
  call DMPlexGetDepthStratum(dm,0,vStart,vEnd,ierr);CHKERRQ(ierr);  ! Start & End vertices
  call DMPlexGetHeightStratum(dm,0,cStart,cEnd,ierr);CHKERRQ(ierr); ! Start & End cells


  !--- Loop over cells
  call PetscPrintf(PETSC_COMM_WORLD,"\n Loop over cells \n\n",ierr)

  do icell = cStart, cEnd-1

     call DMPlexGetTransitiveClosure(dm,icell,PETSC_TRUE,closure,ierr)
     size_closure = size(closure)/2  !closure contains (points,orientation)

     inod = 0
     do ic = 1 , size_closure
        ipoin = closure((ic-1)*2+1)
        call PetscSectionGetDof(section,ipoin,dofs,ierr)
        if (dofs .ne. 0) then  ! To get Vertices
           inod = inod + 1
           !--- Get number of vertex constrained dofs
           call PetscSectionGetConstraintDof(section,ipoin,dofsfix,ierr)
           if (dofsfix /= 0) then
               call PetscSectionGetConstraintIndicesF90(section,ipoin,indices,ierr)
           endif
           !--- Get location of point in local vector
           call DMPlexGetPointLocal(dm,ipoin,pStart,pEnd,ierr)
           !--- Get location of DOF(s) number(s)
           do idof = 1 , dofs
              i = (inod-1)*dofs + idof
              local_ldof(i) = pStart + idof  !Location of DOF in local_locdof()
           enddo
           if (dofsfix /= 0) then ! If constrained dofs
              do indx = 1 , dofsfix
                 idof = indices(indx) + 1
                 i = (inod-1)*dofs + idof
                 ! Change local_dof(i) to -local_dof(i) if constrained
                 ! Will be used to modify RHS for BCs
                 local_ldof(i) = -local_ldof(i)
              enddo
           endif
           if (dofsfix /= 0) then
              call PetscSectionRestoreConstraintIndicesF90(section,ipoin,indices,ierr)
           endif
        endif
     enddo
     call DMPlexRestoreTransitiveClosure(dm,icell,PETSC_TRUE,closure,ierr)


     elstf = 1.0 ! Elementary Stiffness


     do i = 1 , ndof*nodes
        if (local_ldof(i) > 0) then !Set only if DOF not constrained ??
           call MatSetValuesLocal(A,1,local_ldof(i)-1,ndof*nodes,local_ldof(:)-1, &
                                  elstf(i,:),ADD_VALUES,ierr)
        endif
     enddo

  enddo

  call MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY,ierr)
  call MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY,ierr)

  call MatView(A,PETSC_VIEWER_STDOUT_WORLD,ierr)


  !*************************************
  !--- Terminating Program
  !*************************************


  !--- Destroying objects
  call PetscPrintf(PETSC_COMM_WORLD,"\n ... Destroying objects \n",ierr)
  call DMDestroy(dm,ierr)
  call ISDestroy(bcCompIS, ierr)
  call ISDestroy(bcPointIS, ierr)
  !call MatDestroy(A,ierr)
  !call VecDestroy(u_global,ierr)

  call PetscPrintf(PETSC_COMM_WORLD,"\n ---- End Program ---- \n\n",ierr)
  call PetscFinalize(ierr)

  return

ENDPROGRAM