module xstraln_mod

  !  this is a "private" module for xstraln.
  !  only routines in the xstraln library should use this module.

  use xstraln_r8types

  implicit NONE
  save

  ! constants

  real*8, parameter :: xpi = 3.1415926535897931D+00
  real*8, parameter :: twopi = 6.2831853071795862D+00

  ! grid sizes

  integer :: nphi, ntheta,nrho, nr,nz, nbrk1,nbrk2, nrho_brk

  ! grid domain parameters:  limits, ranges, inverse spacing

  real*8 :: phimin,phimax,phi_period,delphi_hinv
  real*8 :: thmin,thmax,th_period,delth_hinv
  real*8 :: rhomin,rhomax,rhobrk,delrho_hinv1,delrho_hinv2
  real*8 :: rmin,rbrk1,rbrk2,rmax
  real*8 :: delr_hinv1,delr_hinv2,delr_hinv3
  real*8 :: zmin,zmax,delz_hinv
  real*8 :: dlgrid_min,dlgrid_max

  ! an inverse map yielding rho.le.rho_small is considered
  ! "near" the axis and will be refined by xplasma

  real*8 rho_small
  real*8 invmap_tol   ! error tolerance on axis refinements

  ! the expanded grids...

  real*8, dimension(:), allocatable :: phigrid,thgrid,rhogrid,rgrid,zgrid

  ! the data fields...

  type (r8_fxfields) :: ffx_r8

  type (r8_rzfields) :: frz_r8

  !  flags

  logical :: no_frz = .TRUE.  ! .TRUE. if no f(R,Z) -- no fast inverse map
  logical :: no_brz = .TRUE.  ! .TRUE. if B(R,Z) is not available

  !  if (no_frz) there are no functions at all...
  !-----------------------------------------------

  contains

    !  lookup routines... r4 and r8 versions... vector and scalar versions
    !  (phi,theta,rho) lookup routines

    subroutine fxgrid1_r8(zphi,zth,zrho, &
         iphi1,zfphi, ith1,zfth, irho1,zfrho, iwarn)

      !  lookup a SINGLE point

      real*8, intent(in) :: zphi,zth,zrho  ! coordinates in

      integer, intent(out) :: iphi1        ! grid cell (phi)
      real*8, intent(out) :: zfphi         ! loc w/in cell, 0->iphi1,1->iphi1+1

      integer, intent(out) :: ith1         ! grid cell (theta)
      real*8, intent(out) :: zfth          ! loc w/in cell, 0->ith1, 1->ith1+1

      integer, intent(out) :: irho1        ! grid cell (rho)
      real*8, intent(out) :: zfrho         ! loc w/in cell, 0->irho1,1->irho1+1

      integer, intent(out) :: iwarn        ! 0=normal,...

      !  set iwarn=1 if out of range (rho)
      !  mod function used to bring angle arguments into range, if necessary

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

      real*8 zwk
      real*8, parameter :: zero = 0
      real*8, parameter :: one = 1
      integer offset

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

      ! phi:

      if(no_frz) then
         iwarn=99
         return
      endif

      iwarn=0

      if(nphi.eq.1) then
         iphi1=1
         zfphi=zero
      else
         if((zphi.lt.phimin).or.(zphi.gt.phimax)) then
            zwk=mod(zphi-phimin,phi_period)
            if(zphi.lt.zero) zwk=zwk+phi_period
            zwk=zwk+phimin
            zwk=max(phimin,min(phimax,zwk))
         else
            zwk=zphi
         endif
         offset=(zwk-phimin)*delphi_hinv
         iphi1=min((nphi-1),1+offset)
         zfphi=(zwk-phigrid(iphi1))*delphi_hinv
      endif

      ! theta:

      if((zth.lt.thmin).or.(zth.gt.thmax)) then
         zwk=mod(zth-thmin,th_period)
         if(zth.lt.zero) zwk=zwk+th_period
         zwk=zwk+thmin
         zwk=max(thmin,min(thmax,zwk))
      else
         zwk=zth
      endif
      offset=(zwk-thmin)*delth_hinv
      ith1=min((ntheta-1),1+offset)
      zfth=(zwk-thgrid(ith1))*delth_hinv

      ! rho:

      if(zrho.lt.rhomin) then
         iwarn=1
         irho1=1
         zfrho=zero
      else if(zrho.gt.rhomax) then
         iwarn=1
         irho1=nrho-1
         zfrho=one
      else if(zrho.lt.rhobrk) then
         zwk=zrho
         offset=(zwk-rhomin)*delrho_hinv1
         irho1=min((nrho_brk-1),1+offset)
         zfrho=(zwk-rhogrid(irho1))*delrho_hinv1
      else
         zwk=zrho
         offset=(zwk-rhobrk)*delrho_hinv2
         irho1=min((nrho-1),nrho_brk+offset)
         zfrho=(zwk-rhogrid(irho1))*delrho_hinv2
      endif

    end subroutine fxgrid1_r8

    subroutine fxgrid_r8(ivec,zphi,zth,zrho, &
         iphi1,zfphi, ith1,zfth, irho1,zfrho, iwarn)

      !  lookup a VECTOR of points

      integer ivec                         ! vector size
      real*8, intent(in) :: zphi(ivec),zth(ivec),zrho(ivec)  ! coordinates in

      integer, intent(out) :: iphi1(ivec)  ! grid cell (phi)
      real*8, intent(out) :: zfphi(ivec)   ! loc w/in cell, 0->iphi1,1->iphi1+1

      integer, intent(out) :: ith1(ivec)   ! grid cell (theta)
      real*8, intent(out) :: zfth(ivec)    ! loc w/in cell, 0->ith1, 1->ith1+1

      integer, intent(out) :: irho1(ivec)  ! grid cell (rho)
      real*8, intent(out) :: zfrho(ivec)   ! loc w/in cell, 0->irho1,1->irho1+1

      integer, intent(out) :: iwarn(ivec)  ! 0=normal,...

      !  set iwarn=1 if out of range (rho)
      !  mod function used to bring angle arguments into range, if necessary

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

      real*8 zwk
      real*8, parameter :: zero = 0
      real*8, parameter :: one = 1
      integer offset,i

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

      ! phi:

      if(no_frz) then
         iwarn=99
         return
      endif

      iwarn=0

      do i=1,ivec

         if(nphi.eq.1) then
            iphi1(i)=1
            zfphi(i)=zero
         else
            if((zphi(i).lt.phimin).or.(zphi(i).gt.phimax)) then
               zwk=mod(zphi(i)-phimin,phi_period)
               if(zphi(i).lt.zero) zwk=zwk+phi_period
               zwk=zwk+phimin
               zwk=max(phimin,min(phimax,zwk))
            else
               zwk=zphi(i)
            endif
            offset=(zwk-phimin)*delphi_hinv
            iphi1(i)=min((nphi-1),1+offset)
            zfphi(i)=(zwk-phigrid(iphi1(i)))*delphi_hinv
         endif

         ! theta:

         if((zth(i).lt.thmin).or.(zth(i).gt.thmax)) then
            zwk=mod(zth(i)-thmin,th_period)
            if(zth(i).lt.zero) zwk=zwk+th_period
            zwk=zwk+thmin
            zwk=max(thmin,min(thmax,zwk))
         else
            zwk=zth(i)
         endif
         offset=(zwk-thmin)*delth_hinv
         ith1(i)=min((ntheta-1),1+offset)
         zfth(i)=(zwk-thgrid(ith1(i)))*delth_hinv

         ! rho:

         if(zrho(i).lt.rhomin) then
            iwarn(i)=1
            irho1(i)=1
            zfrho(i)=zero
         else if(zrho(i).gt.rhomax) then
            iwarn(i)=1
            irho1(i)=nrho-1
            zfrho(i)=one
         else if(zrho(i).lt.rhobrk) then
            zwk=zrho(i)
            offset=(zwk-rhomin)*delrho_hinv1
            irho1(i)=min((nrho_brk-1),1+offset)
            zfrho(i)=(zwk-rhogrid(irho1(i)))*delrho_hinv1
         else
            zwk=zrho(i)
            offset=(zwk-rhobrk)*delrho_hinv2
            irho1(i)=min((nrho-1),nrho_brk+offset)
            zfrho(i)=(zwk-rhogrid(irho1(i)))*delrho_hinv2
         endif

      enddo

    end subroutine fxgrid_r8

    !  (phi,R,Z) lookup routines

    subroutine rzgrid1_r8(zphi,zr,zz, &
         iphi1,zfphi, ir1,zfr, iz1,zfz, iwarn)

      !  lookup a SINGLE point

      real*8, intent(in) :: zphi,zr,zz     ! coordinates in

      integer, intent(out) :: iphi1        ! grid cell (phi)
      real*8, intent(out) :: zfphi         ! loc w/in cell, 0->iphi1,1->iphi1+1

      integer, intent(out) :: ir1          ! grid cell (R)
      real*8, intent(out) :: zfr           ! loc w/in cell, 0->ir1, 1->ir1+1

      integer, intent(out) :: iz1          ! grid cell (Z)
      real*8, intent(out) :: zfz           ! loc w/in cell, 0->iz1, 1->iz1+1

      integer, intent(out) :: iwarn        ! 0=normal,...

      !  set iwarn=1 if out of range (R), 2 if out of range (Z)
      !  set iwarn=3 if out of range both R and Z
      !  mod function used to bring angle arguments into range, if necessary

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

      real*8 zwk
      real*8, parameter :: zero = 0
      real*8, parameter :: one = 1
      integer offset

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

      iwarn=0

      if(no_frz) then
         iwarn=99
         return
      endif

      ! phi:

      if(nphi.eq.1) then
         iphi1=1
         zfphi=zero
      else
         if((zphi.lt.phimin).or.(zphi.gt.phimax)) then
            zwk=mod(zphi-phimin,phi_period)
            if(zphi.lt.zero) zwk=zwk+phi_period
            zwk=zwk+phimin
            zwk=max(phimin,min(phimax,zwk))
         else
            zwk=zphi
         endif
         offset=(zwk-phimin)*delphi_hinv
         iphi1=min((nphi-1),1+offset)
         zfphi=(zwk-phigrid(iphi1))*delphi_hinv
      endif

      ! R:

      if(zr.lt.rmin) then
         iwarn=iwarn+1
         ir1=1
         zfr=zero
      else if(zr.gt.rmax) then
         iwarn=iwarn+1
         ir1=nr-1
         zfr=one
      else
         zwk=zr
         if(zwk.lt.rbrk1) then
            offset=(zwk-rmin)*delr_hinv1
            ir1=min((nbrk1-1),1+offset)
            zfr=(zwk-rgrid(ir1))*delr_hinv1
         else if(zwk.lt.rbrk2) then
            offset=(zwk-rbrk1)*delr_hinv2
            ir1=min((nbrk2-1),nbrk1+offset)
            zfr=(zwk-rgrid(ir1))*delr_hinv2
         else
            offset=(zwk-rbrk2)*delr_hinv3
            ir1=min((nr-1),nbrk2+offset)
            zfr=(zwk-rgrid(ir1))*delr_hinv3
         endif
      endif

      ! Z:

      if(zz.lt.zmin) then
         iwarn=iwarn+2
         iz1=1
         zfz=zero
      else if(zz.gt.zmax) then
         iwarn=iwarn+2
         iz1=nz-1
         zfz=one
      else
         zwk=zz
         offset=(zwk-zmin)*delz_hinv
         iz1=min((nz-1),1+offset)
         zfz=(zwk-zgrid(iz1))*delz_hinv
      endif

    end subroutine rzgrid1_r8

    subroutine rzgrid_r8(ivec,zphi,zr,zz, &
         iphi1,zfphi, ir1,zfr, iz1,zfz, iwarn)

      !  lookup a SINGLE point

      integer, intent(in) :: ivec          ! vector size
      real*8, intent(in) :: zphi(ivec),zr(ivec),zz(ivec)     ! coordinates in

      integer, intent(out) :: iphi1(ivec)  ! grid cell (phi)
      real*8, intent(out) :: zfphi(ivec)   ! loc w/in cell, 0->iphi1, ->iphi1+1

      integer, intent(out) :: ir1(ivec)    ! grid cell (R)
      real*8, intent(out) :: zfr(ivec)     ! loc w/in cell, 0->ir1, 1->ir1+1

      integer, intent(out) :: iz1(ivec)    ! grid cell (Z)
      real*8, intent(out) :: zfz(ivec)     ! loc w/in cell, 0->iz1, 1->iz1+1

      integer, intent(out) :: iwarn(ivec)  ! 0=normal,...

      !  set iwarn=1 if out of range (R), 2 if out of range (Z)
      !  set iwarn=3 if out of range both R and Z
      !  mod function used to bring angle arguments into range, if necessary

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

      real*8 zwk
      real*8, parameter :: zero = 0
      real*8, parameter :: one = 1
      integer offset,i

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

      iwarn=0

      if(no_frz) then
         iwarn=99
         return
      endif

      do i=1,ivec

         ! phi:

         if(nphi.eq.1) then
            iphi1(i)=1
            zfphi(i)=zero
         else
            if((zphi(i).lt.phimin).or.(zphi(i).gt.phimax)) then
               zwk=mod(zphi(i)-phimin,phi_period)
               if(zphi(i).lt.zero) zwk=zwk+phi_period
               zwk=zwk+phimin
               zwk=max(phimin,min(phimax,zwk))
            else
               zwk=zphi(i)
            endif
            offset=(zwk-phimin)*delphi_hinv
            iphi1(i)=min((nphi-1),1+offset)
            zfphi(i)=(zwk-phigrid(iphi1(i)))*delphi_hinv
         endif

      ! R:

         if(zr(i).lt.rmin) then
            iwarn(i)=iwarn(i)+1
            ir1(i)=1
            zfr(i)=zero
         else if(zr(i).gt.rmax) then
            iwarn(i)=iwarn(i)+1
            ir1(i)=nr-1
            zfr(i)=one
         else
            zwk=zr(i)
            if(zwk.lt.rbrk1) then
               offset=(zwk-rmin)*delr_hinv1
               ir1(i)=min((nbrk1-1),1+offset)
               zfr(i)=(zwk-rgrid(ir1(i)))*delr_hinv1
            else if(zwk.lt.rbrk2) then
               offset=(zwk-rbrk1)*delr_hinv2
               ir1(i)=min((nbrk2-1),nbrk1+offset)
               zfr(i)=(zwk-rgrid(ir1(i)))*delr_hinv2
            else
               offset=(zwk-rbrk2)*delr_hinv3
               ir1(i)=min((nr-1),nbrk2+offset)
               zfr(i)=(zwk-rgrid(ir1(i)))*delr_hinv3
            endif
         endif

         ! Z:

         if(zz(i).lt.zmin) then
            iwarn(i)=iwarn(i)+2
            iz1(i)=1
            zfz(i)=zero
         else if(zz(i).gt.zmax) then
            iwarn(i)=iwarn(i)+2
            iz1(i)=nz-1
            zfz(i)=one
         else
            zwk=zz(i)
            offset=(zwk-zmin)*delz_hinv
            iz1(i)=min((nz-1),1+offset)
            zfz(i)=(zwk-zgrid(iz1(i)))*delz_hinv
         endif

      enddo

    end subroutine rzgrid_r8

    subroutine free

      !  clear init flags; free any previously allocated memory

      integer :: istat

      ffx_r8%init=0

       frz_r8%init=0

      deallocate(ffx_r8%r,stat=istat)
      deallocate(ffx_r8%z,stat=istat)

      deallocate(frz_r8%rho,stat=istat)
      deallocate(frz_r8%theta,stat=istat)
      deallocate(frz_r8%dbdy,stat=istat)
      deallocate(frz_r8%dlim,stat=istat)
      deallocate(frz_r8%modb,stat=istat)
      deallocate(frz_r8%binv,stat=istat)
      deallocate(frz_r8%br,stat=istat)
      deallocate(frz_r8%bz,stat=istat)
      deallocate(frz_r8%bphi,stat=istat)

    end subroutine free

    subroutine alloc_r8

      !  allocate REAL memory and init to zero

      ffx_r8%init=1
      frz_r8%init=1

      if(.not.no_fRZ) then
         allocate(ffx_r8%r(nphi,ntheta,nrho))
         allocate(ffx_r8%z(nphi,ntheta,nrho))
         ffx_r8%r=0
         ffx_r8%z=0

         allocate(frz_r8%rho(nphi,nr,nz))
         allocate(frz_r8%theta(nphi,nr,nz))
         frz_r8%rho=0
         frz_r8%theta=0

         allocate(frz_r8%dbdy(nphi,nr,nz))
         allocate(frz_r8%dlim(nphi,nr,nz))
         frz_r8%dbdy=0
         frz_r8%dlim=0

         allocate(frz_r8%modb(nphi,nr,nz))
         allocate(frz_r8%binv(nphi,nr,nz))
         frz_r8%modb=0
         frz_r8%binv=0

         allocate(frz_r8%br(nphi,nr,nz))
         allocate(frz_r8%bz(nphi,nr,nz))
         allocate(frz_r8%bphi(nphi,nr,nz))
         frz_r8%br=0
         frz_r8%bz=0
         frz_r8%bphi=0

      endif

    end subroutine alloc_r8

end module xstraln_mod
