! This collection of routines provide an API to set up portions of
! the data in "get_fbm_module" -- enough to enable fbm_wnetcdf_r4 or _fp
! to be called...

! all floating point arguments are real(fp) ::

subroutine gfbm_set_species(nin,zin,ain)
  ! set fast species list
  use iso_c_binding, only: fp => c_double
  use get_fbm_module
  implicit none

  integer, intent(in) :: nin     ! number of species
  real(fp), intent(in) :: zin(nin) ! Z of species: 1 for H, 2 for He, ...
  real(fp), intent(in) :: ain(nin) ! A of species: ~1 for H, ~2 for D, ...

  !----------------

  nsfast = nin
  nfast = nsfast

  if(allocated(abeams)) deallocate(abeams)
  if(allocated(xzbeams)) deallocate(xzbeams)

  allocate(abeams(nsfast),xzbeams(nsfast))

  abeams = ain
  xzbeams = zin

  ns = 0  ! for now...

  if(allocated(nznbmea)) deallocate(nznbmea)
  allocate(nznbmea(nsfast)); nznbmea = 0  ! for now...

  if(allocated(ktype)) deallocate(ktype)
  allocate(ktype(nsfast)); ktype = 0      ! for now...

  if(allocated(nldep_beam)) deallocate(nldep_beam)
  allocate(nldep_beam(nsfast)); nldep_beam = .FALSE.

  if(allocated(nldep_fusion)) deallocate(nldep_fusion)
  allocate(nldep_fusion(nsfast)); nldep_fusion = .FALSE.

  if(allocated(nlfprod)) deallocate(nlfprod)
  allocate(nlfprod(nsfast)); nlfprod = .FALSE.

end subroutine gfbm_set_species

subroutine gfbm_set_timestep(zta,zdta)
  use iso_c_binding, only: fp => c_double

  ! set time step info
  use get_fbm_module
  implicit none

  real(fp), intent(in) :: zta     ! time of data (s)
  real(fp), intent(in) :: zdta    ! +/- time averaging interval of data (s)

  !----------------

  ta = zta
  dta = zdta

end subroutine gfbm_set_timestep

subroutine gfbm_set_tokamak(ztok)
  use iso_c_binding, only: fp => c_double

  ! set tokamak label -- 3 or 4 character ID, e.g. "TFTR", "NSTX", "JET", ...

  use get_fbm_module
  implicit none

  character(len=*), intent(in) :: ztok

  !----------------
  integer :: ilen
  !----------------

  ilen = min(len(ztok),len(tdev))

  tdev = ztok(1:ilen)

end subroutine gfbm_set_tokamak

subroutine gfbm_set_shot(ishot)
  use iso_c_binding, only: fp => c_double

  ! set shot number (integer :: experimental shot number)

  use get_fbm_module
  implicit none

  integer, intent(in) :: ishot

  !----------------

  nshot = ishot

end subroutine gfbm_set_shot

subroutine gfbm_set_dims(nin,nzns_e,nzns_vpv,nzrows,nztot)
  use iso_c_binding, only: fp => c_double

  ! set dimensions sufficient to declare a distribution function array
  ! CALL GFBM_SET_SPECIES prior to making this call; nin must match the
  ! number of species.

  use get_fbm_module
  implicit none

  integer, intent(in) :: nin       ! number of species
  integer, intent(in) :: nzns_e(nin)  ! number of energy bins, each specie
  integer, intent(in) :: nzns_vpv  ! number of vpll/v bins
  integer, intent(in) :: nzrows    ! number of radial zone rows
  integer, intent(in) :: nztot     ! total number of zones, 2d irregular grid

  !----------------------
  integer :: ii,inume
  !----------------------

  if(nin.ne.nsfast) then

     write(0,*) ' ?gfbm_set_dims: species dimension inconsistency:'
     write(0,*) '  passed variable     nin = ',nin
     write(0,*) '  module variable  nsfast = ',nsfast
     write(0,*) '  These must match.  Prior call to GFBM_SET_SPECIES required.'
     call bad_exit

  end if

  nznbmea(1:nsfast) = nzns_e(1:nsfast)
  nznbma  = nzns_vpv

  nznbmrf = nzrows
  nznbmr  = nznbmrf

  nvol    = nztot
  nznsi   = nztot

  lcentr = 2

  ! XI: normalized radial sqrt(tor.flux) grid: allocation & values

  if(allocated(xi)) deallocate(xi)
  if(allocated(xib)) deallocate(xib)
  allocate(xi(nznbmrf),xib(nznbmrf+1))

  xib(1)= 0.0_fp
  xib(nznbmrf+1)=1.0_fp
  do ii=2,nznbmrf
     xib(ii) = (ii-1)*1.0_fp/(nznbmrf)
  end do
  do ii=1,nznbmrf
     xi(ii) = 0.5_fp*(xib(ii)+xib(ii+1))
  end do

  ! XXKSID: vpll/v grid: allocation & values

  if(allocated(xxksid)) deallocate(xxksid)
  allocate(xxksid(nznbma))
  do ii=1,nznbma
     xxksid(ii) = -1.0_fp + 2.0_fp*(ii-0.5_fp)/(1.0_fp*nznbma)
  end do

  ! EFBMB: allocate only

  if(allocated(efbmb)) deallocate(efbmb)

  inume=maxval(nznbmea)
  allocate(efbmb(inume+1,nsfast))
  efbmb= 0.0_fp  ! for now...

end subroutine gfbm_set_dims

subroutine gfbm_set_trgrid(inrho,rho_in,inmom)
  use iso_c_binding, only: fp => c_double

  ! set dimensions associated with transport code: radial grid;
  ! number of moments in FFT representation of MHD equilibrium

  ! for xitr(...) the values are set from the input data

  ! for other arrays that are to be dimensioned accordingly:
  ! this routine allocates associated arrays and clears them to zero

  use get_fbm_module
  implicit none

  integer, intent(in) :: inrho  ! #bdys in transport grid (nzones + 1)
  real(fp), intent(in) :: rho_in(inrho)  ! the radial flux coord grid.
  integer, intent(in) :: inmom  ! #of FFT moments beyond 0'th moment
  ! ...MHD equilibrium flux surface geometry expansion

  !---------------------------------
  integer :: i
  !---------------------------------

  !  dimension variables:

  mj = inrho + 2
  lcentr = 2
  nzones = inrho - 1

  mimom = inmom
  nmom = inmom

  miz = 2

  !  radial grid "xitr":

  if(allocated(xitr)) deallocate(xitr)
  allocate(xitr(mj,miz))

  xitr(lcentr:lcentr+nzones,1) = rho_in(1:inrho)
  xitr(1,1) = 2*rho_in(1)-rho_in(2)
  xitr(lcentr+nzones+1,1) = 2*rho_in(inrho)-rho_in(inrho-1)
  do i=1,lcentr+nzones
     xitr(i,2)=(xitr(i,1)+xitr(i+1,1))/2
  end do
  xitr(lcentr+nzones+1,2) = 2*xitr(lcentr+nzones,2) - xitr(lcentr+nzones-1,2)

  !  associated arrays:

  if(allocated(omegag)) deallocate(omegag)
  allocate(omegag(mj)); omegag = 0.0_fp

  if(allocated(rmom0)) deallocate(rmom0)
  if(allocated(rmom)) deallocate(rmom)
  if(allocated(ymom)) deallocate(ymom)
  if(allocated(rmc)) deallocate(rmc)
  if(allocated(ymc)) deallocate(ymc)

  allocate(rmom(mimom,mj,miz),ymom(mimom,mj,miz),rmom0(mj,miz))
  allocate(rmc(mj,0:mimom,2),ymc(mj,0:mimom,2))

  rmc = 0.0_fp
  ymc = 0.0_fp

  rmom0= 0.0_fp
  rmom = 0.0_fp
  ymom = 0.0_fp

end subroutine gfbm_set_trgrid

subroutine gfbm_set_omega(inum,omega_in)
  use iso_c_binding, only: fp => c_double

  ! set the angular velocity profile
  ! input profile is expected to have   nzones   zone-centered values.
  ! matching what is set in "gfbm_set_trgrid"

  use get_fbm_module
  implicit none

  integer, intent(in) :: inum          ! number of zones in input profile
  real(fp), intent(in) :: omega_in(inum) ! angular velocity profile @ zone ctrs

  !-----------------
  integer :: ii
  !-----------------

  if(inum.ne.nzones) then

     write(0,*) ' ?gfbm_set_omega: inum = ',inum
     write(0,*) '  in get_fbm_module: nzones = ',nzones
     write(0,*) '  These need to match: call gfbm_set_trgrid first.'
     call bad_exit

  end if

  ! zone ctr -> zone bdy shift

  omegag(lcentr+nzones:)=0.0_fp
  omegag(1:lcentr) = (9.0_fp*omega_in(1) - omega_in(2))/8.0_fp
  do ii=2,nzones
     omegag(lcentr-1+ii) = (omega_in(ii-1)+omega_in(ii))/2.0_fp
  end do

end subroutine gfbm_set_omega

subroutine gfbm_set_udsym(symflag)
  use iso_c_binding, only: fp => c_double
  use physconst_mod, only: pi
  
  ! set variables to indicated up-down symmetric geometry (or not)...

  use get_fbm_module
  implicit none

  logical, intent(in) :: symflag

  !---------------

  nlsym = symflag

  if(nlsym) then
     thbdy0 = 0.0_fp
     thbdy1 = pi

  else
     thbdy0 = -pi
     thbdy1 = -thbdy0
  end if

end subroutine gfbm_set_udsym

subroutine gfbm_set_fgc(gcflag)
  use iso_c_binding, only: fp => c_double

  ! set variables to indicated guiding center (GC) distribution
  ! (as opposed to a particle distribution not at the GC).

  use get_fbm_module
  implicit none

  logical, intent(in) :: gcflag

  !---------------

  nlfbm_gc = gcflag

end subroutine gfbm_set_fgc

subroutine gfbm_set_egridb1(ispecies,inumeb,zebdys)
  use iso_c_binding, only: fp => c_double

  ! set energy grid boundaries for a single species; fill in zone values

  use get_fbm_module
  implicit none

  integer, intent(in) :: ispecies     ! fast ion specie index, btw 1 and nsfast
  integer, intent(in) :: inumeb       ! number of energy grid boundaries
  !  must match specie #energy zones + 1

  real(fp), intent(in) :: zebdys(inumeb)  ! energy grid in units of eV
  !  zone ii energy bin: zebdys(ii) to zebdys(ii+1)

  !------------------------
  integer :: ii,inume
  !------------------------

  if((ispecies.le.0).or.(ispecies.gt.nsfast)) then
     write(0,*) ' ?gfbm_set_egridb1: specie index out of range or nsfast not set:'
     write(0,*) '  #of fast species (nsfast): ',nsfast
     write(0,*) '  specie index (passed):     ',ispecies
     write(0,*) '  value between 1 and nsfast is expected.'
     if(nsfast.eq.0) write(0,*) '  ==> call gfbm_set_species(...) first!'

     call bad_exit
  end if

  inume=inumeb - 1

  if(.not.allocated(nznbmea)) then
     write(0,*) ' ?gfbm_set_egridb1: energy grid sizes NZNBMEA(...) not set.'
     write(0,*) '  call gfbm_set_species(...) first!'
     write(0,*) '  call gfbm_set_dims(...) first!'

     call bad_exit
  end if

  if(inume.ne.nznbmea(ispecies)) then
     write(0,*) ' '
     write(0,*) '  ispecies = ',ispecies
     write(0,*) ' ?gfbm_set_egridb1: #of energy grid bdys not consistent.'
     if(nznbmea(ispecies).eq.0) then
        write(0,*) '  nznbmea(ispecies)=0; call gfbm_set_dims(...) first!'
     else
        write(0,*) '  nznbmea(ispecies) = ',nznbmea(ispecies)
        write(0,*) '  inume = inumeb-1  = ',inume
        write(0,*) '  (these numbers should match).'
     end if

     call bad_exit
  end if

  efbmb(1:inumeb,ispecies) = zebdys

end subroutine gfbm_set_egridb1

subroutine gfbm_set_types(nin,ninfb,ilfprod)
  use iso_c_binding, only: fp => c_double

  ! set species types: NBI ion or fusion product ion

  use get_fbm_module
  implicit none

  integer, intent(in) :: nin     ! number of species
  integer, intent(in) :: ninfb   ! number of beam + fusion product species
  logical, intent(in) :: ilfprod(nin)  ! fusion product flags

  !------------------
  integer :: ii
  !------------------

  fitype(1)='NBI'
  fitype(2)='FUSN'
  fitype(3)='RFI'

  if(nin.ne.nsfast) then

     write(0,*) ' ?gfbm_set_types: species dimension inconsistency:'
     write(0,*) '  passed variable     nin = ',nin
     write(0,*) '  module variable  nsfast = ',nsfast
     write(0,*) '  These must match.  Prior call to GFBM_SET_SPECIES required.'
     call bad_exit

  end if

  nsbeam = ninfb
  nlfprod = .FALSE.

  do ii=1,nsfast
     if(ii.le.nsbeam) then
        if(ilfprod(ii)) then
           nlfprod(ii) = .TRUE.
           ktype(ii) = 2
        else
           ktype(ii) = 1
        end if
     else
        ktype(ii) = 3
     end if
  end do

end subroutine gfbm_set_types

subroutine gfbm_set_nthzns(inumr,inthzns)
  use iso_c_binding, only: fp => c_double

  ! set number of theta zones for each radial zone row

  use get_fbm_module
  implicit none

  integer, intent(in) :: inumr  ! number of radial zone rows: must match:
                                ! nznbmrf
  integer, intent(in) :: inthzns(inumr)  ! number of theta zones, each row

  !------------------
  integer :: ii
  !------------------

  if(inumr.ne.nznbmrf) then

     write(0,*) ' ?gfbm_set_nthzns: zone row dimension inconsistency:'
     write(0,*) '  passed variable      inumr = ',inumr
     write(0,*) '  module variable    nznbmrf = ',nznbmrf
     write(0,*) '  These must match.  Prior call to GFBM_SET_DIMS required.'
     call bad_exit

  end if

  if(allocated(nthzsm)) deallocate(nthzsm)
  allocate(nthzsm(inumr+1))

  lcentr=2
  nthzsm(1)=0
  do ii=1,inumr
     nthzsm(ii+1) = nthzsm(ii) + inthzns(ii)
  end do

  nvol=nthzsm(inumr+1)

end subroutine gfbm_set_nthzns

subroutine gfbm_set_bmvol(inum,zvol)
  use iso_c_binding, only: fp => c_double

  ! set 2d spatial grid zone volumes

  use get_fbm_module
  implicit none

  integer, intent(in) :: inum  ! number of volume elements: must match nvol!
  real(fp), intent(in) :: zvol(inum)  ! zone volumes, cm**3

  !------------------
  integer :: ii
  !------------------

  if(inum.ne.nvol) then

     write(0,*) ' ?gfbm_set_bmvol: zone volumes array size error:'
     write(0,*) '  passed variable  inum = ',inum
     write(0,*) '  module variable  nvol = ',nvol
     write(0,*) '  These must match.  Prior call to GFBM_SET_NTHZNS required.'
     call bad_exit

  end if

  if(allocated(bmvol)) deallocate(bmvol)
  allocate(bmvol(nvol))

  bmvol = zvol

end subroutine gfbm_set_bmvol

subroutine gfbm_set_bdens2(inum,nin,zdens2)
  use iso_c_binding, only: fp => c_double

  ! set beam density 2d profiles over MC grid
  ! note: #/zone, use bmvol to get actual densities

  use get_fbm_module
  implicit none

  integer, intent(in) :: inum   ! number of 2d grid zones: must match nvol!
  integer, intent(in) :: nin    ! number of species, must match nsfast!
  real(fp), intent(in) :: zdens2(inum,nin)  ! #/zone species ptcl counts
  !  (2d profile)

  !------------------

  if(inum.ne.nvol) then

     write(0,*) ' ?gfbm_set_bdens2: zone dimension size error:'
     write(0,*) '  passed variable  inum = ',inum
     write(0,*) '  module variable  nvol = ',nvol
     write(0,*) '  These must match.  Prior call to GFBM_SET_NTHZNS required.'
     call bad_exit

  end if

  if(nin.ne.nsfast) then

     write(0,*) ' ?gfbm_set_bdens2: species dimension inconsistency:'
     write(0,*) '  passed variable     nin = ',nin
     write(0,*) '  module variable  nsfast = ',nsfast
     write(0,*) '  These must match.  Prior call to GFBM_SET_SPECIES required.'
     call bad_exit

  end if

  if(allocated(bdens2)) deallocate(bdens2)
  allocate(bdens2(nvol,nsfast))

  bdens2 = zdens2

end subroutine gfbm_set_bdens2

subroutine gfbm_set_eba2(inum,nin,zeba2pp,zeba2pl)
  use iso_c_binding, only: fp => c_double

  ! set 2d profiles of <Eperp>, <Epll>, eV, over MC grid

  use get_fbm_module
  implicit none

  integer, intent(in) :: inum   ! number of 2d grid zones: must match nvol!
  integer, intent(in) :: nin    ! number of species, must match nsfast!
  real(fp), intent(in) :: zeba2pp(inum,nin)  ! <Eperp>, eV
  real(fp), intent(in) :: zeba2pl(inum,nin)  ! <Epll>, eV

  !------------------

  if(inum.ne.nvol) then

     write(0,*) ' ?gfbm_set_bdens2: zone dimension size error:'
     write(0,*) '  passed variable  inum = ',inum
     write(0,*) '  module variable  nvol = ',nvol
     write(0,*) '  These must match.  Prior call to GFBM_SET_NTHZNS required.'
     call bad_exit

  end if

  if(nin.ne.nsfast) then

     write(0,*) ' ?gfbm_set_bdens2: species dimension inconsistency:'
     write(0,*) '  passed variable     nin = ',nin
     write(0,*) '  module variable  nsfast = ',nsfast
     write(0,*) '  These must match.  Prior call to GFBM_SET_SPECIES required.'
     call bad_exit

  end if

  if(allocated(eba2pp)) deallocate(eba2pp)
  allocate(eba2pp(nvol,nsfast))

  if(allocated(eba2pl)) deallocate(eba2pl)
  allocate(eba2pl(nvol,nsfast))
  eba2pp = zeba2pp
  eba2pl = zeba2pl

end subroutine gfbm_set_eba2

subroutine gfbm_set_fbm(inume,inum_vpv,inum,nin,zfbm)
  use iso_c_binding, only: fp => c_double

  ! set 2d spatial grid zone volumes

  use get_fbm_module
  implicit none

  integer, intent(in) :: inume  ! energy dimension, must equal maxval(nznbmea)
  integer, intent(in) :: inum_vpv  ! vpll/v dimension, must match nznbma
  integer, intent(in) :: inum   ! number of 2d grid zones: must match nvol!
  integer, intent(in) :: nin    ! number of species, must match nsfast!
  real(fp), intent(in) :: zfbm(inume,inum_vpv,inum,nin)  ! distribution fcn
  !  units: #/cm^3/eV/d(omega/4pi) (standard NUBEAM fbm units).

  !  f(E,vpll/v,<spatial position>,<specie>)

  !------------------
  integer :: ie,ivpv,iz,is
  !------------------

  if(.not.allocated(nznbmea)) then
     write(0,*) ' ?gfbm_set_fbm: energy grid sizes NZNBMEA(...) not set.'
     write(0,*) '  call gfbm_set_species(...) first!'
     write(0,*) '  call gfbm_set_dims(...) first!'

     call bad_exit
  end if

  if(inume.ne.maxval(nznbmea)) then
     write(0,*) ' '
     write(0,*) ' ?gfbm_set_fbm: energy dimension not consistent.'
     write(0,*) '  passed variable  inume = ',inume
     write(0,*) '  stored maxval(nznbmea) = ',maxval(nznbmea)
     write(0,*) '  (these numbers should match).'
     write(0,*) '  call gfbm_set_dims(...) first!'

     call bad_exit
  end if

  if(inum_vpv.ne.nznbma) then

     write(0,*) ' ?gfbm_set_fbm: zone dimension size error:'
     write(0,*) '  passed variable  inum_vpv = ',inum_vpv
     write(0,*) '  module variable  nznbma   = ',nznbma
     write(0,*) '  These must match.  Prior call to GFBM_SET_DIMS required.'
     call bad_exit

  end if

  if(inum.ne.nvol) then

     write(0,*) ' ?gfbm_set_fbm: zone dimension size error:'
     write(0,*) '  passed variable  inum = ',inum
     write(0,*) '  module variable  nvol = ',nvol
     write(0,*) '  These must match.  Prior call to GFBM_SET_NTHZNS required.'
     call bad_exit

  end if

  if(nin.ne.nsfast) then

     write(0,*) ' ?gfbm_set_fbm: species dimension inconsistency:'
     write(0,*) '  passed variable     nin = ',nin
     write(0,*) '  module variable  nsfast = ',nsfast
     write(0,*) '  These must match.  Prior call to GFBM_SET_SPECIES required.'
     call bad_exit

  end if

  if(allocated(fbm)) deallocate(fbm)
  allocate(fbm(inume,inum_vpv,nvol,nsfast))

  do is=1,nsfast
     do iz=1,nvol
        do ivpv=1,inum_vpv
           do ie=1,inume

              if(ie.gt.nznbmea(is)) then
                 fbm(ie,ivpv,iz,is) = 0.0_fp
              else
                 fbm(ie,ivpv,iz,is) = zfbm(ie,ivpv,iz,is)
              end if

           end do
        end do
     end do
  end do

end subroutine gfbm_set_fbm
