! This file is part of GTC version 4.4 
! GTC version 4.4 is released under the 3-Clause BSD license:

! Copyright (c) 2002,2010,2016, GTC Team (team leader: Zhihong Lin, zhihongl@uci.edu)
! All rights reserved.

! 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 GTC Team nor the names of its contributors may be 
!    used to endorse or promote products derived from this software without 
!    specific prior written permission.

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Radial Boundary Condition Treatment
!
! Initialized by Sam Taimourzadeh, polished by Lei Shi.
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

!> \file
!> Radial Boundary Condition Treatment


!> \brief Radial Boundary Decay functions
module boundary_decay
  use utility, only: check_allocation, mstat
  use precision, only: lk
  use global_parameters, only: BC_type, BC_drive, BC_field, BC_weight, &
   BC_RMP, BC_nbound, BC_nboundR, mpsi, gtcout,mype, rg0
  use equilibrium, only: lsp
  implicit none


  private
  public BC_initial, radial_bound, field_bound, weight_bound, decay_coeff,&
    decay_coeff_sp

  ! Gaussian decay scale parameter
  ! The parameter specifies how many Half Width at Half Maximum(HWHM)
  ! corresponds to the given boundary grids.
  real(lk), parameter:: hwhm_number=4.0_lk

  ! Gaussian decay parameters for field and drive
  real(lk) sigma_inv, sigmaR_inv

  ! Calculated decay coefficients on radial mesh
  real(lk), dimension(:), allocatable:: decay_coeff, decay_coeff_sp
  !$acc declare create(decay_coeff, decay_coeff_sp)

  save
  contains

    !> \brief Initialize the BC module, calculate the key parameters

    subroutine BC_initial
      use equilibrium, only: spdpsi
      use spline_function, only: sprgpsi
      use field_array, only: deltar
      integer:: i, idx
      real(lk):: rg, wr0, deltar_inv

      if (mype==0) write(gtcout,*) "==============================="

      ! no initialization required if BC is turned off
      if(BC_type==0) then
        if(mype == 0) write(gtcout,*) "No Radial Boundary Decay."
        BC_drive=0
        BC_field=0
        BC_weight=0
        BC_RMP=0
        return
      endif

      ! Set default BC appliable quantities 
      !> @note
      !> Programmer needs to set these paramters **MANUALLY IN SOURCE CODE**
      !> to control what quantities are subjected to the Radial Boundary
      !> decay.\n
      !> The Default values are BC_drive=1, BC_field=1, BC_weight=0,BC_RMP=1, since
      !> modification to the weights is more dangerous. Programmer should reset
      !> to default before **Committing and pushing back to Git Repository**
      !> after personal modifications.
      BC_drive=1
      BC_field=1
      BC_weight=0
      BC_RMP=1

      ! Write out Boundary Condition information
      if(mype == 0) then
        write(gtcout,*) "Radial Boundary Decay initialization."
        if (BC_type == 1) then
          write(gtcout,*) "Radial Boundary type: Linear"
        elseif (BC_type == 2) then
          write(gtcout,*) "Radial Boundary type: Gaussian"
        else
          write(gtcout,*) "WARNING: UNKNOWN BOUNDARY TYPE CODE, BC_type=",&
            BC_type
        endif

        write(gtcout,*) "Radial Boundary Decay applied to:"
        if(BC_drive == 1) then
          write(gtcout,*) "Equilibrium Gradient Drives: YES"
        else
          write(gtcout,*) "Equilibrium Gradient Drives: NO"
        endif

        if(BC_field == 1) then
          write(gtcout,*) "Perturbed Field Solutions: YES"
        else
          write(gtcout,*) "Perturbed Field Solutions: NO"
        endif

        if(BC_weight == 1) then
          write(gtcout,*) "Perturbed Particle Weights: YES"
        else
          write(gtcout,*) "Perturbed Particle Weights: NO"
        endif

        if(BC_RMP == 1) then
          write(gtcout,*) "Equilibrium RMP Fields: YES"
        else
          write(gtcout,*) "Equilibrium RMP Fields: NO"
        endif
      endif

      ! set symmetric boundary layer if BC_nboundR==0
      if(BC_nboundR==0) then
        if (mype==0) write(gtcout,*) "Symmetric Radial BC range:", BC_nbound
        BC_nboundR = BC_nbound
      else
        if (mype==0) write(gtcout,*) "Asymmetric Radial BC range: inner", &
          BC_nbound, "outer", BC_nboundR
      endif

      ! BC information output finishes
      if (mype==0) write(gtcout,*) &
        "==================================================================="

      ! allocate decay coefficient array with the same indexing convention as
      ! other radial field array: index starts from 0 to mpsi
      allocate(decay_coeff(0:mpsi), STAT=mstat)
      call check_allocation(mstat, "decay_coeff in BC_function")

      ! allocate decay coefficients on spline grid if RMP decay is on
      if (BC_RMP==1) then
        allocate(decay_coeff_sp(lsp), STAT=mstat)
        call check_allocation(mstat, "decay_coeff_sp in BC_function")
      endif

      ! Initialize module variables based on BC_type
      if(BC_type==1) then
        ! calculate the decay coefficients on radial mesh
        decay_coeff(:) = 1.0_lk
        do i=0,BC_nbound
          decay_coeff(i)=real(i,lk)/max(1.0_lk,real(BC_nbound,lk))
        enddo
        do i=0,BC_nboundR
          decay_coeff(mpsi-i)=real(i,lk)/max(1.0_lk,real(BC_nboundR,lk))
        enddo
      else
        ! Gaussian BC
        ! Field and drive are decaying on radial grid
        sigma_inv = (hwhm_number*sqrt(2.0_lk*log(2.0_lk)))/real(BC_nbound,lk)
        sigmaR_inv =(hwhm_number*sqrt(2.0_lk*log(2.0_lk)))/real(BC_nboundR,lk)

        ! calculate the decay coefficients on radial mesh
        decay_coeff(:) = 1.0_lk
        do i=0,BC_nbound
          decay_coeff(i)=exp(-(((i-BC_nbound)*sigma_inv)**2)*0.5_lk)
        enddo
        do i=0,BC_nboundR
          decay_coeff(mpsi-i)=exp(-(((BC_nboundR-i)*sigmaR_inv)**2)*0.5_lk)
        enddo
      endif
      
      ! Now generate the decay coefficients on the equilibrium spline mesh
      ! based on the decay coefficients calculated on the simulation mesh
      ! This is only used for RMP boundary decay
      if(BC_RMP == 1) then
        deltar_inv = 1.0_lk/deltar
        !$omp parallel do private(i,idx,rg,wr0)
        do i=1,lsp
          rg = sprgpsi(real(i-1,lk)*spdpsi)
          ! radial grid index on outer flux surface
          idx=ceiling((rg-rg0)*deltar_inv)
          ! screen out of range points
          if(idx < 1 .or. idx > mpsi) then
            decay_coeff_sp(i) = 0
            continue
          else
            ! linear interpolation weight on inner surface,
            ! depends on the distance in rg from particle to outer grid
            wr0 = real(idx,lk) - (rg-rg0)*deltar_inv

            ! apply decay
            decay_coeff_sp(i) = ( wr0*decay_coeff(idx-1) + (1-wr0)*decay_coeff(idx) )
          endif
        enddo
     endif
     !$acc update device(decay_coeff, decay_coeff_sp)
    end subroutine BC_initial


    !> \brief Apply 1D radial boundary decay on the given drive array
    !
    !> BC_type determines which kind of decay function is applied. The radial
    !boundary widths are given by BC_nbound, and BC_nboundR
    !> @param[inout] field 
    !> A 1D radial array containing data needed to be suppressed near both boundareis.
    !> The decay is applied IN PLACE.

    subroutine radial_bound(field)
      real(lk) field(0:mpsi)
      integer i
      ! no decay needed for BC_type==0
      if(BC_type == 0) return
      ! apply decay coefficients on drive
      !$omp parallel do private(i)
      do i=0,mpsi
        field(i) = field(i)*decay_coeff(i)
      enddo
    end subroutine radial_bound 

    !> \brief Apply radial boundary decay on the given field array
    !
    !> @param[inout] field
    !> A 1D array arranged as the GTC perturbation mesh, containing data needed to be suppressed near both radial boundareis.
    !> The decay is applied IN PLACE.
    !
    !> @see fieldinitial, igrid, mtheta
    subroutine field_bound(field)
      use global_parameters, only: mgrid
      use field_array, only: mtheta, igrid
      real(lk) field(mgrid)
      integer i,j

      ! no decay needed for BC_type==0
      if(BC_type == 0) return
      ! apply decay coefficients on drive_array
#ifdef _OPENACC
      !$acc parallel loop gang
#else
      !$omp parallel do private(i,j)
#endif
      do i=0,mpsi
        !$acc loop vector
        do j=igrid(i),igrid(i)+mtheta(i)
          field(j) = field(j)*decay_coeff(i)
        enddo
      enddo
    endsubroutine field_bound

    !> \brief Decay function on particle weights
    !
    !> Unlike the field quantities, particles are NOT on radial mesh. We need
    !> to calculate the local decay coefficient for each particle. Since we
    !> already have an pretty fine array of coefficients on the radial mesh, a
    !> linear interpolation is accurate enough.
    !
    !> @param [in,out] zpart
    !> The simulation particle array ITSELF. The weights are modified IN PLACE.
    !> @param [in] mp
    !> number of particle in current MPI
    !
    !> \see zion, mi, PushParticle
    !> @todo
    !> Particle loop included, optimize with GPU!
    subroutine weight_bound(zpart, mp)
      use spline_function, only: sprgpsi
      use field_array, only: deltar
      integer mp, nparamp
      real(lk),dimension(:,:) :: zpart
      integer m,idx
      real(lk) rg, wr0, deltar_inv

      deltar_inv = 1.0_lk/deltar
      !$omp parallel do private(m,idx,rg,wr0)
      do m=1,mp
        rg = sprgpsi(zpart(1,m))
        ! radial grid index on outer flux surface
        idx=max(1,min(mpsi,ceiling((rg-rg0)*deltar_inv)))
        ! linear interpolation weight on inner surface,
        ! depends on the distance in rg from particle to outer grid
        wr0 = real(idx,lk) - (rg-rg0)*deltar_inv

        ! apply decay
        zpart(5,m) = zpart(5,m)*&
          ( wr0*decay_coeff(idx-1) + (1-wr0)*decay_coeff(idx) )
      enddo
    end subroutine weight_bound
end module boundary_decay


