!!******************************************************************************
!!
!!  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-2021 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: REFINEMENT
!!
!!  This module handles the error estimation and refinement criterion
!!  determination.
!!
!!
!!******************************************************************************
!
module refinement

#ifdef PROFILE
! import external subroutines
!
  use timers, only : set_timer, start_timer, stop_timer
#endif /* PROFILE */

! module variables are not implicit by default
!
  implicit none

#ifdef PROFILE
! timer indices
!
  integer            , save :: iri, irc
#endif /* PROFILE */

! refinement criterion parameters
!
  real(kind=8), save :: crefmin = 2.0d-01
  real(kind=8), save :: crefmax = 8.0d-01
  real(kind=8), save :: vortmin = 1.0d-03
  real(kind=8), save :: vortmax = 1.0d-01
  real(kind=8), save :: currmin = 1.0d-03
  real(kind=8), save :: currmax = 1.0d-01
  real(kind=8), save :: epsref  = 1.0d-02

! flags for variable included in the refinement criterion calculation
!
  logical, dimension(:), allocatable, save :: qvar_ref
  logical                           , save :: vort_ref = .false.
  logical                           , save :: curr_ref = .false.

! by default everything is private
!
  private

! declare public subroutines
!
  public :: initialize_refinement, finalize_refinement, print_refinement
  public :: check_refinement_criterion

!- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
!
  contains
!
!===============================================================================
!!
!!***  PUBLIC SUBROUTINES  *****************************************************
!!
!===============================================================================
!
!===============================================================================
!
! subroutine INITIALIZE_REFINEMENT:
! --------------------------------
!
!   Subroutine initializes module REFINEMENT.
!
!   Arguments:
!
!     verbose - flag determining if the subroutine should be verbose;
!     status  - return flag of the procedure execution status;
!
!===============================================================================
!
  subroutine initialize_refinement(verbose, status)

! import external procedures and variables
!
    use equations , only : magnetized, nv, pvars
    use parameters, only : get_parameter

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    logical, intent(in)  :: verbose
    integer, intent(out) :: status

! local variables
!
    logical            :: test
    integer            :: p
    character(len=255) :: variables = "dens pres"
!
!-------------------------------------------------------------------------------
!
#ifdef PROFILE
! set timer descriptions
!
    call set_timer('refinement:: initialization', iri)
    call set_timer('refinement:: criterion'     , irc)

! start accounting time for module initialization/finalization
!
    call start_timer(iri)
#endif /* PROFILE */

! reset the status flag
!
    status = 0

! get the refinement parameters
!
    call get_parameter("crefmin", crefmin)
    call get_parameter("crefmax", crefmax)
    call get_parameter("vortmin", vortmin)
    call get_parameter("vortmax", vortmax)
    call get_parameter("currmin", currmin)
    call get_parameter("currmax", currmax)
    call get_parameter("epsref" , epsref )

! get variables to include in the refinement criterion calculation
!
    call get_parameter("refinement_variables", variables)

! allocate vector for indicators, which variables are taken into account in
! calculating the refinement criterion
!
    allocate(qvar_ref(nv))

! check which primitive variable is used to determine the refinement criterion
!
    do p = 1, nv
      qvar_ref(p) = index(variables, trim(pvars(p))) > 0
    end do ! p = 1, nv

! turn on refinement based on vorticity if specified
!
    vort_ref = index(variables, 'vort') > 0

! check if any resonable variable controls the refinement
!
    test = any(qvar_ref(:)) .or. vort_ref

! turn on refinement based on current density if specified
!
    if (magnetized) then
      curr_ref = index(variables, 'curr') > 0 .or. index(variables, 'jabs') > 0
      test = test .or. curr_ref
    end if

! check if there is any variable selected to control refinement
!
    if (.not. test) then
      if (verbose) then
        write(*,*)
        write(*,"(1x,a)") "ERROR!"
        write(*,"(1x,a)") "No or wrong variable selected to control"   //      &
                          " the refinement."
        write(*,"(1x,a)") "Available control variables: any primitive" //      &
                          " variable, or 'vort' for the magnitude of"  //      &
                          " vorticity, or 'curr' for the magnitude of" //      &
                          " current density."
      end if
      status = 1
    end if

    if (crefmin > crefmax .or. crefmin < 0.0d+00) then
      if (verbose) then
        write(*,*)
        write(*,"(1x,a)") "ERROR!"
        write(*,"(1x,a)") "Wrong 'crefmin' or 'crefmax' parameters."
        write(*,"(1x,a)") "Both should be positive and 'crefmin' <= 'crefmax'."
      end if
      status = 1
    end if

    if (epsref < 0.0d+00) then
      if (verbose) then
        write(*,*)
        write(*,"(1x,a)") "ERROR!"
        write(*,"(1x,a)") "Wrong 'epsref' parameters."
        write(*,"(1x,a)") "It should be positive."
      end if
      status = 1
    end if

    if (vortmin > vortmax .or. vortmin < 0.0d+00) then
      if (verbose) then
        write(*,*)
        write(*,"(1x,a)") "ERROR!"
        write(*,"(1x,a)") "Wrong 'vortmin' or 'vortmax' parameters."
        write(*,"(1x,a)") "Both should be positive and 'vortmin' <= 'vortmax'."
      end if
      status = 1
    end if

    if ((currmin > currmax .or. currmin < 0.0d+00) .and. magnetized) then
      if (verbose) then
        write(*,*)
        write(*,"(1x,a)") "ERROR!"
        write(*,"(1x,a)") "Wrong 'currmin' or 'currmax' parameters."
        write(*,"(1x,a)") "Both should be positive and 'currmin' <= 'currmax'."
      end if
      status = 1
    end if

#ifdef PROFILE
! stop accounting time for module initialization/finalization
!
    call stop_timer(iri)
#endif /* PROFILE */

!-------------------------------------------------------------------------------
!
  end subroutine initialize_refinement
!
!===============================================================================
!
! subroutine FINALIZE_REFINEMENT:
! ------------------------------
!
!   Subroutine releases memory used by the module variables.
!
!   Arguments:
!
!     status - return flag of the procedure execution status;
!
!===============================================================================
!
  subroutine finalize_refinement(status)

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    integer, intent(out) :: status
!
!-------------------------------------------------------------------------------
!
#ifdef PROFILE
! start accounting time for module initialization/finalization
!
    call start_timer(iri)
#endif /* PROFILE */

! reset the status flag
!
    status = 0

! deallocate refined variable indicators
!
    if (allocated(qvar_ref)) deallocate(qvar_ref)

#ifdef PROFILE
! stop accounting time for module initialization/finalization
!
    call stop_timer(iri)
#endif /* PROFILE */

!-------------------------------------------------------------------------------
!
  end subroutine finalize_refinement
!
!===============================================================================
!
! subroutine PRINT_REFINEMENT:
! ---------------------------
!
!   Subroutine prints module parameters.
!
!   Arguments:
!
!     verbose - flag determining if the subroutine should be verbose;
!
!===============================================================================
!
  subroutine print_refinement(verbose)

! import external procedures and variables
!
    use helpers  , only : print_section, print_parameter
    use equations, only : magnetized, pvars, nv

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    logical, intent(in) :: verbose

! local variables
!
    character(len=80) :: rvars = "", msg
    character(len=64) :: sfmt
    integer           :: p
!
!-------------------------------------------------------------------------------
!
    if (verbose) then

      rvars = ""
      do p = 1, nv
        if (qvar_ref(p)) rvars = adjustl(trim(rvars) // ' ' // trim(pvars(p)))
      end do
      if (vort_ref) then
        rvars = adjustl(trim(rvars) // ' vort')
      end if
      if (magnetized .and. curr_ref) then
        rvars = adjustl(trim(rvars) // ' curr')
      end if

      call print_section(verbose, "Refinement")
      call print_parameter(verbose, "refined variables", rvars)
      sfmt = "(es9.2,1x,'...',1x,es9.2)"
      write(msg,sfmt) crefmin, crefmax
      call print_parameter(verbose, "2nd order error limits", msg)
      if (vort_ref) then
        write(msg,sfmt) vortmin, vortmax
        call print_parameter(verbose, "vorticity limits"      , msg)
      end if
      if (magnetized .and. curr_ref) then
        write(msg,sfmt) currmin, currmax
        call print_parameter(verbose, "current density limits", msg)
      end if

    end if

!-------------------------------------------------------------------------------
!
  end subroutine print_refinement
!
!===============================================================================
!
! function CHECK_REFINEMENT_CRITERION:
! -----------------------------------
!
!   Function scans the given data block and checks for the refinement
!   criterion.  It returns +1 if the criterion is met, which indicates that
!   the block needs to be refined, 0 if there is no need for the refinement,
!   and -1 if the block can be derefined.
!
!   Arguments:
!
!     pdata - pointer to the data block for which the refinement criterion
!             has to be determined;
!
!===============================================================================
!
  function check_refinement_criterion(pdata) result(criterion)

! import external procedures and variables
!
    use blocks         , only : block_data
    use equations      , only : nv

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    type(block_data), pointer, intent(in) :: pdata

! return variable
!
    integer(kind=4)                       :: criterion

! local variables
!
    integer      :: p
    real(kind=8) :: cref
!
!-------------------------------------------------------------------------------
!
#ifdef PROFILE
! start accounting time for the refinement criterion estimation
!
    call start_timer(irc)
#endif /* PROFILE */

! reset the refinement criterion flag
!
    criterion = -1

! check the second derivative error from selected primitive variables
!
    do p = 1, nv
      if (qvar_ref(p)) then

        cref = second_derivative_error(p, pdata)

        if (cref > crefmin) criterion = max(criterion, 0)
        if (cref > crefmax) criterion = max(criterion, 1)

      end if
    end do ! p = 1, nv

! check vorticity criterion
!
    if (vort_ref) then
      cref = vorticity_magnitude(pdata)

      if (cref > vortmin) criterion = max(criterion, 0)
      if (cref > vortmax) criterion = max(criterion, 1)
    end if

! check current density criterion
!
    if (curr_ref) then
      cref = current_density_magnitude(pdata)

      if (cref > currmin) criterion = max(criterion, 0)
      if (cref > currmax) criterion = max(criterion, 1)
    end if

#ifdef PROFILE
! stop accounting time for the refinement criterion estimation
!
    call stop_timer(irc)
#endif /* PROFILE */

! return the refinement flag
!
    return

!-------------------------------------------------------------------------------
!
  end function check_refinement_criterion
!
!===============================================================================
!!
!!***  PRIVATE SUBROUTINES  ****************************************************
!!
!===============================================================================
!
!===============================================================================
!
! function SECOND_DERIVATIVE_ERROR:
! --------------------------------
!
!   Function calculate the second derivative error for a given data block
!   and selected primitive variables.  The total error is returned then.
!
!   Arguments:
!
!     iqt   - the index of primitive variable;
!     pdata - pointer to the data block for which error is calculated;
!
!===============================================================================
!
  function second_derivative_error(iqt, pdata) result(error)

! import external procedures and variables
!
    use blocks     , only : block_data
    use coordinates, only : nbl, neu

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    integer                  , intent(in) :: iqt
    type(block_data), pointer, intent(in) :: pdata

! return variable
!
    real(kind=8) :: error

! local variables
!
    integer      :: i, im1, ip1
    integer      :: j, jm1, jp1
    integer      :: k = 1
#if NDIMS == 3
    integer      :: km1, kp1
#endif /* NDIMS == 3 */
    real(kind=8) :: fl, fr, fc, fx, fy
#if NDIMS == 3
    real(kind=8) :: fz
#endif /* NDIMS == 3 */

! local parameters
!
    real(kind=8), parameter :: eps = epsilon(1.0d+00)
!
!-------------------------------------------------------------------------------
!
! reset indicators
!
    error = 0.0e+00

! calculate local refinement criterion for variable which exists
!
    if (iqt > 0) then

! find the maximum smoothness indicator over all cells
!
#if NDIMS == 3
      do k = nbl, neu
        km1 = k - 1
        kp1 = k + 1
#endif /* NDIMS == 3 */
        do j = nbl, neu
          jm1 = j - 1
          jp1 = j + 1
          do i = nbl, neu
            im1 = i - 1
            ip1 = i + 1

! calculate the second derivative error the X direction
!
            fr   = pdata%q(iqt,ip1,j,k) - pdata%q(iqt,i  ,j,k)
            fl   = pdata%q(iqt,im1,j,k) - pdata%q(iqt,i  ,j,k)
            fc   = abs(pdata%q(iqt,ip1,j,k)) + abs(pdata%q(iqt,im1,j,k))       &
                                           + 2.0d+00 * abs(pdata%q(iqt,i,j,k))
            fx   = abs(fr + fl) / (abs(fr) + abs(fl) + epsref * fc + eps)

! calculate the second derivative error along the Y direction
!
            fr   = pdata%q(iqt,i,jp1,k) - pdata%q(iqt,i,j  ,k)
            fl   = pdata%q(iqt,i,jm1,k) - pdata%q(iqt,i,j  ,k)
            fc   = abs(pdata%q(iqt,i,jp1,k)) + abs(pdata%q(iqt,i,jm1,k))       &
                                           + 2.0d+00 * abs(pdata%q(iqt,i,j,k))
            fy   = abs(fr + fl) / (abs(fr) + abs(fl) + epsref * fc + eps)

#if NDIMS == 3
! calculate the second derivative error along the Z direction
!
            fr   = pdata%q(iqt,i,j,kp1) - pdata%q(iqt,i,j,k  )
            fl   = pdata%q(iqt,i,j,km1) - pdata%q(iqt,i,j,k  )
            fc   = abs(pdata%q(iqt,i,j,kp1)) + abs(pdata%q(iqt,i,j,km1))       &
                                           + 2.0d+00 * abs(pdata%q(iqt,i,j,k))
            fz   = abs(fr + fl) / (abs(fr) + abs(fl) + epsref * fc + eps)
#endif /* NDIMS == 3 */

! take the maximum second derivative error
!
#if NDIMS == 2
            error = max(error, fx, fy)
#endif /* NDIMS == 2 */
#if NDIMS == 3
            error = max(error, fx, fy, fz)
#endif /* NDIMS == 3 */

          end do ! i = nbl, neu
        end do ! j = nbl, neu
#if NDIMS == 3
      end do ! k = nbl, neu
#endif /* NDIMS == 3 */

    end if ! iqt > 0

! return the refinement flag
!
    return

!-------------------------------------------------------------------------------
!
  end function second_derivative_error
!
!===============================================================================
!
! function VORTICITY_MAGNITUDE:
! ----------------------------
!
!   Function finds the maximum values of the vorticity magnitude
!   for the current data block.
!
!   Arguments:
!
!     pdata - pointer to the data block for which error is calculated;
!
!===============================================================================
!
  function vorticity_magnitude(pdata) result(wmax)

    use blocks         , only : block_data
    use coordinates    , only : nn => bcells
    use coordinates    , only : nbl, neu
    use equations      , only : ivx, ivz
    use iso_fortran_env, only : error_unit
    use operators      , only : curl
    use workspace      , only : work, nwork, work_in_use

    implicit none

    type(block_data), pointer, intent(in) :: pdata

    real(kind=4)  :: wmax

    logical, save :: first = .true.

    integer       :: i, j, k = 1, status
    real(kind=8)  :: vort

    real(kind=8), dimension(3), save :: dh

    real(kind=8), dimension(:,:,:,:), pointer, save :: w

    character(len=*), parameter :: loc = 'REFINEMENT:vorticity_magnitude()'

!-------------------------------------------------------------------------------
!
    wmax = 0.0e+00

    if (first) then
      status = 0
      if (work_in_use) then
        write(error_unit,"('[',a,']: ',a,3i4,a)") trim(loc),                   &
                         "Workspace is already occupied!"
        return
      else
        i = 3 * nn**NDIMS

        if (i > nwork) then
          write(error_unit,"('[',a,']: ',a,3i4,a)") trim(loc),                 &
                           "Workspace has too be increased!"
          nwork = i
          deallocate(work)
          allocate(work(nwork), stat=status)
          if (status /= 0) then
            write(error_unit,"('[',a,']: ',a)") trim(loc),                     &
                             "Cannot increase the workspace's size!"
            return
          end if
        end if

#if NDIMS == 3
        w(1:3,1:nn,1:nn,1:nn) => work(1:i)
#else /* NDIMS == 3 */
        w(1:3,1:nn,1:nn,1: 1) => work(1:i)
#endif /* NDIMS == 3 */

        dh(:) = 1.0d+00
      end if

      first = .false.
    end if

    work_in_use = .true.

    call curl(dh(:), pdata%q(ivx:ivz,:,:,:), w(:,:,:,:))

#if NDIMS == 3
    do k = nbl, neu
#endif /* NDIMS == 3 */
      do j = nbl, neu
        do i = nbl, neu

          vort = sum(w(:,i,j,k)**2)

          wmax = max(wmax, real(vort, kind=4))

        end do
      end do
#if NDIMS == 3
    end do
#endif /* NDIMS == 3 */

    work_in_use = .false.

    wmax = sqrt(wmax)

!-------------------------------------------------------------------------------
!
  end function vorticity_magnitude
!
!===============================================================================
!
! function CURRENT_DENSITY_MAGNITUDE:
! ----------------------------------
!
!   Function finds the maximum values of the current density magnitude
!   for the current data block.
!
!   Arguments:
!
!     pdata - pointer to the data block for which error is calculated;
!
!===============================================================================
!
  function current_density_magnitude(pdata) result(jmax)

    use blocks         , only : block_data
    use coordinates    , only : nn => bcells
    use coordinates    , only : nbl, neu
    use equations      , only : magnetized
    use equations      , only : ibx, ibz
    use iso_fortran_env, only : error_unit
    use operators      , only : curl
    use workspace      , only : work, nwork, work_in_use

    implicit none

    type(block_data), pointer, intent(in) :: pdata

    real(kind=4)  :: jmax

    logical, save :: first = .true.

    integer       :: i, j, k = 1, status
    real(kind=8)  :: jabs

    real(kind=8), dimension(3), save :: dh

    real(kind=8), dimension(:,:,:,:), pointer, save :: w

    character(len=*), parameter :: loc = 'REFINEMENT:current_density_magnitude()'

!-------------------------------------------------------------------------------
!
    jmax = 0.0e+00

    if (.not. magnetized) return

    if (first) then
      status = 0
      if (work_in_use) then
        write(error_unit,"('[',a,']: ',a,3i4,a)") trim(loc),                   &
                         "Workspace is already occupied!"
        return
      else
        i = 3 * nn**NDIMS

        if (i > nwork) then
          write(error_unit,"('[',a,']: ',a,3i4,a)") trim(loc),                 &
                           "Workspace has too be increased!"
          nwork = i
          deallocate(work)
          allocate(work(nwork), stat=status)
          if (status /= 0) then
            write(error_unit,"('[',a,']: ',a)") trim(loc),                     &
                             "Cannot increase the workspace's size!"
            return
          end if
        end if

#if NDIMS == 3
        w(1:3,1:nn,1:nn,1:nn) => work(1:i)
#else /* NDIMS == 3 */
        w(1:3,1:nn,1:nn,1: 1) => work(1:i)
#endif /* NDIMS == 3 */

        dh(:) = 1.0d+00
      end if

      first = .false.
    end if

    work_in_use = .true.

    call curl(dh(:), pdata%q(ibx:ibz,:,:,:), w(:,:,:,:))

#if NDIMS == 3
    do k = nbl, neu
#endif /* NDIMS == 3 */
      do j = nbl, neu
        do i = nbl, neu

          jabs = sum(w(:,i,j,k)**2)

          jmax = max(jmax, real(jabs, kind=4))

        end do ! i = nbl, neu
      end do ! j = nbl, neu
#if NDIMS == 3
    end do ! k = nbl, neu
#endif /* NDIMS == 3 */

    work_in_use = .false.

    jmax = sqrt(jmax)

!-------------------------------------------------------------------------------
!
  end function current_density_magnitude

!===============================================================================
!
end module refinement