module beamdata
  use const_and_precisions, only : wp_

  implicit none

  integer, save :: nray,nrayr,nrayth,nstep,jkray1
  real(wp_), save :: dst,h,hh,h6,rwmax,twodr2

contains

  subroutine init_btr(rtrparam,ywork,ypwork,xc,du1,gri,ggri,psjki,ppabs,ccci, &
                      tau0,alphaabs0,dids0,ccci0,p0jk,ext,eyt,iiv)
    use gray_params, only : raytracing_parameters
    use const_and_precisions, only : half,two
    type(raytracing_parameters), intent(inout) :: rtrparam
    real(wp_),  dimension(:,:), intent(out), pointer :: ywork,ypwork, &
       gri,psjki,ppabs,ccci
    real(wp_), dimension(:,:,:), intent(out), pointer :: xc,du1,ggri
    real(wp_),    dimension(:), intent(out), pointer :: tau0,alphaabs0, &
       dids0,ccci0,p0jk
    complex(wp_), dimension(:), intent(out), pointer :: ext, eyt
    integer,      dimension(:), intent(out), pointer :: iiv
    
    integer :: jray1

    dst = rtrparam%dst
    h   = dst
    hh  = h*half
    h6  = h/6.0_wp_
    
    nrayr  = rtrparam%nrayr
    nrayth = rtrparam%nrayth
    rwmax  = rtrparam%rwmax

    if (nrayr==1) then
      nrayth = 1
      jray1  = 1
    else
      jray1 = 1 + max(nint((nrayr-1)/rwmax),1)
      rwmax = dble(nrayr-1)/dble(jray1-1)
    end if
    nray   = (nrayr-1)*nrayth + 1
    jkray1 = (jray1-2)*nrayth + 2
    rtrparam%nray = nray

    if(nrayr>1) then
      twodr2 = two*(rwmax/(nrayr-1))**2
    else
      twodr2 = two
    end if

    nstep=rtrparam%nstep
    
    ! Allocate for each ray:
    ! y = (x, k),
    ! yp = dy/dt,
    ! etc.
    call alloc_beam(ywork,ypwork,xc,du1,gri,ggri,psjki,ppabs,ccci, &
                    tau0,alphaabs0,dids0,ccci0,p0jk,ext,eyt,iiv)
  end subroutine init_btr



  subroutine pweight(p0,p0jk)
! power associated to jk-th ray p0jk(j) for total beam power p0
    use const_and_precisions, only : wp_, zero, one, half, two
! arguments
    real(wp_), intent(in) :: p0
    real(wp_), dimension(:), intent(out) :: p0jk
! local variables
    integer :: j,jk,jkn
    real(wp_) :: dr,r,w,r0,w0,summ
    real(wp_), dimension(nrayr) :: q

    if(nray==1) then
      q(1) = one
    else
      dr = rwmax/dble(nrayr - 1)
      summ = zero
      w0 = one
      do j = 1, nrayr
        r    = (dble(j) - half)*dr
        w    = exp(-two*r**2)
        q(j) = w - w0
        summ = summ + q(j)
        r0   = r
        w0   = w
      end do
      q = q/summ
      q(2:) = q(2:)/nrayth
    end if

    p0jk(1)=q(1)*p0
    jk=2
    do j=2,nrayr
      jkn=jk+nrayth
      p0jk(jk:jkn-1)=q(j)*p0
      jk=jkn
    end do
  end subroutine pweight



  ! Functions to map radial (j), poloidal (k) ray
  ! indices to a single global index (i)

  function rayi2jk(i) result(jk)
    integer, intent(in) :: i
    integer, dimension(2) :: jk
    integer :: ioff

    if (i>1) then
      ioff  = i - 2
      jk(1) = ioff/nrayth              ! jr-2
      jk(2) = ioff - jk(1)*nrayth + 1  ! kt
!     jk(2) = mod(ioff,nrayth) + 1     ! kt
      jk(1) = jk(1) + 2                ! jr
    else
      jk = 1
    end if
  end function rayi2jk



  function rayi2j(i) result(jr)
    integer, intent(in) :: i
    integer :: jr

!   jr = max(1, (i-2)/nrayth + 2)
    if (i>1) then
      jr = (i-2)/nrayth + 2
    else
      jr = 1
    end if
  end function rayi2j



  function rayi2k(i) result(kt)
    integer, intent(in) :: i
    integer :: kt

!   kt = max(1, mod(i-2,nrayth) + 1)
    if (i>1) then
      kt = mod(i-2,nrayth) + 1
    else
      kt = 1
    end if
  end function rayi2k



  function rayjk2i(jr,kt) result(i)
    integer, intent(in) :: jr,kt
    integer :: i
    
!    i = max(1, (jr-2)*nrayth + kt + 1)
    if (jr>1) then
      i = (jr-2)*nrayth + kt + 1
    else
      i = 1
    end if
  end function rayjk2i



  subroutine alloc_beam(ywork,ypwork,xc,du1,gri,ggri,psjki,ppabs,ccci, &
                        tau0,alphaabs0,dids0,ccci0,p0jk,ext,eyt,iiv)
    real(wp_),  dimension(:,:), intent(out), pointer :: ywork,ypwork, &
       gri,psjki,ppabs,ccci
    real(wp_), dimension(:,:,:), intent(out), pointer :: xc,du1,ggri
    real(wp_),    dimension(:), intent(out), pointer :: tau0,alphaabs0, &
       dids0,ccci0,p0jk
    complex(wp_), dimension(:), intent(out), pointer :: ext, eyt
    integer,      dimension(:), intent(out), pointer :: iiv
    
    call dealloc_beam(ywork,ypwork,xc,du1,gri,ggri,psjki,ppabs,ccci, &
                      tau0,alphaabs0,dids0,ccci0,p0jk,ext,eyt,iiv)

    allocate(ywork(6,nray), ypwork(6,nray), gri(3,nray), ggri(3,3,nray), &
             xc(3,nrayth,nrayr),  du1(3,nrayth,nrayr),                   &
             psjki(nray,nstep),   ppabs(nray,nstep),   ccci(nray,nstep), &
             tau0(nray),   alphaabs0(nray), dids0(nray),    ccci0(nray), &
             p0jk(nray),         ext(nray),   eyt(nray),      iiv(nray))
  end subroutine alloc_beam



  subroutine dealloc_beam(ywork,ypwork,xc,du1,gri,ggri,psjki,ppabs,ccci, &
                          tau0,alphaabs0,dids0,ccci0,p0jk,ext,eyt,iiv)
    real(wp_),  dimension(:,:), intent(out), pointer :: ywork,ypwork, &
       gri,psjki,ppabs,ccci
    real(wp_), dimension(:,:,:), intent(out), pointer :: xc,du1,ggri
    real(wp_),    dimension(:), intent(out), pointer :: tau0,alphaabs0, &
       dids0,ccci0,p0jk
    complex(wp_), dimension(:), intent(out), pointer :: ext, eyt
    integer,      dimension(:), intent(out), pointer :: iiv

    if (associated(ywork))     deallocate(ywork)
    if (associated(ypwork))    deallocate(ypwork)
    if (associated(xc))        deallocate(xc)
    if (associated(du1))       deallocate(du1)
    if (associated(gri))       deallocate(gri)
    if (associated(ggri))      deallocate(ggri)
    if (associated(psjki))     deallocate(psjki)
    if (associated(ppabs))     deallocate(ppabs)
    if (associated(ccci))      deallocate(ccci)
    if (associated(tau0))      deallocate(tau0)
    if (associated(alphaabs0)) deallocate(alphaabs0)
    if (associated(dids0))     deallocate(dids0)
    if (associated(ccci0))     deallocate(ccci0)
    if (associated(p0jk))      deallocate(p0jk)
    if (associated(ext))       deallocate(ext)
    if (associated(eyt))       deallocate(eyt)
    if (associated(iiv))       deallocate(iiv)
  end subroutine dealloc_beam

end module beamdata