module beamdata
  use const_and_precisions, only : wp_

  implicit none

contains

  subroutine pweight(params, p0, p0jk)
    ! Power associated to jk-th ray p0jk(j) for total beam power p0
    use const_and_precisions, only : half
    use gray_params,          only : raytracing_parameters

    ! subroutine arguments
    type(raytracing_parameters), intent(in)  :: params
    real(wp_),                   intent(in)  :: p0
    real(wp_),                   intent(out) :: p0jk(:)

    ! local variables
    integer   :: j, jk, jkn
    real(wp_) :: dr, r, w, r0, w0, summ
    real(wp_) :: q(params%nrayr)

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

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


  pure function unfold_index(params, jk)
    ! Maps the ray index jk=1,nray to a pair
    ! (j, k) of radial (j) and azimuthal (k) indices.
    use gray_params, only : raytracing_parameters

    type(raytracing_parameters), intent(in) :: params
    integer,                     intent(in) :: jk
    integer                                 :: unfold_index(2)

    if (jk == 1) then
      unfold_index = [1, 1]
    else
      unfold_index(1) = 2 + (jk - 2) / params%nrayth    ! j
      unfold_index(2) = 1 + mod(jk - 2, params%nrayth)  ! k
    end if
  end function unfold_index


  pure function fold_indices(params, j, k)
    ! Maps the pair (j, k) of radial (j) and azimuthal (k)
    ! indices to a single index jk=1,nray.
    use gray_params, only : raytracing_parameters

    type(raytracing_parameters), intent(in) :: params
    integer,                     intent(in) :: j, k
    integer                                 :: fold_indices

    fold_indices = (j - 2) * params%nrayth + k + 1  ! jk
  end function fold_indices


  pure function tokamak2beam(n, inverse) result(M)
    ! Returns the change-of-basis matrix for switching
    ! from the standard tokamak and the local beam basis.
    use const_and_precisions, only : zero, one
    use utils,                only : diag

    ! function arguments
    real(wp_), intent(in)           :: n(3)    ! central ray direction
    logical,   intent(in), optional :: inverse ! do the opposite transformation
    real(wp_)                       :: M(3,3)  ! change-of-basis matrix

    ! local variables
    real(wp_) :: c
    logical   :: inverse_

    inverse_ = .false.
    if (present(inverse)) inverse_ = inverse

    ! If n̅ is the normalised direction of the central ray,
    ! and the {e̅₁,e̅₂,e̅₃} the standard basis, the unit vectors of
    ! the beam basis are given by:
    !
    !   x̅ = n̅×e̅₃ / |n̅×e̅₃|     =   (n₂/c)e̅₁ +  (-n₁/c)e̅₂
    !   y̅ = n̅×(n̅×e̅₃) / |n̅×e̅₃| = (n₁n₃/c)e̅₁ + (n₂n₃/c)e̅₂ - ce̅₃
    !   z̅ = n̅                 =       n₁e̅₁ +       n₂e̅₂ + n₃e̅₃
    !
    !  where c = |n̅×e̅₃| = √(n₁² + n₂²)
    !
    ! Or in the case where n̅ ∥ e̅₃ by:
    !
    !   x̅ = e̅₁
    !   y̅ = (n₃e̅₃)×e̅₁ = n₃e̅₂
    !   z̅ = n₃e̅₃
    !
    ! So, the change of basis matrix is
    !
    !       [n₂/c    -n₁/c      0]
    !   M = [n₁n₃/c   n₂n₃/c   -c],   c > 0
    !       [n₁       n₂       n₃]
    !
    ! or M = diag(1, n₃, n₃) if c = 0.
    !
    ! By definition then we have: x̅_loc = M x̅.
    !
    c = norm2(n(1:2))
    if (c > 0) then
      M = reshape([ n(2)/c,  n(1)*n(3)/c,  n(1), &
                   -n(1)/c,  n(2)*n(3)/c,  n(2), &
                      zero,           -c,  n(3)  ], [3, 3])
    else
      M = diag([one, n(3), n(3)])
    end if

    ! The basis is orthonormal, so M⁻¹ is just the transpose
    if (inverse_) M = transpose(M)

  end function tokamak2beam

end module beamdata