! //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// !
!
! Maintainers : support@fluidnumerics.com
! Official Repository : https://github.com/FluidNumerics/self/
!
! Copyright © 2024 Fluid Numerics LLC
!
! Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
!
! 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
!
! 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in
!    the documentation and/or other materials provided with the distribution.
!
! 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from
!    this software without specific prior written permission.
!
! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
! LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
! HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
! LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
! THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
! THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
!
! //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// !

module SELF_ElementSetMesh_3D
!! Partition a 3-D mesh into element sets and extract one sub-mesh per set.
!!
!! An element set is the collection of all global elements that share a set id
!! s(e) in {1, ..., nSets}, for example one set per material type. Each set is run by
!! its own DG model on its own Mesh3D, and a parent model (MultiSetModel3D) couples the
!! sets where they share faces. This module builds those sub-meshes.
!!
!! Numbering. Global elements are numbered e = 1..nGlobalElem and rank r (0-based) owns
!! the contiguous range offsetElem(r+1)+1 .. offsetElem(r+2) of the global mesh. Within
!! set k, the set-global id of element e with s(e) = k is
!!
!!   g_k(e) = #{ e' <= e : s(e') = k },
!!
!! i.e. set-global ids follow global element order. Each rank's elements of set k
!! therefore form one contiguous block of set-global ids, and every sub-mesh keeps the
!! global mesh's element ownership: no element changes rank.
!!
!! Faces. A face between two elements of the same set stays an interior face of that
!! set's sub-mesh, with its neighbour renumbered to g_k. A face between elements of
!! different sets becomes an interface face of both sub-meshes:
!!
!!   sideInfo(1) = 0, sideInfo(3) = 0, sideInfo(4) = 0,
!!   sideInfo(5) = SELF_BC_SET_INTERFACE,
!!
!! with sideInfo(2) keeping the global side id. An interface face is never a physical
!! boundary; the parent model fills its exterior state. Domain boundary faces are copied
!! unchanged. Every sub-mesh keeps the global nUniqueSides and global side ids, because
!! the CPU halo exchange derives its message tags from them.
!!
!! Mortars. A 2:1 or anisotropic mortar must lie entirely inside one set
!! (ValidateSetPartition). Mortars inside a set are copied with their element ids
!! renumbered to g_k, and the mortar index in sideInfo(1) is renumbered to the
!! sub-mesh's mortar tables.
!!
!! Collectives. GatherSetOfElem is a collective over the global mesh's communicator, and
!! so is ExtractSubMesh (through the decomposition it builds). Every rank must call
!! them for every set, including sets in which it owns no elements. The other routines
!! are rank-local and return the same result on every rank.

  use SELF_Constants
  use SELF_Mesh_3D
  use mpi

  implicit none

  !! Boundary-condition id that marks an interface face between two element sets in a
  !! sub-mesh. It is not a physical boundary condition: the parent model writes the
  !! exterior state on these faces. The value is outside the range of HOHQMesh boundary
  !! indices and of the ids defined in SELF_Mesh.
  integer,parameter :: SELF_BC_SET_INTERFACE = 9100

contains

  subroutine GatherSetOfElem(globalMesh,setOfElemLocal,setOfElem)
    !! Gather the set id of every global element onto every rank.
    !!
    !! On return setOfElem(e) = s(e) for e = 1..nGlobalElem, in global element order,
    !! where rank r contributed setOfElemLocal(1:nElem) for its own global range. This
    !! is one MPI_Allgatherv over globalMesh%decomp%mpiComm, run at setup only.
    !!
    !!  Input
    !!    - globalMesh : the decomposed global mesh
    !!    - setOfElemLocal(1:globalMesh%nElem) : set id of each rank-local element
    !!  Output
    !!    - setOfElem(1:nGlobalElem) : set id of each global element, where
    !!      nGlobalElem = globalMesh%decomp%offsetElem(nRanks+1)
    implicit none
    type(Mesh3D),intent(in) :: globalMesh
    integer,intent(in) :: setOfElemLocal(1:globalMesh%nElem)
    integer,intent(out) :: setOfElem(1:)
    ! Local
    integer :: r,nRanks,ierror
    integer,allocatable :: counts(:),displs(:)

    nRanks = globalMesh%decomp%nRanks
    if(size(setOfElem) /= globalMesh%decomp%offsetElem(nRanks+1)) then
      error stop "GatherSetOfElem : setOfElem must have one entry per global element"
    endif

    if(nRanks == 1) then
      setOfElem(1:globalMesh%nElem) = setOfElemLocal(1:globalMesh%nElem)
      return
    endif

    allocate(counts(1:nRanks),displs(1:nRanks))
    do r = 1,nRanks
      counts(r) = globalMesh%decomp%offsetElem(r+1)-globalMesh%decomp%offsetElem(r)
      displs(r) = globalMesh%decomp%offsetElem(r)
    enddo
    call mpi_allgatherv(setOfElemLocal,counts(globalMesh%decomp%rankId+1),MPI_INTEGER, &
                        setOfElem,counts,displs,MPI_INTEGER, &
                        globalMesh%decomp%mpiComm,ierror)
    deallocate(counts,displs)

  endsubroutine GatherSetOfElem

  function ValidateSetPartition(globalMesh,setOfElem,nSets) result(ok)
    !! Check that a global set map can be split into sub-meshes.
    !!
    !! The partition is valid when
    !!   1. 1 <= s(e) <= nSets for every global element e;
    !!   2. every set k = 1..nSets holds at least one element;
    !!   3. every 2:1 mortar (big element and its four small elements) and every
    !!      anisotropic mortar (big element and its two small elements) lies in a single
    !!      set, because the interface exchange between sets handles conforming faces
    !!      only.
    !! The mortar tables hold global element ids and are replicated, and setOfElem is
    !! global, so every rank returns the same answer without communication. The first
    !! violation found is printed on rank 0.
    !!
    !!  Input
    !!    - globalMesh : the decomposed global mesh
    !!    - setOfElem(1:nGlobalElem) : set id of each global element (GatherSetOfElem)
    !!    - nSets : number of sets, >= 1
    !!  Output
    !!    - ok : .true. when all three conditions hold
    implicit none
    type(Mesh3D),intent(in) :: globalMesh
    integer,intent(in) :: setOfElem(1:)
    integer,intent(in) :: nSets
    logical :: ok
    ! Local
    integer :: e,k,m,q,s0
    integer,allocatable :: nInSet(:)
    logical :: report

    ok = .true.
    report = (globalMesh%decomp%rankId == 0)

    if(nSets < 1) then
      ok = .false.
      if(report) print*,__FILE__," : Error : nSets must be at least 1, got ",nSets
      return
    endif

    allocate(nInSet(1:nSets))
    nInSet = 0
    do e = 1,size(setOfElem)
      if(setOfElem(e) < 1 .or. setOfElem(e) > nSets) then
        ok = .false.
        if(report) print*,__FILE__," : Error : element ",e," has set id ",setOfElem(e), &
          " outside 1..",nSets
        deallocate(nInSet)
        return
      endif
      nInSet(setOfElem(e)) = nInSet(setOfElem(e))+1
    enddo
    do k = 1,nSets
      if(nInSet(k) == 0) then
        ok = .false.
        if(report) print*,__FILE__," : Error : element set ",k," is empty"
        deallocate(nInSet)
        return
      endif
    enddo
    deallocate(nInSet)

    ! Mortar element rows: 1 (big) and 3,5,7,9 (small) in mortarInfo; 1 (big) and 3,5
    ! (small) in anisoMortarInfo.
    do m = 1,globalMesh%nMortars
      s0 = setOfElem(globalMesh%mortarInfo(1,m))
      do q = 1,4
        if(setOfElem(globalMesh%mortarInfo(1+2*q,m)) /= s0) then
          ok = .false.
          if(report) print*,__FILE__," : Error : mortar ",m," crosses element sets"
          return
        endif
      enddo
    enddo
    do m = 1,globalMesh%nAnisoMortars
      s0 = setOfElem(globalMesh%anisoMortarInfo(1,m))
      do q = 1,2
        if(setOfElem(globalMesh%anisoMortarInfo(1+2*q,m)) /= s0) then
          ok = .false.
          if(report) print*,__FILE__," : Error : anisotropic mortar ",m, &
            " crosses element sets"
          return
        endif
      enddo
    enddo

  endfunction ValidateSetPartition

  subroutine BuildSetNumbering(globalMesh,setOfElem,nSets,setGlobalId,setOffsets)
    !! Number the elements of every set and partition each set over the ranks.
    !!
    !! For each global element e with k = s(e), setGlobalId(e) = g_k(e), its id within
    !! set k (see the module documentation). For each set k, setOffsets(1:nRanks+1,k)
    !! are the decomposition offsets of the set's sub-mesh: rank r owns set-global ids
    !! setOffsets(r+1,k)+1 .. setOffsets(r+2,k), which are exactly the set-k elements in
    !! rank r's range of the global mesh. A rank that owns no element of set k has
    !! setOffsets(r+1,k) = setOffsets(r+2,k). Rank-local; every rank gets the same result.
    !!
    !!  Input
    !!    - globalMesh : the decomposed global mesh
    !!    - setOfElem(1:nGlobalElem) : set id of each global element, validated
    !!    - nSets : number of sets
    !!  Output
    !!    - setGlobalId(1:nGlobalElem) : set-global id of each global element
    !!    - setOffsets(1:nRanks+1,1:nSets) : per-set decomposition offsets
    implicit none
    type(Mesh3D),intent(in) :: globalMesh
    integer,intent(in) :: setOfElem(1:)
    integer,intent(in) :: nSets
    integer,intent(out) :: setGlobalId(1:)
    integer,intent(out) :: setOffsets(1:,1:)
    ! Local
    integer :: e,k,r,nRanks
    integer,allocatable :: nInSet(:)

    nRanks = globalMesh%decomp%nRanks
    allocate(nInSet(1:nSets))
    nInSet = 0
    setOffsets(1:nRanks+1,1:nSets) = 0

    do r = 1,nRanks
      do e = globalMesh%decomp%offsetElem(r)+1,globalMesh%decomp%offsetElem(r+1)
        k = setOfElem(e)
        nInSet(k) = nInSet(k)+1
        setGlobalId(e) = nInSet(k)
      enddo
      setOffsets(r+1,1:nSets) = nInSet(1:nSets)
    enddo

    deallocate(nInSet)

  endsubroutine BuildSetNumbering

  subroutine ExtractSubMesh(globalMesh,iset,setOfElem,setGlobalId,setOffsets,comm,subMesh)
    !! Build the sub-mesh of element set iset from this rank's slice of the global mesh.
    !!
    !! The sub-mesh holds this rank's elements of set iset, in global element order, with
    !! the global mesh's geometry, materials and boundary-condition table. Its
    !! decomposition is built from setOffsets(:,iset), so element ownership is unchanged;
    !! a rank may own zero elements. Faces are renumbered as described in the module
    !! documentation, and mortars inside the set are copied with renumbered element ids.
    !!
    !! The decomposition is initialised on comm, which must be a duplicate of
    !! globalMesh%decomp%mpiComm made for this sub-mesh alone (MPI_Comm_dup). The
    !! decomposition does not own it: the caller frees comm with MPI_Comm_free after
    !! subMesh%Free. A separate communicator per sub-mesh keeps the halo messages of
    !! different sets from matching each other.
    !!
    !! Collective over comm: every rank calls it for every set.
    !!
    !!  Input
    !!    - globalMesh : the decomposed global mesh
    !!    - iset : the set to extract, 1 <= iset <= size(setOffsets,2)
    !!    - setOfElem, setGlobalId, setOffsets : from GatherSetOfElem and
    !!      BuildSetNumbering, with the partition validated by ValidateSetPartition
    !!    - comm : this sub-mesh's duplicate communicator
    !!  Output
    !!    - subMesh : the sub-mesh, device copy up to date
    implicit none
    type(Mesh3D),intent(in) :: globalMesh
    integer,intent(in) :: iset
    integer,intent(in) :: setOfElem(1:)
    integer,intent(in) :: setGlobalId(1:)
    integer,intent(in) :: setOffsets(1:,1:)
    integer,intent(in) :: comm
    type(Mesh3D),intent(out) :: subMesh
    ! Local
    integer :: nGeo,nBCs,nLocal,rank,goff,el,eg,e2,s,n,k,m,q,a,nm,na
    integer,allocatable :: subLocal(:)
    integer,allocatable :: mortarNew(:),anisoNew(:)

    nGeo = globalMesh%nGeo
    nBCs = globalMesh%nBCs

    ! The decomposition pointers carry no null default, and GenerateDecompositionFromOffsets
    ! tests their association; subMesh may live in uninitialised (heap) memory.
    nullify(subMesh%decomp%elemToRank,subMesh%decomp%offSetElem)
    call subMesh%decomp%Init(comm=comm)
    ! Start the request scratch at one entry: the global mesh's size would be copied into
    ! every set on every rank. MultiSetModel3D%SetupInterfaces reserves each set's actual
    ! requirement once its variable count is known, and every exchange reserves before
    ! posting in any case.
    call subMesh%decomp%GenerateDecompositionFromOffsets(setOffsets(:,iset),1)

    rank = subMesh%decomp%rankId
    nLocal = setOffsets(rank+2,iset)-setOffsets(rank+1,iset)

    call subMesh%Init(nGeo,nLocal,6*nLocal,8*nLocal,nBCs)
    subMesh%nGlobalElem = setOffsets(subMesh%decomp%nRanks+1,iset)
    subMesh%nUniqueSides = globalMesh%nUniqueSides ! global side count: the CPU halo tag stride
    subMesh%nUniqueNodes = globalMesh%nUniqueNodes
    subMesh%nCornerNodes = globalMesh%nCornerNodes
    subMesh%quadrature = globalMesh%quadrature
    subMesh%nLayers = 0 ! a set need not contain whole columns of an extruded mesh

    ! Boundary-condition and material tables (replicated on every rank).
    if(nBCs > 0) then
      subMesh%BCType(1:4,1:nBCs) = globalMesh%BCType(1:4,1:nBCs)
      do k = 1,nBCs
        subMesh%BCNames(k) = globalMesh%BCNames(k)
      enddo
    endif
    subMesh%nMaterials = globalMesh%nMaterials
    if(allocated(subMesh%materialNames)) deallocate(subMesh%materialNames)
    allocate(subMesh%materialNames(1:globalMesh%nMaterials))
    subMesh%materialNames(1:globalMesh%nMaterials) = globalMesh%materialNames(1:globalMesh%nMaterials)

    ! Per-element data. subLocal(el) is the sub-mesh local id of global-mesh local
    ! element el, or 0 when el belongs to another set.
    goff = globalMesh%decomp%offsetElem(rank+1)
    allocate(subLocal(1:max(globalMesh%nElem,1)))
    subLocal = 0
    n = 0
    do el = 1,globalMesh%nElem
      eg = goff+el
      if(setOfElem(eg) /= iset) cycle
      n = n+1
      subLocal(el) = n
      subMesh%nodeCoords(1:3,1:nGeo+1,1:nGeo+1,1:nGeo+1,n) = &
        globalMesh%nodeCoords(1:3,1:nGeo+1,1:nGeo+1,1:nGeo+1,el)
      subMesh%globalNodeIDs(1:nGeo+1,1:nGeo+1,1:nGeo+1,n) = &
        globalMesh%globalNodeIDs(1:nGeo+1,1:nGeo+1,1:nGeo+1,el)
      subMesh%elemInfo(1:6,n) = globalMesh%elemInfo(1:6,el)
      subMesh%elemMaterial(n) = globalMesh%elemMaterial(el)

      do s = 1,6
        e2 = globalMesh%sideInfo(3,s,el)
        subMesh%sideInfo(1:5,s,n) = globalMesh%sideInfo(1:5,s,el)
        if(e2 > 0) then
          if(setOfElem(e2) == iset) then
            ! Same-set neighbour: renumber; side, flip and global side id are unchanged.
            subMesh%sideInfo(3,s,n) = setGlobalId(e2)
          else
            ! Neighbour in another set: an interface face of this sub-mesh.
            subMesh%sideInfo(1,s,n) = 0
            subMesh%sideInfo(3,s,n) = 0
            subMesh%sideInfo(4,s,n) = 0
            subMesh%sideInfo(5,s,n) = SELF_BC_SET_INTERFACE
          endif
        endif
      enddo
    enddo

    if(n /= nLocal) then
      error stop "ExtractSubMesh : local element count disagrees with setOffsets"
    endif

    ! Mortars wholly inside this set (ValidateSetPartition rejects any other), renumbered
    ! in their global order. The tables are replicated, so every rank copies all of them.
    nm = 0
    allocate(mortarNew(1:max(globalMesh%nMortars,1)))
    mortarNew = 0
    do m = 1,globalMesh%nMortars
      if(setOfElem(globalMesh%mortarInfo(1,m)) == iset) then
        nm = nm+1
        mortarNew(m) = nm
      endif
    enddo
    na = 0
    allocate(anisoNew(1:max(globalMesh%nAnisoMortars,1)))
    anisoNew = 0
    do a = 1,globalMesh%nAnisoMortars
      if(setOfElem(globalMesh%anisoMortarInfo(1,a)) == iset) then
        na = na+1
        anisoNew(a) = na
      endif
    enddo

    subMesh%nMortars = nm
    if(nm > 0) then
      allocate(subMesh%mortarInfo(1:14,1:nm))
      do m = 1,globalMesh%nMortars
        if(mortarNew(m) == 0) cycle
        subMesh%mortarInfo(1:14,mortarNew(m)) = globalMesh%mortarInfo(1:14,m)
        subMesh%mortarInfo(1,mortarNew(m)) = setGlobalId(globalMesh%mortarInfo(1,m))
        do q = 1,4
          subMesh%mortarInfo(1+2*q,mortarNew(m)) = setGlobalId(globalMesh%mortarInfo(1+2*q,m))
        enddo
        ! Mortar index on this rank's faces of the mortar (big face, then sub-faces).
        call MarkMortarFace(globalMesh%mortarInfo(1,m),globalMesh%mortarInfo(2,m),mortarNew(m))
        do q = 1,4
          call MarkMortarFace(globalMesh%mortarInfo(1+2*q,m), &
                              globalMesh%mortarInfo(2+2*q,m)/10,mortarNew(m))
        enddo
      enddo
    endif

    subMesh%nAnisoMortars = na
    if(na > 0) then
      allocate(subMesh%anisoMortarInfo(1:9,1:na))
      do a = 1,globalMesh%nAnisoMortars
        if(anisoNew(a) == 0) cycle
        subMesh%anisoMortarInfo(1:9,anisoNew(a)) = globalMesh%anisoMortarInfo(1:9,a)
        subMesh%anisoMortarInfo(1,anisoNew(a)) = setGlobalId(globalMesh%anisoMortarInfo(1,a))
        do q = 1,2
          subMesh%anisoMortarInfo(1+2*q,anisoNew(a)) = &
            setGlobalId(globalMesh%anisoMortarInfo(1+2*q,a))
        enddo
        ! Anisotropic mortar faces carry sideInfo(1) = nMortars + a.
        call MarkMortarFace(globalMesh%anisoMortarInfo(1,a),globalMesh%anisoMortarInfo(2,a), &
                            nm+anisoNew(a))
        do q = 1,2
          call MarkMortarFace(globalMesh%anisoMortarInfo(1+2*q,a), &
                              globalMesh%anisoMortarInfo(2+2*q,a)/10,nm+anisoNew(a))
        enddo
      enddo
    endif

    deallocate(subLocal,mortarNew,anisoNew)

    call subMesh%UpdateDevice()

  contains

    subroutine MarkMortarFace(eGlobal,face,mortarIndex)
      !! Set sideInfo(1) = mortarIndex on local side face of global element eGlobal when this
      !! rank owns it; elements on other ranks are marked by their owners.
      implicit none
      integer,intent(in) :: eGlobal
      integer,intent(in) :: face
      integer,intent(in) :: mortarIndex
      ! Local
      integer :: elLocal

      elLocal = eGlobal-goff
      if(elLocal < 1 .or. elLocal > globalMesh%nElem) return
      subMesh%sideInfo(1,face,subLocal(elLocal)) = mortarIndex

    endsubroutine MarkMortarFace

  endsubroutine ExtractSubMesh

endmodule SELF_ElementSetMesh_3D