!!******************************************************************************
!!
!!  This file is part of the AMUN source code, a program to perform
!!  Newtonian or relativistic magnetohydrodynamical simulations on uniform or
!!  adaptive mesh.
!!
!!  Copyright (C) 2008-2020 Grzegorz Kowal <grzegorz@amuncode.org>
!!
!!  This program is free software: you can redistribute it and/or modify
!!  it under the terms of the GNU General Public License as published by
!!  the Free Software Foundation, either version 3 of the License, or
!!  (at your option) any later version.
!!
!!  This program is distributed in the hope that it will be useful,
!!  but WITHOUT ANY WARRANTY; without even the implied warranty of
!!  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!!  GNU General Public License for more details.
!!
!!  You should have received a copy of the GNU General Public License
!!  along with this program.  If not, see <http://www.gnu.org/licenses/>.
!!
!!******************************************************************************
!!
!! module: DOMAINS
!!
!!  This module handles the initialization of the problem domains.
!!
!!
!!******************************************************************************
!
module domains

! module variables are not implicit by default
!
  implicit none

! interfaces for procedure pointers
!
  abstract interface
    subroutine setup_domain_iface(status)
      integer, intent(out) :: status
    end subroutine
  end interface

! pointer to the problem setup subroutine
!
  procedure(setup_domain_iface), pointer, save :: setup_domain => null()

! by default everything is private
!
  private

! declare public subroutines
!
  public :: initialize_domains, finalize_domains
  public :: setup_domain

!- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!
  contains
!
!===============================================================================
!!
!!***  PUBLIC SUBROUTINES  *****************************************************
!!
!===============================================================================
!
!===============================================================================
!
! subroutine INITIALIZE_DOMAINS:
! -----------------------------
!
!   Subroutine prepares module DOMAINS.
!
!   Arguments:
!
!     problem - the problem name
!     verbose - a logical flag turning the information printing;
!     status  - the subroutine call status: 0 for success, otherwise failure;
!
!===============================================================================
!
  subroutine initialize_domains(problem, verbose, status)

! include external procedures and variables
!
    use parameters, only : get_parameter

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    character(len=64), intent(in)  :: problem
    logical          , intent(in)  :: verbose
    integer          , intent(out) :: status
!
!-------------------------------------------------------------------------------
!
! reset the status flag
!
    status = 0

! associate the setup_domain pointer with the respective problem setup
! subroutine
!
    select case(trim(problem))

    case default
      setup_domain => setup_domain_default

    end select

!-------------------------------------------------------------------------------
!
  end subroutine initialize_domains
!
!===============================================================================
!
! subroutine FINALIZE_DOMAINS:
! ---------------------------
!
!   Subroutine releases memory used by the module.
!
!   Arguments:
!
!     status  - the subroutine call status: 0 for success, otherwise failure;
!
!===============================================================================
!
  subroutine finalize_domains(status)

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    integer, intent(out) :: status
!
!-------------------------------------------------------------------------------
!
! reset the status flag
!
    status = 0

!-------------------------------------------------------------------------------
!
  end subroutine finalize_domains
!
!===============================================================================
!!
!!***  PRIVATE SUBROUTINES  ****************************************************
!!
!===============================================================================
!
!===============================================================================
!
! subroutine SETUP_DOMAIN_DEFAULT:
! -------------------------------
!
!   Subroutine sets the default domain of N₁xN₂xN₃ blocks in the right
!   configuration.
!
!   Arguments:
!
!     status  - the subroutine call status: 0 for success, otherwise failure;
!
!===============================================================================
!
  subroutine setup_domain_default(status)

! include external procedures and variables
!
    use blocks         , only : pointer_meta, block_meta, block_data
    use blocks         , only : append_metablock, append_datablock, link_blocks
    use blocks         , only : metablock_set_leaf, metablock_set_level
    use blocks         , only : metablock_set_configuration
    use blocks         , only : metablock_set_coordinates, metablock_set_bounds
    use blocks         , only : nsides
    use coordinates    , only : xmin, ymin, zmin, xlen, ylen, zlen
    use coordinates    , only : ddims => domain_base_dims
    use coordinates    , only : periodic
    use iso_fortran_env, only : error_unit

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    integer, intent(out) :: status

! local variables
!
    integer      :: i, j, k, n, p, ic, jc, kc
    real(kind=8) :: xl, xmn, xmx, yl, ymn, ymx, zl, zmn, zmx

! local arrays
!
    integer, dimension(3) :: pos, del

! local pointers
!
    type(block_meta), pointer :: pmeta, pnext

! allocatable arrays
!
    integer, dimension(:,:,:), allocatable :: cfg
    integer, dimension(:)    , allocatable :: im, jm, km
    integer, dimension(:)    , allocatable :: ip, jp, kp

! local pointer array
!
    type(pointer_meta), dimension(:,:,:), allocatable :: block_array

! local parameters
!
    character(len=*), parameter :: loc = 'DOMAINS::setup_domain_default()'
!
!-------------------------------------------------------------------------------
!
! reset the status flag
!
    status = 0

! obtain the number of blocks
!
    n = product(ddims(1:NDIMS))

!! PREPARE BLOCK CONFIGURATION ARRAY
!!
! allocate the configuration array
!
    allocate(cfg(ddims(1),ddims(2),ddims(3)), stat = status)

! check if the allocation succeeded
!
    if (status /= 0) then
      write(error_unit,"('[',a,']: ',a)") trim(loc)                            &
                      , "Cannot allocate the configuration array!"
      status = 1
      return
    end if ! allocation

! set the block configurations
!
    cfg(1:ddims(1),1:ddims(2):2,1:ddims(3):2) = 12

    if (ddims(2) > 1) then
      cfg(1:ddims(1),2:ddims(2):2,1:ddims(3):2) = 43
      cfg(  ddims(1),1:ddims(2)  ,1:ddims(3):2) = 13
    end if

    if (ddims(3) > 1) then
      cfg(1:ddims(1),1:ddims(2):2,2:ddims(3):2) = 65
      if (ddims(2) > 1) then
        cfg(1:ddims(1),2:ddims(2):2,2:ddims(3):2) = 78
        cfg(  ddims(1),1:ddims(2)  ,2:ddims(3):2) = 75
      end if
      if (ddims(1) == 1 .or. mod(ddims(2),2) == 1) then
        cfg(  ddims(1),  ddims(2)  ,1:ddims(3)  ) = 15
      else
        cfg(        1 ,  ddims(2)  ,1:ddims(3)  ) = 48
      end if
    end if

!! ALLOCATE AND GENERATE META BLOCK CHAIN AND SET BLOCK CONFIGURATIONS
!!
! allocate the block pointer array
!
    allocate(block_array(ddims(1),ddims(2),ddims(3)), stat = status)

! check if the allocation succeeded
!
    if (status /= 0) then
      write(error_unit,"('[',a,']: ',a)") trim(loc)                            &
                      , "Cannot allocate the pointer array for domain base!"
      if (allocated(cfg)) deallocate(cfg, stat = status)
      status = 1
      return
    end if ! allocation

! generate the gray code for a given configuration and link the block in
! the proper order
!
    pos(:) = (/ 0, 0, 0 /)
    del(:) = (/ 1, 1, 1 /)

    p = 1
    do k = 1, ddims(3)
      if (del(3) == 1) pos(3) = pos(3) + del(3)
      do j = 1, ddims(2)
        if (del(2) == 1) pos(2) = pos(2) + del(2)
        do i = 1, ddims(1)
          if (del(1) == 1) pos(1) = pos(1) + del(1)

! append a new metablock
!
          call append_metablock(block_array(pos(1),pos(2),pos(3))%ptr, status)

! set the configuration type, if new block was appended successfully
!
          if (status == 0) then
            call metablock_set_configuration(                                  &
                                        block_array(pos(1),pos(2),pos(3))%ptr  &
                                              , cfg(pos(1),pos(2),pos(3)))
          else ! status
            write(error_unit,"('[',a,']: ',a)") trim(loc)                      &
                            , "Cannot append new metablock to domain base!"
            if (allocated(block_array)) deallocate(block_array, stat = status)
            if (allocated(cfg))         deallocate(cfg        , stat = status)
            status = 1
            return
          end if ! status

          p = p + 1

          if (del(1) == -1) pos(1) = pos(1) + del(1)
        end do
        if (del(2) == -1) pos(2) = pos(2) + del(2)
        del(1) = - del(1)
      end do
      if (del(3) == -1) pos(3) = pos(3) + del(3)
      del(2) = - del(2)
    end do

! deallocate the configuration array
!
    deallocate(cfg, stat = status)

! check if the deallocation succeeded
!
    if (status /= 0) then
      write(error_unit,"('[',a,']: ',a)") trim(loc)                            &
                      , "Cannot deallocate the configuration array!"
      status = 1
    end if ! deallocation

!! FILL OUT THE REMAINING FIELDS AND ALLOCATE AND ASSOCIATE DATA BLOCKS
!!
! calculate block sizes
!
    xl = xlen / ddims(1)
    yl = ylen / ddims(2)
    zl = zlen / ddims(3)

! fill out block structure fields
!
    do k = 1, ddims(3)

! claculate the block position along Z
!
      kc  = k - 1

! calculate the Z bounds
!
      zmn = zmin + kc * zl
      zmx = zmin + k  * zl

      do j = 1, ddims(2)

! claculate the block position along Y
!
        jc  = j - 1

! calculate the Y bounds
!
        ymn = ymin + jc * yl
        ymx = ymin + j  * yl

        do i = 1, ddims(1)

! claculate the block position along Y
!
          ic  = i - 1

! calculate the Z bounds
!
          xmn = xmin + ic * xl
          xmx = xmin + i  * xl

! assign a pointer
!
          pmeta => block_array(i,j,k)%ptr

! mark it as the leaf
!
          call metablock_set_leaf(pmeta)

! set the level
!
          call metablock_set_level(pmeta, 1)

! set block coordinates
!
          call metablock_set_coordinates(pmeta, ic, jc, kc)

! set the bounds
!
          call metablock_set_bounds(pmeta, xmn, xmx, ymn, ymx, zmn, zmx)
        end do
      end do
    end do

!! ASSIGN THE BLOCK NEIGHBORS
!!
! allocate indices
!
#if NDIMS == 3
   allocate(im(ddims(1)), ip(ddims(1)), jm(ddims(2)), jp(ddims(2)),            &
            km(ddims(3)), kp(ddims(3)), stat = status)
#else /* NDIMS == 3 */
   allocate(im(ddims(1)), ip(ddims(1)), jm(ddims(2)), jp(ddims(2)),            &
            stat = status)
#endif /* NDIMS == 3 */

! check if the allocation succeeded
!
    if (status == 0) then

! generate indices
!
     im(:) = cshift((/(i, i = 1, ddims(1))/),-1)
     ip(:) = cshift((/(i, i = 1, ddims(1))/), 1)
     jm(:) = cshift((/(j, j = 1, ddims(2))/),-1)
     jp(:) = cshift((/(j, j = 1, ddims(2))/), 1)
#if NDIMS == 3
     km(:) = cshift((/(k, k = 1, ddims(3))/),-1)
     kp(:) = cshift((/(k, k = 1, ddims(3))/), 1)
#endif /* NDIMS == 3 */

! check periodicity and reset the edge indices if box is not periodic
!
     if (.not. periodic(1)) then
       im(       1) = 0
       ip(ddims(1)) = 0
     end if
     if (.not. periodic(2)) then
       jm(       1) = 0
       jp(ddims(2)) = 0
     end if
#if NDIMS == 3
     if (.not. periodic(3)) then
       km(       1) = 0
       kp(ddims(3)) = 0
     end if
#endif /* NDIMS == 3 */

! iterate over all initial blocks
!
      do k = 1, ddims(3)
        do j = 1, ddims(2)
          do i = 1, ddims(1)

! assign pmeta with the current block
!
            pmeta => block_array(i,j,k)%ptr

#if NDIMS == 3
! assign face neighbor pointers
!
            if (im(i) > 0) then
              pmeta%faces(1,1,1,1)%ptr => block_array(im(i),j,k)%ptr
              pmeta%faces(1,2,1,1)%ptr => block_array(im(i),j,k)%ptr
              pmeta%faces(1,1,2,1)%ptr => block_array(im(i),j,k)%ptr
              pmeta%faces(1,2,2,1)%ptr => block_array(im(i),j,k)%ptr
            end if
            if (ip(i) > 0) then
              pmeta%faces(2,1,1,1)%ptr => block_array(ip(i),j,k)%ptr
              pmeta%faces(2,2,1,1)%ptr => block_array(ip(i),j,k)%ptr
              pmeta%faces(2,1,2,1)%ptr => block_array(ip(i),j,k)%ptr
              pmeta%faces(2,2,2,1)%ptr => block_array(ip(i),j,k)%ptr
            end if

            if (jm(j) > 0) then
              pmeta%faces(1,1,1,2)%ptr => block_array(i,jm(j),k)%ptr
              pmeta%faces(2,1,1,2)%ptr => block_array(i,jm(j),k)%ptr
              pmeta%faces(1,1,2,2)%ptr => block_array(i,jm(j),k)%ptr
              pmeta%faces(2,1,2,2)%ptr => block_array(i,jm(j),k)%ptr
            end if
            if (jp(j) > 0) then
              pmeta%faces(1,2,1,2)%ptr => block_array(i,jp(j),k)%ptr
              pmeta%faces(2,2,1,2)%ptr => block_array(i,jp(j),k)%ptr
              pmeta%faces(1,2,2,2)%ptr => block_array(i,jp(j),k)%ptr
              pmeta%faces(2,2,2,2)%ptr => block_array(i,jp(j),k)%ptr
            end if

            if (km(k) > 0) then
              pmeta%faces(1,1,1,3)%ptr => block_array(i,j,km(k))%ptr
              pmeta%faces(2,1,1,3)%ptr => block_array(i,j,km(k))%ptr
              pmeta%faces(1,2,1,3)%ptr => block_array(i,j,km(k))%ptr
              pmeta%faces(2,2,1,3)%ptr => block_array(i,j,km(k))%ptr
            end if
            if (kp(k) > 0) then
              pmeta%faces(1,1,2,3)%ptr => block_array(i,j,kp(k))%ptr
              pmeta%faces(2,1,2,3)%ptr => block_array(i,j,kp(k))%ptr
              pmeta%faces(1,2,2,3)%ptr => block_array(i,j,kp(k))%ptr
              pmeta%faces(2,2,2,3)%ptr => block_array(i,j,kp(k))%ptr
            end if
#endif /* NDIMS == 3 */

! assign edge neighbor pointers
!
#if NDIMS == 2
            if (im(i) > 0) then
              pmeta%edges(1,1,2)%ptr => block_array(im(i),j,k)%ptr
              pmeta%edges(1,2,2)%ptr => block_array(im(i),j,k)%ptr
            end if
            if (ip(i) > 0) then
              pmeta%edges(2,1,2)%ptr => block_array(ip(i),j,k)%ptr
              pmeta%edges(2,2,2)%ptr => block_array(ip(i),j,k)%ptr
            end if
            if (jm(j) > 0) then
              pmeta%edges(1,1,1)%ptr => block_array(i,jm(j),k)%ptr
              pmeta%edges(2,1,1)%ptr => block_array(i,jm(j),k)%ptr
            end if
            if (jp(j) > 0) then
              pmeta%edges(1,2,1)%ptr => block_array(i,jp(j),k)%ptr
              pmeta%edges(2,2,1)%ptr => block_array(i,jp(j),k)%ptr
            end if
#endif /* NDIMS == 2 */
#if NDIMS == 3
            if (jm(j) > 0 .and. km(k) > 0) then
              pmeta%edges(1,1,1,1)%ptr => block_array(i,jm(j),km(k))%ptr
              pmeta%edges(2,1,1,1)%ptr => block_array(i,jm(j),km(k))%ptr
            end if
            if (jp(j) > 0 .and. km(k) > 0) then
              pmeta%edges(1,2,1,1)%ptr => block_array(i,jp(j),km(k))%ptr
              pmeta%edges(2,2,1,1)%ptr => block_array(i,jp(j),km(k))%ptr
            end if
            if (jm(j) > 0 .and. kp(k) > 0) then
              pmeta%edges(1,1,2,1)%ptr => block_array(i,jm(j),kp(k))%ptr
              pmeta%edges(2,1,2,1)%ptr => block_array(i,jm(j),kp(k))%ptr
            end if
            if (jp(j) > 0 .and. kp(k) > 0) then
              pmeta%edges(1,2,2,1)%ptr => block_array(i,jp(j),kp(k))%ptr
              pmeta%edges(2,2,2,1)%ptr => block_array(i,jp(j),kp(k))%ptr
            end if

            if (im(i) > 0 .and. km(k) > 0) then
              pmeta%edges(1,1,1,2)%ptr => block_array(im(i),j,km(k))%ptr
              pmeta%edges(1,2,1,2)%ptr => block_array(im(i),j,km(k))%ptr
            end if
            if (ip(i) > 0 .and. km(k) > 0) then
              pmeta%edges(2,1,1,2)%ptr => block_array(ip(i),j,km(k))%ptr
              pmeta%edges(2,2,1,2)%ptr => block_array(ip(i),j,km(k))%ptr
            end if
            if (im(i) > 0 .and. kp(k) > 0) then
              pmeta%edges(1,1,2,2)%ptr => block_array(im(i),j,kp(k))%ptr
              pmeta%edges(1,2,2,2)%ptr => block_array(im(i),j,kp(k))%ptr
            end if
            if (ip(i) > 0 .and. kp(k) > 0) then
              pmeta%edges(2,1,2,2)%ptr => block_array(ip(i),j,kp(k))%ptr
              pmeta%edges(2,2,2,2)%ptr => block_array(ip(i),j,kp(k))%ptr
            end if

            if (im(i) > 0 .and. jm(j) > 0) then
              pmeta%edges(1,1,1,3)%ptr => block_array(im(i),jm(j),k)%ptr
              pmeta%edges(1,1,2,3)%ptr => block_array(im(i),jm(j),k)%ptr
            end if
            if (ip(i) > 0 .and. jm(j) > 0) then
              pmeta%edges(2,1,1,3)%ptr => block_array(ip(i),jm(j),k)%ptr
              pmeta%edges(2,1,2,3)%ptr => block_array(ip(i),jm(j),k)%ptr
            end if
            if (im(i) > 0 .and. jp(j) > 0) then
              pmeta%edges(1,2,1,3)%ptr => block_array(im(i),jp(j),k)%ptr
              pmeta%edges(1,2,2,3)%ptr => block_array(im(i),jp(j),k)%ptr
            end if
            if (ip(i) > 0 .and. jp(j) > 0) then
              pmeta%edges(2,2,1,3)%ptr => block_array(ip(i),jp(j),k)%ptr
              pmeta%edges(2,2,2,3)%ptr => block_array(ip(i),jp(j),k)%ptr
            end if
#endif /* NDIMS == 3 */

! assign corner neighbor pointers
!
#if NDIMS == 2
            if (im(i) > 0 .and. jm(j) > 0)                                     &
                      pmeta%corners(1,1)%ptr => block_array(im(i),jm(j),k)%ptr
            if (ip(i) > 0 .and. jm(j) > 0)                                     &
                      pmeta%corners(2,1)%ptr => block_array(ip(i),jm(j),k)%ptr
            if (im(i) > 0 .and. jp(j) > 0)                                     &
                      pmeta%corners(1,2)%ptr => block_array(im(i),jp(j),k)%ptr
            if (ip(i) > 0 .and. jp(j) > 0)                                     &
                      pmeta%corners(2,2)%ptr => block_array(ip(i),jp(j),k)%ptr
#endif /* NDIMS == 2 */
#if NDIMS == 3
            if (im(i) > 0 .and. jm(j) > 0 .and. km(k) > 0)                     &
                pmeta%corners(1,1,1)%ptr => block_array(im(i),jm(j),km(k))%ptr
            if (ip(i) > 0 .and. jm(j) > 0 .and. km(k) > 0)                     &
                pmeta%corners(2,1,1)%ptr => block_array(ip(i),jm(j),km(k))%ptr
            if (im(i) > 0 .and. jp(j) > 0 .and. km(k) > 0)                     &
                pmeta%corners(1,2,1)%ptr => block_array(im(i),jp(j),km(k))%ptr
            if (ip(i) > 0 .and. jp(j) > 0 .and. km(k) > 0)                     &
                pmeta%corners(2,2,1)%ptr => block_array(ip(i),jp(j),km(k))%ptr
            if (im(i) > 0 .and. jm(j) > 0 .and. kp(k) > 0)                     &
                pmeta%corners(1,1,2)%ptr => block_array(im(i),jm(j),kp(k))%ptr
            if (ip(i) > 0 .and. jm(j) > 0 .and. kp(k) > 0)                     &
                pmeta%corners(2,1,2)%ptr => block_array(ip(i),jm(j),kp(k))%ptr
            if (im(i) > 0 .and. jp(j) > 0 .and. kp(k) > 0)                     &
                pmeta%corners(1,2,2)%ptr => block_array(im(i),jp(j),kp(k))%ptr
            if (ip(i) > 0 .and. jp(j) > 0 .and. kp(k) > 0)                     &
                pmeta%corners(2,2,2)%ptr => block_array(ip(i),jp(j),kp(k))%ptr
#endif /* NDIMS == 3 */
          end do ! over i
        end do ! over j
      end do ! over k

! deallocate indices
!
#if NDIMS == 3
      deallocate(im, ip, jm, jp, km, kp, stat = status)
#else /* NDIMS == 3 */
      deallocate(im, ip, jm, jp, stat = status)
#endif /* NDIMS == 3 */

! check if the deallocation succeeded
!
      if (status /= 0) then
        write(error_unit,"('[',a,']: ',a)") trim(loc)                          &
                        , "Cannot deallocate the configuration indices!"
        status = 1
      end if ! deallocation

    else ! allocation
      write(error_unit,"('[',a,']: ',a)") trim(loc)                            &
                      , "Cannot allocate the configuration indices!"
      status = 1
    end if ! allocation

! deallocate the block pointer array
!
    deallocate(block_array, stat = status)

! check if the deallocation succeeded
!
    if (status /= 0) then
      write(error_unit,"('[',a,']: ',a)") trim(loc)                            &
                      , "Cannot deallocate the pointer array for domain base!"
      status = 1
    end if ! deallocation

!-------------------------------------------------------------------------------
!
  end subroutine setup_domain_default

!===============================================================================
!
end module domains