! 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.

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Treatment of antenna to drive modes in place of fast particles
! external anntena standing wave cos(m*theta-n*zeta)*cos(omega*t)
! Implemented by Valentin Aslanyan, Lei Shi, Sam Taimourzadeh
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!


! Every routine assumes antenna_on>0 condition is satisfied
module antenna
  use system_env, only: ierror, MPI_INTEGER, MPI_COMM_WORLD 
  use utility, only: check_allocation, mstat
  use precision, only: lk, mpi_Rsize
  use global_parameters,only: mgrid, mpsi, istep, irk, pi2, mstep, mype, gtcout, utime, tstep, irun, fielddir
  use field_array, only:  mtheta, zeta0, zeta1, igrid, deltat, qtinv
!#TODO - external variables

  implicit none

  private
  public particle_initialization, field_initialization, phi_ext, grad_phi_ext, antenna_setup, antenna_advance_tstep

  !> Maximum number of antenna toroidal/poloidal harmonics that can be set
  integer, parameter :: antenna_max_modes=8



  !> 1: antenna creates a delta_phi perturbation; 2: antenna creates a delta_Aparallel perturbation
  integer :: antenna_type
  !> actual number of antenna toroidal/poloidal harmonics used in calculation
  integer :: antenna_num_modes
  !> antenna toroidal mode number - set mode to negative to disable harmonic
  integer :: antenna_n(antenna_max_modes)
  !> antenna poloidal mode number
  integer :: antenna_m(antenna_max_modes)
  !> Location on GTC grid where the antenna signal peaks
  integer :: antenna_center
  !> Distance (in points on GTC grid) to where the antenna signal goes to zero
  integer :: antenna_boundary
  !> 0: the initial perturbations go to zero (only antenna drive leads to modes)
  integer :: particle_initialization
  integer :: field_initialization

  !> frequency of antenna [Hz]
  real(lk) :: antenna_freq

  !> Angular frequency of antenna [rad s^-1]
  real(lk) :: antenna_omega
  !> Gaussian width of the antenna signal - if negative defaults to recommended "optimal"
  real(lk) :: antenna_sigma
  !> Peak amplitude at antenna_center location
  real(lk) :: antenna_amplitude

  !> Gaussian envelope for the sinusoidal antenna perturbations
  real(lk),dimension(:),allocatable :: envelope

  !> antenna delta_phi - in e phi/T_e unit
  real(lk),dimension(:,:),allocatable :: phi_ext
  !$acc declare create(phi_ext)
  !> Gradient of antenna delta_phi - in e phi/T_e unit
  real(lk),dimension(:,:,:),allocatable :: grad_phi_ext
  !$acc declare create(grad_phi_ext)
  save
  contains

!=============================================================================
  Subroutine antenna_setup
!=============================================================================

!> constant downshift of the gaussian envelope, so that it is continuous at the
!> cutoff location 
    real(lk) :: gauss_downshift
    integer :: i

    call antenna_read

    ! convert omega to GTC unit
    antenna_omega=antenna_omega*utime
    if(antenna_sigma<0.0_lk)then
       antenna_sigma=0.5_lk*real(antenna_boundary,lk)
    endif

    call broadcast_antenna_params

    allocate( envelope(0:mpsi), STAT=mstat )
    call check_allocation(mstat, "envelope in antenna.F90")

    gauss_downshift = exp(-0.5_lk*(real(antenna_boundary,lk)/antenna_sigma)**2)

!~~~~~~~~~~~Gaussian envelope~~~~~~~~~~~~~~~~~~~~~~
    do i=0,mpsi
      if ( i>(antenna_center-antenna_boundary) .and. i<(antenna_center+antenna_boundary)) then
        envelope(i)=antenna_amplitude/(1.0_lk-gauss_downshift)*( exp(-0.5_lk*( (real(i,lk)-real(antenna_center,lk))/antenna_sigma )**2)&
          - gauss_downshift)
	    else
	      envelope(i)=0.0_lk
	    endif
    enddo
!~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

     grad_phi_ext=0.0_lk

  end Subroutine antenna_setup

!=============================================================================
!> time advance the antenna induced phi_ext and grad_phi_ext arrays. 
!> These arrays are later used in either poisson or pushfield to add external
!> drive to the self-consistent phi or Apara
  Subroutine antenna_advance_tstep
!=============================================================================

     integer :: i,j,ij,idx_antmode
     real(lk) :: tdum0, tdum1, wt
     real(lk) :: buf(mgrid)
     !$acc declare create(buf)
     
     phi_ext=0.0_lk

     if(mype==0 .and. istep==1 .and. irk==1)then
        write(gtcout,*)'Antenna freq=',antenna_omega/(pi2*utime),'Hz'
        flush(gtcout)
     endif

! phase of driving source
     wt=cos(antenna_omega*tstep*(real(istep+mstep*irun,lk)+0.5_lk*irk)) !irun = # of restart runs


     do i=0,mpsi
!$omp parallel do private(j,ij,tdum0,tdum1)
        do j=1,mtheta(i)
           ij=igrid(i)+j
           if(fielddir==1 .or. fielddir==3)then
             tdum0=deltat(i)*real(j,lk)+(zeta0-pi2)*qtinv(i)
             tdum1=deltat(i)*real(j,lk)+(zeta1-pi2)*qtinv(i)
           else
             tdum0=deltat(i)*real(j,lk)+zeta0*qtinv(i)
             tdum1=deltat(i)*real(j,lk)+zeta1*qtinv(i)
           endif
           do idx_antmode=1,antenna_num_modes
             phi_ext(0,ij)=phi_ext(0,ij)+envelope(i)*cos(real(antenna_m(idx_antmode),lk)*tdum0-real(antenna_n(idx_antmode),lk)*zeta0)*wt
             phi_ext(1,ij)=phi_ext(1,ij)+envelope(i)*cos(real(antenna_m(idx_antmode),lk)*tdum1-real(antenna_n(idx_antmode),lk)*zeta1)*wt
           enddo
        enddo
     enddo

! write out driving source
     if(mype==0 .and. istep==1 .and. irk==1)then
        open(9991,file='time_ext.out',status='replace')
        open(9992,file='phi_ext.out',status='replace')
     endif
     if(mype==0 .and. irk==1)then
        write(9991,*)wt
        write(9992,*)phi_ext(1,igrid(mpsi/2)+1)
     endif
     if(mype==0 .and. istep==mstep .and. irk==1)then
        close(9991)
        close(9992)
     endif

! delta_phi antenna version - do nothing at this point
! delta_Aparallel antenna version - calculate gradient for Aparallel, but reset phi_ext
     if(antenna_type==2)then
        buf=phi_ext(1,:)
        !$acc update device(buf)
        call gradient_gpu(buf,grad_phi_ext)
        !$acc update host(grad_phi_ext)
        phi_ext=0.0_lk
     endif

  end Subroutine antenna_advance_tstep

!=============================================================================
!> read antenna.in file
  Subroutine antenna_read
!=============================================================================
  use system_env, only: GTC_ABORT
  use global_parameters, only: magnetic

  integer :: idx_antmode
  logical :: file_exist

  namelist /antenna_parameters/ antenna_type,antenna_n,antenna_m,&
      antenna_freq,antenna_amplitude,&
      antenna_center,antenna_sigma,antenna_boundary,&
      particle_initialization,field_initialization

    ! default values of the antenna parameters
    antenna_type=1 		! 1: antenna creates a delta_phi perturbation; 2: antenna creates a delta_Aparallel perturbation
    antenna_num_modes=1         ! actual number of toroidal/poloidal antenna harmonics used
    antenna_n=(/5,-5,-5,-5,-5,-5,-5,-5/)	! antenna toroidal mode numbers - set to negative to not use a given mode
    antenna_m=(/7,8,9,10,11,12,13,14/)		! antenna poloidal mode numbers
    antenna_freq=120000.0_lk  	! antenna frequency [Hz]
    antenna_amplitude=1e-7_lk  	! 
    antenna_center=50     	! GTC grid midpoint of antenna profile
    antenna_sigma=-1.0_lk     	! Gaussian sigma in GTC grid units
				! Defaults to predefined (antenna_boundary/2) if negative
    antenna_boundary=10      	! Radial size of the antenna signal (the number of psi flux points)
    particle_initialization=0	! 0 to set all initial perturbation to zero
    field_initialization=0     	! 


    antenna_num_modes=0         ! actual number of toroidal/poloidal antenna harmonics used

    do idx_antmode=1,antenna_max_modes
      if(antenna_m(idx_antmode)>=0)then
        antenna_num_modes=antenna_num_modes+1
      endif
    enddo


! Test if the antenna input file antenna.in exists
    inquire(file='antenna.in',exist=file_exist)
    if (file_exist) then
      open(550,file='antenna.in',status='old')
      read(550,nml=antenna_parameters,iostat=ierror)
      if(ierror > 0)write(*,*)'WARNING: check input_parameters'
      close(550)
    else
      write(*,*)'Cannot find file antenna.in !!!'
      call GTC_ABORT
    endif

    ! convert to angular frequency
    antenna_omega = antenna_freq * pi2 

    if(mype==0) then
      write(gtcout,*) "===================================================="
      write(gtcout,*) " External Antenna is used. "
      write(gtcout,*) " See antenna.in for details."
      if(magnetic == 0 .and. antenna_type == 2) then
        write(gtcout,*) "WARNING: Antenna type set to Apara in electrostatic",&
        " simulation does NOT have effect. Antenna type is converted to phi.",&
        " If this is not desired, please update gtc.in and/or antenna.in."
        antenna_type = 1
      endif
      write(gtcout,antenna_parameters)
      write(gtcout,*) "===================================================="
    endif
  end Subroutine antenna_read


!=============================================================================
!> broadcast the read-in antenna parameters to all the MPIs
  Subroutine broadcast_antenna_params
!=============================================================================

  implicit none

  integer,parameter :: n_integers=6+antenna_max_modes*2, n_reals=3
  integer  :: integer_params(n_integers)
  real(lk) :: real_params(n_reals)


  if(mype==0)then
    integer_params(1)=antenna_type
    integer_params(2)=antenna_center
    integer_params(3)=antenna_boundary
    integer_params(4)=particle_initialization
    integer_params(5)=field_initialization
    integer_params(6)=antenna_num_modes
    integer_params(7:6+antenna_max_modes)=antenna_n(1:antenna_max_modes)
    integer_params(7+antenna_max_modes:6+antenna_max_modes*2)=antenna_m(1:antenna_max_modes)


    real_params(1)=antenna_omega
    real_params(2)=antenna_amplitude
    real_params(3)=antenna_sigma
  endif

! Send input parameters to all processes
  call MPI_BCAST(integer_params,n_integers,MPI_INTEGER,0,MPI_COMM_WORLD,ierror)
  call MPI_BCAST(real_params,n_reals,mpi_Rsize,0,MPI_COMM_WORLD,ierror)

  if(mype/=0)then
    antenna_type=integer_params(1)
    antenna_center=integer_params(2)
    antenna_boundary=integer_params(3)
    particle_initialization=integer_params(4)
    field_initialization=integer_params(5)
    antenna_num_modes=integer_params(6)
    antenna_n(1:antenna_max_modes)=integer_params(7:6+antenna_max_modes)
    antenna_m(1:antenna_max_modes)=integer_params(7+antenna_max_modes:6+antenna_max_modes*2)

    antenna_omega=real_params(1)
    antenna_amplitude=real_params(2)
    antenna_sigma=real_params(3)
  endif

  end Subroutine broadcast_antenna_params


end module antenna
