!!******************************************************************************
!!
!!  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) 2017 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: FORCING
!!
!!  This module provides energy injection, e.g. turbulence driving, supernova
!!  explosions, etc.
!!
!!*****************************************************************************
!
module forcing

#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 :: ifi, ifu
#endif /* PROFILE */

! flag indicating if the energy injection is enabled
!
  logical, save :: forcing_enabled = .false.

! forcing parameters
!
  real(kind=8), save :: power = 1.0d+00
  real(kind=8), save :: rate  = 1.0d+03
  real(kind=8), save :: amp   = 1.0d-03
  real(kind=8), save :: del   = 1.0d-01

! module parameters
!
  real(kind=8), save :: dinj = 0.0d+00

! by default everything is private
!
  private

! declare public subroutines
!
  public :: initialize_forcing, finalize_forcing
  public :: update_forcing

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

! import external procedures and variables
!
    use parameters, only : get_parameter_string, get_parameter_real

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    logical, intent(in)    :: verbose
    integer, intent(inout) :: iret

! local variables
!
    character(len=64) :: injection_method = "none"
!
!-------------------------------------------------------------------------------
!
#ifdef PROFILE
! set timer descriptions
!
    call set_timer('forcing:: initialization', ifi)
    call set_timer('forcing:: update'        , ifu)

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

! obtain the chosen injection method
!
    call get_parameter_string("injection_method", injection_method)
    call get_parameter_real  ("injection_power" , power           )
    call get_parameter_real  ("injection_rate"  , rate            )
    call get_parameter_real  ("eddy_amplitude"  , amp             )
    call get_parameter_real  ("eddy_width"      , del             )

! select the energy injection method
!
    select case(trim(injection_method))
    case ("eddy", "random")
      forcing_enabled = .true.
    case default
      injection_method = "none"
    end select

! print information about the energy injection
!
    if (verbose) then

      write (*,"(4x,a16, 7x,'=',1x,a)") "energy injection"                     &
                                                      , trim(injection_method)

    end if

#ifdef PROFILE
! stop accounting time
!
    call stop_timer(ifi)
#endif /* PROFILE */

!-------------------------------------------------------------------------------
!
  end subroutine initialize_forcing
!
!===============================================================================
!
! subroutine FINALIZE_FORCING:
! ---------------------------
!
!   Subroutine releases memory used by the module variables.
!
!   Arguments:
!
!     iret    - return flag of the procedure execution status;
!
!===============================================================================
!
  subroutine finalize_forcing(iret)

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    integer, intent(inout) :: iret
!
!-------------------------------------------------------------------------------
!
#ifdef PROFILE
! start accounting time for initialization/finalization
!
    call start_timer(ifi)
#endif /* PROFILE */

#ifdef PROFILE
! stop accounting time
!
    call stop_timer(ifi)
#endif /* PROFILE */

!-------------------------------------------------------------------------------
!
  end subroutine finalize_forcing
!
!===============================================================================
!
! subroutine UPDATE_FORCING:
! -------------------------
!
!   Subroutine adds the energy injection terms.
!
!   Arguments:
!
!     t, dt - time and its increment;
!
!===============================================================================
!
  subroutine update_forcing(t, dt)

! import external procedures and variables
!
    use coordinates, only : xmin, ymin, zmin, xlen, ylen, zlen
    use random     , only : randomu, randomn

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    real(kind=8), intent(in) :: t, dt

! local variables
!
    integer                    :: ni, n
    real(kind=8)               :: tmp
    real(kind=8), dimension(3) :: xp, ap
!
!-------------------------------------------------------------------------------
!
#ifdef PROFILE
! start accounting time for forcing term update
!
    call start_timer(ifu)
#endif /* PROFILE */

! proceed only if forcing is enabled
!
    if (forcing_enabled) then

! calculate the number of eddies to be injected during this interval
!
      dinj = dinj + rate * dt
      ni   = int(floor(dinj))
      dinj = dinj - 1.0d+00 * ni

! inject the required number of eddies
!
      if (ni > 0) then

        do n = 1, ni

! get random position
!
          xp(1) = xmin + xlen * randomu()
          xp(2) = ymin + ylen * randomu()
          xp(3) = zmin + zlen * randomu()

! get random orientation
!
#if NDIMS == 3
          tmp = 0.0d+00
          do while(tmp == 0.0d+00)
            ap(1) = randomn()
            ap(2) = randomn()
            ap(3) = randomn()
            tmp   = sqrt(ap(1)**2 + ap(2)**2 + ap(3)**2)
          end do
          ap(:) = amp * ap(:) / tmp / del * dt
#else /* NDIMS == 3 */
          ap(:) = sign(1.0d+00, randomn()) * (/ 0.0d+00, 0.0d+00, amp / del /) * dt
#endif /* NDIMS == 3 */

! iterate over data blocks and add forcing components
!
          call inject_eddy(xp(:), ap(:))

        end do

      end if

    end if ! forcing enabled

#ifdef PROFILE
! stop accounting time
!
    call stop_timer(ifu)
#endif /* PROFILE */

!-------------------------------------------------------------------------------
!
  end subroutine update_forcing
!
!===============================================================================
!!
!!***  PRIVATE SUBROUTINES  ****************************************************
!!
!===============================================================================
!
!===============================================================================
!
! subroutine INJECT_EDDY:
! ----------------------
!
!   Subroutine injects one eddy to the domain at given position.
!
!   Arguments:
!
!     xp, ap - eddy position and amplitude vectors;
!
!===============================================================================
!
  subroutine inject_eddy(xp, ap)

! include external variables
!
    use blocks, only : block_data, list_data

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    real(kind=8), dimension(3), intent(in) :: xp, ap

! local pointers
!
    type(block_data), pointer :: pdata
!
!-------------------------------------------------------------------------------
!
! assign pdata with the first block on the data block list
!
    pdata => list_data

! iterate over all data blocks
!
    do while (associated(pdata))

! inject eddy into the current block
!
      call inject_eddy_block(pdata, xp(:), ap(:))

! assign pdata to the next block
!
      pdata => pdata%next

    end do ! over data blocks

!-------------------------------------------------------------------------------
!
  end subroutine inject_eddy
!
!===============================================================================
!
! subroutine INJECT_EDDY_BLOCK:
! ----------------------------
!
!   Subroutine injects one eddy to the local block.
!
!   Arguments:
!
!     xp, ap - eddy position and amplitude vectors;
!
!===============================================================================
!
  subroutine inject_eddy_block(pdata, xp, ap)

! include external variables
!
    use blocks     , only : block_data
    use coordinates, only : im, jm, km
    use coordinates, only : ax, ay, az
    use coordinates, only : xlen, ylen, zlen
    use equations  , only : nv
    use equations  , only : idn, imx, imy, imz, ien
    use mpitools   , only : periodic

! local variables are not implicit by default
!
    implicit none

! subroutine arguments
!
    type(block_data), pointer , intent(inout) :: pdata
    real(kind=8), dimension(3), intent(in)    :: xp, ap

! local variables
!
    integer      :: i, j, k
    real(kind=8) :: fx, fy, fz, fp, ek

! local arrays
!
    real(kind=8), dimension(im)    :: x
    real(kind=8), dimension(jm)    :: y
    real(kind=8), dimension(km)    :: z
!
!-------------------------------------------------------------------------------
!
! prepare block coordinates
!
    if (periodic(1)) then
      fx      = 0.5d+00 * xlen
      x(1:im) = pdata%meta%xmin + ax(pdata%meta%level,1:im) - xp(1)
      x(1:im) = abs(x(1:im) / fx)
      x(1:im) = (abs(1.0d+00 - abs(1.0d+00 - abs(x(1:im))))) * fx / del
    else
      x(1:im) = (pdata%meta%xmin + ax(pdata%meta%level,1:im) - xp(1)) / del
    end if
    if (periodic(2)) then
      fy      = 0.5d+00 * ylen
      y(1:jm) = pdata%meta%ymin + ay(pdata%meta%level,1:jm) - xp(2)
      y(1:jm) = abs(y(1:jm) / fy)
      y(1:jm) = (abs(1.0d+00 - abs(1.0d+00 - abs(y(1:jm))))) * fy / del
    else
      y(1:jm) = (pdata%meta%ymin + ay(pdata%meta%level,1:jm) - xp(2)) / del
    end if
#if NDIMS == 3
    if (periodic(3)) then
      fz      = 0.5d+00 * zlen
      z(1:km) = pdata%meta%zmin + az(pdata%meta%level,1:km) - xp(3)
      z(1:km) = abs(z(1:km) / fz)
      z(1:km) = (abs(1.0d+00 - abs(1.0d+00 - abs(z(1:km))))) * fz / del
    else
      z(1:km) = (pdata%meta%zmin + az(pdata%meta%level,1:km) - xp(3)) / del
    end if
#else /* NDIMS == 3 */
    z(1:km) = 0.0d+00
#endif /* NDIMS == 3 */

! iterate over the block coordinates
!
    if (ien > 0) then
      do k = 1, km
        do j = 1, jm
          do i = 1, im

            fp = pdata%u(idn,i,j,k) * exp(- 0.5d+00 * (x(i)**2 + y(j)**2 + z(k)**2))

#if NDIMS == 3
            fx = (- ap(3) * y(j) + ap(2) * z(k)) * fp
            fy = (- ap(1) * z(k) + ap(3) * x(i)) * fp
            fz = (- ap(2) * x(i) + ap(1) * y(j)) * fp
#else /* NDIMS == 3 */
            fx = (- ap(3) * y(j)) * fp
            fy = (+ ap(3) * x(i)) * fp
            fz = 0.0d+00
#endif /* NDIMS == 3 */

            ek = 0.5d+00 * (pdata%u(imx,i,j,k)**2 + pdata%u(imy,i,j,k)**2 + pdata%u(imz,i,j,k)**2) / pdata%u(idn,i,j,k)

            pdata%u(imx,i,j,k) = pdata%u(imx,i,j,k) + fx
            pdata%u(imy,i,j,k) = pdata%u(imy,i,j,k) + fy
            pdata%u(imz,i,j,k) = pdata%u(imz,i,j,k) + fz

            ek = 0.5d+00 * (pdata%u(imx,i,j,k)**2 + pdata%u(imy,i,j,k)**2 + pdata%u(imz,i,j,k)**2) / pdata%u(idn,i,j,k) - ek

            pdata%u(ien,i,j,k) = pdata%u(ien,i,j,k) + ek

          end do
        end do
      end do
    else
      do k = 1, km
        do j = 1, jm
          do i = 1, im

            fp = pdata%u(idn,i,j,k) * exp(- 0.5d+00 * (x(i)**2 + y(j)**2 + z(k)**2))

#if NDIMS == 3
            fx = (- ap(3) * y(j) + ap(2) * z(k)) * fp
            fy = (- ap(1) * z(k) + ap(3) * x(i)) * fp
            fz = (- ap(2) * x(i) + ap(1) * y(j)) * fp
#else /* NDIMS == 3 */
            fx = (- ap(3) * y(j)) * fp
            fy = (+ ap(3) * x(i)) * fp
            fz = 0.0d+00
#endif /* NDIMS == 3 */

            pdata%u(imx,i,j,k) = pdata%u(imx,i,j,k) + fx
            pdata%u(imy,i,j,k) = pdata%u(imy,i,j,k) + fy
            pdata%u(imz,i,j,k) = pdata%u(imz,i,j,k) + fz

          end do
        end do
      end do
    end if

!-------------------------------------------------------------------------------
!
  end subroutine inject_eddy_block

!===============================================================================
!
end module forcing