module nbdrive_common
  use iso_c_binding, only: fp => c_double

  implicit none
  private :: fp
  SAVE
  PUBLIC

  ! shared information -- nubeam_comp_sub library
  !   (reduced version of nbdrive from an earlier NUBEAM interface)

  ! note methods.  Arrays can only be allocated once!  Allocate species
  ! arrays first, then v-space grids!

  integer :: ngmax = 0  ! no. of non-impurity thermal species
  integer :: nsfast= 0  ! no. of NUBEAM fast species (beam + fusion products)
                        ! nsfast may include RF species eventually
  integer :: nsbfp = 0  ! number of beam + number of fusion product species

  integer :: nznbme= 0  ! max no. of energy zones, of any fast specie
  integer :: nznbma= 0  ! no. of pitch zones (vpll/v), assumed same for all
                        ! fast species


  integer :: id_mcgrid = 0 ! old_xplasma MCgrid ID

  integer :: lun_msgs = 6  ! LUN for message output
  integer :: lun_summary = 6  ! LUN for summary output (nbdrive_summary*.f90)

  logical :: nlsym=.FALSE. ! MHD equilibrium up-down symmetry flag

  integer :: nzones_fb = 0 ! number of radial zone rows in MC grid
  integer :: nzones_mc = 0 ! total number of zones in MC grid

  !----------------------
  ! internal use only:

  logical, private :: ialloc_species = .FALSE.
  logical, private :: ialloc_vgrids = .FALSE.

  !----------------------
  ! allocatable arrays:

  ! species data:

  integer, dimension(:), allocatable :: indx_fi  ! NUBEAM species indices
  !  (output of NUBEAM indx_nbfi(iz,ia,itype) calls)

  logical, dimension(:), allocatable :: nlfprod  ! flag fusion product species

  integer, dimension(:), allocatable :: nznbmea  ! #energy zones, by specie
  !  nznbmea(1:nsfast)

  ! species labels:

  character(len=40) :: runid = 'undefined_runid'

  character(len=40), dimension(:), allocatable :: label_fastions

  character(len=40), dimension(:), allocatable :: label_thions

  !  Z & A of species:

  real(fp), dimension(:), allocatable :: aplasm  ! A of thermal species
  !  aplasm(1:ngmax)
  real(fp), dimension(:), allocatable :: backz   ! Z of thermal species
  !  backz(1:ngmax)

  real(fp), dimension(:), allocatable :: abeams  ! A of fast species
  !  abeams(1:nsfast)
  real(fp), dimension(:), allocatable :: xzbeams ! Z of fast species
  !  xzbeams(1:nsfast)

  !  v-space grids
  !     zone-centers and zone boundaries (+1 dim)

  real(fp), dimension(:), allocatable :: xxksid,xxksidb  ! pitch (vpll/v)
  !  xxksid(nznbma),xxksidb(nznbma+1)

  real(fp), dimension(:,:), allocatable :: efbm,efbmb    ! energy
  !  efbm(nznbme,nsfast),efbmb(nznbme+1,nsfast)
  !    for species ns, efbm(1:nznbmea(ns),ns) are actual energy zone ctrs.

  !  v-space data poloidal zone indexing

  integer, dimension(:), allocatable :: nthzns,nthzsum
  !  nthzns(nzones_fb)    ! #th zones per zone row (e.g. 4,8,12,16,...)
  !  nthzsum(0:nzones_fb) ! #th zones total to each zone bdy
  !                          (e.g. 0,4,12,24,40,...)

  !  v-space data arrays

  real(fp), dimension(:), allocatable :: bmvol
  real(fp), dimension(:,:), allocatable :: bdens2
  real(fp), dimension(:,:), allocatable :: eba2pl
  real(fp), dimension(:,:), allocatable :: eba2pp
  real(fp), dimension(:,:,:,:), allocatable :: fbm

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

CONTAINS

  subroutine alloc_species(ierr)
    use iso_c_binding, only: fp => c_double

    integer, intent(out) :: ierr  ! completion status, 0=OK

    ! allocate species arrays

    ierr = 0

    if(ialloc_species) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_species -- already allocated.'
       ierr=2
       return
    end if

    if(ngmax.eq.0) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_species -- ngmax=0.'
       ierr=1
    end if

    if(nsfast.eq.0) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_species -- nsfast=0.'
       ierr=1
    end if

    if(ierr.ne.0) return

    ! OK...................

    allocate(indx_fi(nsfast)); indx_fi = 0

    allocate(nznbmea(nsfast)); nznbmea = 0

    allocate(nlfprod(nsfast)); nlfprod = .FALSE.

    allocate(label_fastions(nsfast)); label_fastions=' '
    allocate(label_thions(ngmax)); label_thions=' '

    allocate(aplasm(ngmax)); aplasm = 0.0_fp
    allocate(backz(ngmax));   backz = 0.0_fp

    allocate(abeams(nsfast));   abeams = 0.0_fp
    allocate(xzbeams(nsfast)); xzbeams = 0.0_fp

    ialloc_species = .TRUE.

  end subroutine alloc_species


  subroutine alloc_vgrids(ierr)
    use iso_c_binding, only: fp => c_double

    integer, intent(out) :: ierr  ! completion status, 0=OK

    ! allocate v-grid arrays

    !-------------
    real(fp) :: aztest
    !-------------

    ierr = 0

    ! error checks: (a) not already allocated; (b) species data set;
    ! (c) grid array sizes set; (d) individual species nznbmea(...) set.

    if(.NOT.ialloc_species) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_vgrids -- call alloc_species first!'
       ierr=3
       return
    end if

    if(ialloc_vgrids) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_vgrids -- already allocated.'
       ierr=2
       return
    end if

    if(nznbme.eq.0) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_vgrids -- nznbme=0.'
       ierr=1
    end if

    if(nznbma.eq.0) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_vgrids -- nznbma=0.'
       ierr=1
    end if

    if(ierr.ne.0) return

    aztest = 1.0_fp

    aztest = min(aztest,minval(backz))
    aztest = min(aztest,minval(aplasm))
    aztest = min(aztest,minval(abeams))
    aztest = min(aztest,minval(xzbeams))

    if(aztest.lt.(0.75_fp)) then
       ierr=4
       write(lun_msgs,*) ' ?nbdrive_common:alloc_vgrids -- species data not set:'
       write(lun_msgs,*) '  APLASM = ',aplasm
       write(lun_msgs,*) '  BACKZ  = ',backz
       write(lun_msgs,*) '  ABEAMS = ',abeams
       write(lun_msgs,*) '  XZBEAMS= ',xzbeams

       return
    end if

    if(minval(nznbmea).eq.0) then
       ierr=5
       write(lun_msgs,*) ' ?nbdrive_common:alloc_vgrids -- nznbmea(...) not set:'
       write(lun_msgs,*) '  nznbmea = ',nznbmea

       return
    end if

    if(maxval(nznbmea).gt.nznbme) then
       ierr=6
       write(lun_msgs,*) ' ?nbdrive_common:alloc_vgrids -- maxval(nznbmea) exceeds nznbme:'
       write(lun_msgs,*) '  nznbmea = ',nznbmea

       return
    end if

    ! OK...................

    allocate(xxksid(nznbma),xxksidb(nznbma+1))
    xxksid = 0.0_fp
    xxksidb = 0.0_fp

    allocate(efbm(nznbme,nsfast),efbmb(nznbme+1,nsfast))
    efbm = 0.0_fp
    efbmb = 0.0_fp

  end subroutine alloc_vgrids

  subroutine alloc_nthzns(ierr)
    use iso_c_binding, only: fp => c_double

    integer, intent(out) :: ierr

    if(nzones_fb.le.0) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_nthzns -- nzones_fb = ',nzones_fb
       ierr = 1
    else if(allocated(nthzns)) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_nthzns -- nthzns already allocated.'
       ierr = 2
    else
       ierr = 0
       allocate(nthzns(nzones_fb),nthzsum(0:nzones_fb))
       nthzns = 0
       nthzsum = 0
    end if

  end subroutine alloc_nthzns

  subroutine alloc_data(ierr)
    use iso_c_binding, only: fp => c_double

    integer, intent(out) :: ierr

    if(min(nzones_fb,nzones_mc,nznbme,nznbma,nsfast).eq.0) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_data -- some dimensions not set:'
       write(lun_msgs,*) '  nzones_fb = ',nzones_fb
       write(lun_msgs,*) '  nzones_mc = ',nzones_mc
       write(lun_msgs,*) '  nznbme    = ',nznbme
       write(lun_msgs,*) '  nznbma    = ',nznbma
       write(lun_msgs,*) '  nsfast    = ',nsfast
       ierr = 1
    else if(allocated(fbm)) then
       write(lun_msgs,*) ' ?nbdrive_common:alloc_data -- fbm already allocated.'
       ierr = 2
    else
       ierr = 0
       allocate(bmvol(nzones_mc),bdens2(nzones_mc,nsfast))
       allocate(eba2pl(nzones_mc,nsfast),eba2pp(nzones_mc,nsfast))
       bmvol = 0.0_fp
       bdens2 = 0.0_fp
       eba2pl = 0.0_fp
       eba2pp = 0.0_fp
       allocate(fbm(nznbme,nznbma,nzones_mc,nsfast))
       fbm = 0.0_fp
    end if

  end subroutine alloc_data

end module nbdrive_common
