! This module provides the ray_projector object that is used to
! compute the final ECCD power and current density radial profiles
module gray_project

  use const_and_precisions, only : wp_

  implicit none

  type ray_projector
    ! Projects the volumetric rays data onto a radial grid to
    ! obtain the ECCD power and current density profiles
    real(wp_), public,  allocatable :: rho_p(:)           ! radial bin edges (ρ_p)
    real(wp_), public,  allocatable :: rho_t(:)           ! radial bin edges (ρ_t)
    real(wp_), private, allocatable :: psi_mid(:)         ! radial bin midpoints (ψ_n)
    real(wp_), private, allocatable :: dA(:), dV(:)       ! area, volume elements
    real(wp_), private, allocatable :: ratio_J_cd_phi(:)  ! J_cd_jintrac / J_φ
    contains
      procedure :: init        ! Initialises the grids etc.
      procedure :: project     ! Projects rays to radial profile
      procedure :: statistics  ! Compute statistics of a profile
  end type

  private
  public ray_projector

contains

  subroutine init(self, params, equil)
    ! Initialises the radial grid and other quantities needed to
    ! compute the radial profiles (dP/dV, J_φ, J_cd)
    use const_and_precisions, only : one
    use gray_params,          only : output_parameters, RHO_POL
    use gray_equil,           only : abstract_equil

    ! subroutine arguments
    class(ray_projector),    intent(inout) :: self
    type(output_parameters), intent(in)    :: params
    class(abstract_equil),   intent(in)    :: equil

    ! local variables
    integer   :: i
    real(wp_) :: drho, rho, rho_mid
    real(wp_) :: V0, V1, A0, A1

    associate (n => params%nrho)

    allocate(self%rho_p(n), self%rho_t(n), self%psi_mid(n))
    allocate(self%dV(n), self%dA(n), self%ratio_J_cd_phi(n))

    V0 = 0
    A0 = 0

    do i = 1, n
      if(allocated(params%grid)) then
        ! Read points from the input grid
        rho  = params%grid(i)
        if (i < n) drho = params%grid(i + 1) - rho
      else
        ! Build a uniform grid
        drho = one/(n - 1)
        rho  = (i - 1)*drho
      end if

      ! Switch between uniform grid in ρ_p or ρ_t
      if (params%ipec == RHO_POL) then
        self%rho_p(i) = rho
        self%rho_t(i) = equil%pol2tor(rho)
        rho_mid = rho + drho/2
      else
        self%rho_t(i) = rho
        self%rho_p(i) = equil%tor2pol(rho)
        rho_mid = equil%tor2pol(rho + drho/2)
      end if

      ! ψ_n at the midpoint of the bin
      !
      ! Note: each bin is associated with a volume element ΔV₁=V₁-V₀,
      ! with V₁,₀ computed at the bin midpoint. For consistency the ΔP
      ! is also computed at the midpoint: ΔP₁ = P_in(ψ₁) - P_in(ψ₀),
      ! where ψ₁,₀ are shifted by Δρ/2 w.r.t ρ₁.
      self%psi_mid(i) = rho_mid**2

      ! Compute ΔV, ΔA for each bin
      call equil%flux_average(rho_mid, area=A1, volume=V1)
      self%dA(i) = abs(A1 - A0)
      self%dV(i) = abs(V1 - V0)
      V0 = V1
      A0 = A1

      ! Compute conversion factors between J_φ and J_cd
      call equil%flux_average(self%rho_p(i), ratio_jintrac_tor=self%ratio_J_cd_phi(i))
    end do

    end associate
  end subroutine init


  subroutine project(self, psi_n, P_abs, I_cd, last_step, &
                     dPdV, J_phi, J_cd, P_in, I_in)
    ! Projects the volumetric rays data ψ_n, P_abs, I_cd onto a radial grid
    ! to obtain the density profiles dP/dV, J_φ, J_cd and the cumulatives
    ! P_in and I_in
    use const_and_precisions, only : zero
    use utils,                only : locate_unord, linear_interp

    ! subroutine arguments
    class(ray_projector), intent(in) :: self
    real(wp_),  intent(in)  :: psi_n(:, :)       ! normalised flux (ray,step)
    real(wp_),  intent(in)  :: P_abs(:, :)       ! absorbed power (ray,step)
    real(wp_),  intent(in)  :: I_cd(:, :)        ! driven current (ray,step)
    integer,    intent(in)  :: last_step(:)      ! last step absorption took place
    real(wp_),  intent(out) :: dPdV(:), J_phi(:) ! power, toroidal current density
    real(wp_),  intent(out) :: J_cd(:)           ! current drive density, JINTRAC def.
    real(wp_),  intent(out) :: P_in(:), I_in(:)  ! power, current within ρ

    ! local variables
    integer   :: i, ray, bin
    integer   :: cross(maxval(last_step)+1), ncross
    real(wp_) :: I_entry, I_exit, P_entry, P_exit

    ! We start by computing the cumulatives first, and obtain the densities
    ! as finite differences later. This should avoid precision losses.
    P_in = 0
    I_in = 0

    ! For each bin in the radial grid
    bins: do concurrent (bin = 1 : size(self%rho_t))
      ! For each ray
      rays: do ray = 1, size(psi_n, dim=1)
        ! Skip rays with no absorption
        if (last_step(ray) == 1) cycle

        ! Find where the ray crosses the flux surface of the current bin
        call locate_unord(psi_n(ray, :), self%psi_mid(bin), cross, last_step(ray), ncross)

        if (psi_n(ray, 1) < 0) then
          ! Skip unphysical crossing from -1 (ψ undefined)
          cross = cross(2:ncross)
          ncross = ncross - 1
        else if (psi_n(ray, 1) < self%psi_mid(bin)) then
          ! Add a crossing when starting within ψ (reflections!)
          cross = [1, cross(1:ncross)]
        end if

        ! For each crossing
        crossings: do i = 1, ncross, 2
          ! Interpolate P_abs(ψ) and I_cd(ψ) at the entry
          P_entry = interp(psi_n(ray,:), P_abs(ray,:), cross(i), self%psi_mid(bin))
          I_entry = interp(psi_n(ray,:),  I_cd(ray,:), cross(i), self%psi_mid(bin))

          if (i+1 <= ncross) then
            ! Interpolate P_abs(ψ) and I_cd(ψ) at the exit
            P_exit = interp(psi_n(ray,:), P_abs(ray,:), cross(i+1), self%psi_mid(bin))
            I_exit = interp(psi_n(ray,:),  I_cd(ray,:), cross(i+1), self%psi_mid(bin))
          else
            ! The ray never left, use last_step as exit point
            P_exit = P_abs(ray, last_step(ray))
            I_exit = I_cd(ray, last_step(ray))
          end if

          ! Accumulate the deposition from this crossing
          P_in(bin) = P_in(bin) + (P_exit - P_entry)
          I_in(bin) = I_in(bin) + (I_exit - I_entry)
        end do crossings
      end do rays
    end do bins

    ! Compute densities
    dPdV(1)  = P_in(1) / self%dV(1)
    J_phi(1) = I_in(1) / self%dA(1)

    do concurrent (bin = 2 : size(self%rho_t))
      dPdV(bin)  = max(P_in(bin) - P_in(bin-1), zero) / self%dV(bin)
      J_phi(bin) = (I_in(bin) - I_in(bin-1)) / self%dA(bin)
    end do

    J_cd = J_phi * self%ratio_J_cd_phi

  contains

    pure real(wp_) function interp(x, y, i, x1)
      ! Interpolate linearly y(x) at x=x(i)
      real(wp_), intent(in) :: x(:), y(:), x1
      integer,   intent(in) :: i
      interp = linear_interp(x(i), y(i), x(i+1), y(i+1), x1)
    end function interp

  end subroutine project


  subroutine statistics(self, equil, dPdV, J_phi, rho_avg_P, drho_avg_P, &
                        rho_avg_J, drho_avg_J, dPdV_peak, J_phi_peak,    &
                        rho_max_P,  drho_max_P,  rho_max_J, drho_max_J,  &
                        dPdV_max, J_phi_max, ratio_Ja_max, ratio_Jb_max)
    ! Given the radial profiles dP/dV and J_φ computes the following statics:

    use const_and_precisions, only : pi
    use gray_equil,           only : abstract_equil

    class(ray_projector),  intent(in)  :: self
    class(abstract_equil), intent(in)  :: equil

    real(wp_), intent(in)  :: dPdV(:), J_phi(:)                  ! dP/dV(ρ), J_φ(ρ) profiles
    real(wp_), intent(out) :: rho_avg_P, drho_avg_P, dPdV_peak   ! ⟨ρ_t⟩, Δρ_t, max dP/dV using dP/dV⋅dV as a PDF
    real(wp_), intent(out) :: rho_avg_J, drho_avg_J, J_phi_peak  ! ⟨ρ_t⟩, Δρ_t, max |J_φ| using |J_φ|⋅dA as a PDF
    real(wp_), intent(out) :: rho_max_P, drho_max_P, dPdV_max    ! argmax, full-width at max/e and max of dPdV
    real(wp_), intent(out) :: rho_max_J, drho_max_J, J_phi_max   ! argmax, full-width at max/e and max of |J_φ|
    real(wp_), intent(out) :: ratio_Ja_max, ratio_Jb_max         ! max of J_cd_astra/J_φ, J_cd_jintrac/J_φ

    ! local variables
    real(wp_) :: P_tot, I_tot
    real(wp_) :: dVdrho_t, dAdrho_t
    integer   :: i_max

    ! ⟨ρ_t⟩, Δρ_t using dP/dV⋅dV as a PDF
    ! Note: the factor √8 is to match the full-width at max/e of a Gaussian
    P_tot = sum(dPdV*self%dV)
    if (P_tot > 0) then
      rho_avg_P  = sum(self%rho_t * dPdV*self%dV)/P_tot
      drho_avg_P = sqrt(8 * sum((self%rho_t - rho_avg_P)**2 * dPdV*self%dV)/P_tot)
    else
      rho_avg_P  = 1
      drho_avg_P = 0
    end if

    ! ⟨ρ_t⟩, Δρ_t using |J_φ|⋅dA as a PDF
    I_tot = sum(abs(J_phi)*self%dA)
    if (I_tot > 0) then
      rho_avg_J  = sum(self%rho_t * abs(J_phi)*self%dA)/I_tot
      drho_avg_J = sqrt(8* sum((self%rho_t - rho_avg_J)**2 * abs(J_phi)*self%dA)/I_tot)
    else
      rho_avg_J  = 1
      drho_avg_J = 0
    end if

    if (P_tot > 0) then
      ! Compute max of dP/dV as if it were a Gaussian
      call equil%flux_average(equil%tor2pol(rho_avg_P), dVdrho_t=dVdrho_t)
      dPdV_peak = P_tot * 2/(sqrt(pi) * drho_avg_P * dVdrho_t)

      ! Compute numerically argmax, max and full-width at max/e of dP/dV
      call profwidth(self%rho_t, dPdV, i_max, drho_max_P)
      rho_max_P = self%rho_t(i_max)
      dPdV_max = dPdV(i_max)
    else
      dPdV_peak  = 0
      rho_max_P  = 1
      dPdV_max   = 0
      drho_max_P = 0
    end if

    if (I_tot > 0) then
      ! Compute max of J_φ as if it were a Gaussian
      call equil%flux_average(equil%tor2pol(rho_avg_J), dAdrho_t=dAdrho_t, &
                              ratio_astra_tor=ratio_Ja_max, ratio_jintrac_tor=ratio_Jb_max)
      J_phi_peak = sum(J_phi*self%dA) * 2/(sqrt(pi)*drho_avg_J*dAdrho_t)

      ! Compute numerically argmax, max and full-width at max/e of J_φ
      call profwidth(self%rho_t, abs(J_phi), i_max, drho_max_J)
      rho_max_J = self%rho_t(i_max)
      J_phi_max = J_phi(i_max)
    else
      J_phi_peak   = 0
      rho_max_J    = 1
      J_phi_max    = 0
      drho_max_J   = 0
      ratio_Ja_max = 0
      ratio_Jb_max = 0
    end if

    ! Signal that J_φ has alternating sign
    if ((I_tot - abs(sum(J_phi*self%dA)))/ I_tot > 0.1) then
      rho_avg_J  = -rho_avg_J
      drho_avg_J = -drho_avg_J
      rho_max_J  = -rho_max_J
      drho_max_J = -drho_max_J
    end if
  end subroutine statistics


  subroutine profwidth(x, y, i_max, width)
    ! Computes the argmax and full-width at max/e of the curve y(x)

    use const_and_precisions, only : emn1
    use utils, only : locate_unord, linear_interp

    ! subroutine arguments
    real(wp_), intent(in)  :: x(:), y(:)
    integer, intent(out)   :: i_max
    real(wp_), intent(out) :: width

    ! local variables
    integer :: i, left, right
    integer :: nsols, sols(size(y))
    real(wp_) :: x_left, x_right

    ! Find maximum
    i_max = maxloc(y, dim=1)

    ! Find width
    width = 0

    ! Find solutions to y(x) = y_max⋅e⁻¹
    call locate_unord(y, y(i_max)*emn1, sols, size(y), nsols)
    if (nsols < 2) return

    ! Find the two closest solutions to the maximum
    left  = sols(1)
    right = sols(nsols)
    do i = 1, nsols
      if (sols(i) > left  .and. sols(i) < i_max) left = sols(i)
      if (sols(i) < right .and. sols(i) > i_max) right = sols(i)
    end do

    ! Interpolate linearly x(y) for better accuracy
    x_left  = linear_interp(y(left),   x(left),  y(left+1),  x(left+1), y(i_max)*emn1)
    x_right = linear_interp(y(right), x(right), y(right+1), x(right+1), y(i_max)*emn1)

    width = x_right - x_left
  end subroutine profwidth

end module gray_project