module multipass
  use const_and_precisions, only : wp_
  use polarization,         only : pol_limit, field_to_ellipse
  use reflections,          only : wall_refl
  use gray_equil,           only : abstract_equil

  implicit none

contains

  subroutine plasma_in(i, equil, x, N, Bres, sox, cpl, &
                       psi, chi, iop, ext, eyt, perfect)
    ! Computes the ray polarisation and power couplings when it enteres the plasma
    use const_and_precisions, only : cm

    ! subroutine arguments
    integer,               intent(in)    :: i              ! ray index
    class(abstract_equil), intent(in)    :: equil          ! MHD equilibrium
    real(wp_),             intent(in)    :: x(3), N(3)     ! position, refactive index
    real(wp_),             intent(in)    :: Bres           ! resonant B field
    integer,               intent(in)    :: sox            ! sign of polarisation mode: -1 ⇒ O, +1 ⇒ X
    real(wp_),             intent(out)   :: cpl(2)         ! power coupling vector (O, X)
    real(wp_),             intent(out)   :: psi, chi       ! polarisation ellipse angles
    integer,               intent(inout) :: iop(:)         ! inside/outside plasma flag
    complex(wp_),          intent(inout) :: ext(:), eyt(:) ! ray polarisation vector (e_x, e_y)
    logical,               intent(in)    :: perfect        ! whether to assume perfect coupling

    ! local variables
    real(wp_) :: R, z, phi
    real(wp_) :: B(3)
    real(wp_) :: c
    complex(wp_) :: e_mode(2), e_ray(2)

    ! Update the inside/outside flag
    iop(i) = iop(i) + 1

    ! Compute magnetic field
    R = norm2(x(1:2)) * cm
    z = x(3) * cm
    phi = atan2(x(2), x(1))
    call equil%b_field(R, z, phi, B=B)

    ! Get the polarisation vector of the given mode
    call pol_limit(N, B, Bres, sox, e_mode(1), e_mode(2))

    if(i == 1) then
      ! For the central ray, compute the polarization ellipse
      call field_to_ellipse(e_mode(1), e_mode(2), psi, chi)
    else
      psi = 0
      chi = 0
    end if

    if (perfect) then
      ! Ignore the given vector and use the expected one
      ! Note: this will give 100% coupling to the current mode
      ext(i) = e_mode(1)
      eyt(i) = e_mode(2)
    end if

    ! Compute the power coupling with the current mode
    e_ray = [ext(i), eyt(i)]
    c = abs(dot_product(e_mode, e_ray))**2

    ! Store both O and X couplings, in this order
    c = merge(c, 1-c, sox == -1)
    cpl = [c, 1-c]
  end subroutine plasma_in


  subroutine plasma_out(i, equil, x, N, Bres, sox, iop, ext, eyt)
    ! Ray exits plasma
    use const_and_precisions, only : cm

    ! subroutine arguments
    integer,               intent(in)    :: i              ! ray index
    class(abstract_equil), intent(in)    :: equil          ! MHD equilibrium
    real(wp_),             intent(in)    :: x(3), N(3)
    real(wp_),             intent(in)    :: Bres
    integer,               intent(in)    :: sox
    integer,               intent(inout) :: iop(:)         ! in/out plasma flag
    complex(wp_),          intent(out)   :: ext(:), eyt(:)

    ! local variables
    real(wp_) :: R, z, phi, B(3)

    iop(i) = iop(i)+1 ! in->out

    ! Compute magnetic field
    R = norm2(x(1:2)) * cm
    z = x(3) * cm
    phi = atan2(x(2), x(1))
    call equil%b_field(R, z, phi, B=B)

    ! Compute polarisation on the boundary
    call pol_limit(N, B, Bres, sox, ext(i), eyt(i))
  end subroutine plasma_out


  subroutine wall_out(i, equil, wall, ins, xv, anv, dst, bres, &
                      sox, psipol1, chipol1, iow, iop, ext, eyt)
    ! Ray exits vessel
    use types, only : contour

    ! subroutine arguments
    integer,               intent(in)    :: i                 ! ray index
    class(abstract_equil), intent(in)    :: equil             ! MHD equilibrium
    type(contour),         intent(in)    :: wall              ! wall contour
    logical,               intent(in)    :: ins               ! inside plasma? (ins=1 plasma/wall overlap)
    real(wp_),             intent(inout) :: xv(3), anv(3)
    real(wp_),             intent(in)    :: dst               ! step size
    real(wp_),             intent(in)    :: bres
    integer,               intent(in)    :: sox
    real(wp_),             intent(out)   :: psipol1, chipol1
    integer,               intent(inout) :: iow(:)            ! in/out vessel and plasma flags
    integer,               intent(inout) :: iop(:)            ! in/out vessel and plasma flags
    complex(wp_),          intent(inout) :: ext(:), eyt(:)

    ! local variables
    integer :: irfl
    real(wp_), dimension(3) :: xvrfl,anvrfl,walln
    complex(wp_) :: ext1,eyt1

    iow(i) = iow(i) + 1                                                   ! out->in
    if(ins) call plasma_out(i, equil, xv, anv, bres, sox, iop, ext, eyt)  ! plasma-wall overlapping
    call wall_refl(wall, xv-dst*anv, anv, ext(i), eyt(i), &
                   xvrfl, anvrfl, ext1, eyt1, walln, irfl)                ! ray reflects at wall
    ext(i) = ext1                                                         ! save parameters at wall reflection
    eyt(i) = eyt1
    xv  = xvrfl
    anv = anvrfl

    if(i == 1) then
      call field_to_ellipse(ext1, eyt1, psipol1, chipol1)
    else
      psipol1 = 0
      chipol1 = 0
    end if

  end subroutine wall_out


  subroutine initbeam(i, iroff, iboff, iwait, stv, jphi_beam, pins_beam, &
                      currins_beam, dpdv_beam, jcd_beam)
    ! Initialization at beam propagation start

    use logger, only : log_info, log_warning

    ! subroutine arguments
    integer,                 intent(in)  :: i              ! beam index
    logical,                 intent(in)  :: iroff(:,:)     ! global ray status (F=active, T=inactive)
    logical,                 intent(out) :: iboff
    logical,                 intent(out) :: iwait(:)
    real(wp_), dimension(:), intent(out) :: jphi_beam, pins_beam, currins_beam
    real(wp_), dimension(:), intent(out) :: dpdv_beam, jcd_beam, stv
    character(256) :: msg ! buffer for formatting log messages

    iboff = .false.      ! beam status (F = active, T = inactive)
    iwait = iroff(:,i)   ! copy ray status for current beam from global ray status
    if(all(iwait)) then  ! no rays active => stop beam
      iboff = .true.
    else if (any(iwait)) then
      ! only some rays active
      write (msg,'("  beam ",g0,": only some rays are active!")') i
      call log_warning(msg, mod='multipass', proc='initbeam')
    end if

    stv          = 0  ! starting step
    jphi_beam    = 0  ! 1D beam profiles
    pins_beam    = 0
    currins_beam = 0
    dpdv_beam    = 0
    jcd_beam     = 0
  end subroutine initbeam


  subroutine initmultipass(params, iroff, yynext, yypnext, yw0, ypw0, stnext, &
                           p0ray, taus, tau1, etau1, cpls, cpl1, lgcpl1, psipv, chipv)
    ! Initialization before pass loop
    use gray_params, only : gray_parameters

    ! subroutine arguments
    type(gray_parameters),       intent(in)  :: params
    logical,   dimension(:,:),   intent(out) :: iroff  ! global ray status (F = active, T = inactive)
    real(wp_), dimension(:),     intent(out) :: p0ray, tau1, etau1, cpl1
    real(wp_), dimension(:),     intent(out) :: lgcpl1, psipv, chipv
    real(wp_), dimension(:,:),   intent(out) :: yw0, ypw0, stnext, taus, cpls
    real(wp_), dimension(:,:,:), intent(out) :: yynext, yypnext

    iroff = .false. ! global ray status (F = active, T = inactive)

    if(.not. params%raytracing%ipol) then
      ! stop other mode (iox=1/2 -> stop ib=2/1 at first pass)
      call turnoffray(0, 1, params%raytracing%ipass, 3-params%antenna%iox, iroff)
    end if

    yynext  = 0  ! starting beam coordinates (1)
    yypnext = 0  ! starting beam coordinates (2)
    yw0     = 0  ! temporary beam coordinates (1)
    ypw0    = 0  ! temporary beam coordinates (2)
    stnext  = 0  ! starting beam step
    p0ray   = 0  ! starting beam power
    taus    = 0  ! beam tau from previous passes
    tau1    = 0
    etau1   = 1
    cpls    = 1  ! beam coupling from previous passes
    cpl1    = 1
    lgcpl1  = 0
    psipv   = 0  ! psi polarization angle at vacuum-plasma boundary
    chipv   = 0  ! chi polarization angle at vacuum-plasma boundary
  end subroutine initmultipass


  subroutine turnoffray(jk, ip, npass, ib, iroff)
    ! Turn off ray propagation

    ! subroutine arguments
    integer, intent(in) :: jk, ip, ib               ! ray (0=all rays), pass, beam indexes
    integer, intent(in) :: npass                    ! total number of passes
    logical, dimension(:,:), intent(out) :: iroff   ! global ray status (F = active, T = inactive)

    ! local variables
    integer :: ipx, i1, i2

    if(jk==0) then                                  ! stop all rays
      do ipx=ip,npass                               ! from pass ip to last pass
        i1 = 2**ipx-2+2**(ipx-ip)*(ib-1)+1          ! first derived beam at pass ipx
        i2 = 2**ipx-2+2**(ipx-ip)*ib                ! last derived beam at pass ipx (i1=i2 for ipx=ip)
        iroff(:,i1:i2) = .true.
      end do
    else                                            ! only stop ray jk
      do ipx=ip,npass
        i1 = 2**ipx-2+2**(ipx-ip)*(ib-1)+1
        i2 = 2**ipx-2+2**(ipx-ip)*ib
        iroff(jk,i1:i2) = .true.
      end do
    end if
  end subroutine turnoffray

end module multipass