! This module contains the gray ouput tables and routines to operate on them
module gray_tables

  use const_and_precisions, only : wp_
  use types,                only : queue, table, wrap

  implicit none

  private
  public init_tables                                  ! Initialising main tables
  public find_flux_surfaces, store_flux_surface       ! Filling the main tables
  public store_ray_data, store_ec_profiles            !
  public store_beam_shape                             !
  public kinetic_profiles, ec_resonance, inputs_maps  ! Extra tables

contains

  pure function is_active(params, id)
    ! Helper to check whether a table is active
    use gray_params,  only : gray_parameters

    ! function arguments
    type(gray_parameters), intent(in) :: params
    integer,               intent(in) :: id
    logical :: is_active

    associate (ids => params%misc%active_tables)
      is_active = any(ids == id)
      if (size(ids) == 1) then
        ! special case, all enabled
        is_active = is_active .or. (ids(1) == -1)
      end if
    end associate
  end function


  subroutine init_tables(params, tbls)
    ! Initialises the gray_main output tables

    use gray_params, only : gray_parameters, gray_tables, ptr=>table_ptr

    ! subroutine arguments
    type(gray_parameters),     intent(in)  :: params
    type(gray_tables), target, intent(out) :: tbls

    call tbls%flux_averages%init( &
      title='flux-averages', id=56, active=is_active(params, 56), &
      labels=[character(64) :: 'ρ_p', 'ρ_t', 'B_avg',             &
              'B_max', 'B_min', 'area', 'vol', 'I_pl',            &
              'J_φ_avg', 'fc', 'ratio_Ja', 'ratio_Jb', 'q'])

    call tbls%flux_surfaces%init( &
      title='flux-surfaces', id=71, active=is_active(params, 71), &
      labels=[character(64) :: 'i', 'ψ_n', 'R', 'z'])

    call tbls%ec_profiles%init( &
      title='ec-profiles', id=48, active=is_active(params, 48), &
      labels=[character(64) :: 'ρ_p', 'ρ_t', 'J_φ',             &
              'J_cdb', 'dPdV', 'I_cd_inside',                   &
              'P_inside', 'index_rt'])

    call tbls%summary%init( &
      title='summary', id=7, active=is_active(params, 7),    &
      labels=[character(64) :: 'I_cd', 'P_abs', 'J_φ_peak',  &
              'dPdV_peak', 'ρ_max_J', 'ρ_avg_J', 'ρ_max_P',  &
              'ρ_avg_P', 'Δρ_avg_J', 'Δρ_avg_P',             &
              'ratio_Ja_max', 'ratio_Jb_max', 's_max',       &
              'ψ', 'χ', 'index_rt', 'J_φ_max', 'dPdV_max',   & 
              'Δρ_J', 'Δρ_P', 'P0', 'cpl1', 'cpl2'])

    call tbls%central_ray%init( &
      title='central-ray', id=4, active=is_active(params, 4), &
      labels=[character(64) :: 's', 'R', 'z', 'φ',            &
              'ψ_n', 'ρ_t', 'n_e', 'T_e', 'B', 'b_x',         &
              'b_y', 'b_z', 'N_⊥', 'N_∥', 'N_x', 'N_y',       &
              'N_z', 'k_im', 'α', 'τ', 'P_t', 'dI/ds',        &
              'n_harm_min', 'n_harm_max', 'i_okhawa',         &
              'index_rt', 'Λ_r', 'X', 'Y',                    &
              '∇X_x', '∇X_y', '∇X_z'])

    call tbls%outer_rays%init( &
      title='outer-rays', id=33, active=is_active(params, 33), &
      labels=[character(64) :: 'i', 'k', 's', 'x', 'y',        &
              'R', 'z', 'ψ_n', 'τ', 'N_∥', 'α', 'index_rt'])

    call tbls%dispersion%init( &
      title='dispersion', id=17, active=is_active(params, 17), &
      labels=[character(64) :: 's', 'Λ_r', 'Λ_i'])

    call tbls%beam_shape%init( &
      title='beam-shape', id=8, active=is_active(params, 8), &
      labels=['s', 'j', 'k', 'x', 'y', 'z', 'r'])

    call tbls%beam_final%init( &
      title='beam-shape-final', id=9, active=is_active(params, 9), &
      labels=['s', 'j', 'k', 'x', 'y', 'z', 'r'])

    call tbls%beam_size%init( &
      title='beam-size', id=12, active=is_active(params, 12), &
      labels=[character(64) :: 's', 'r_min', 'r_max'])

    ! List of pointers to the tables for iterating over
    tbls%iter = [ &
      ptr(tbls%flux_averages), ptr(tbls%flux_surfaces),                   &
      ptr(tbls%summary), ptr(tbls%ec_profiles),                           &
      ptr(tbls%central_ray), ptr(tbls%outer_rays), ptr(tbls%dispersion),  &
      ptr(tbls%beam_shape), ptr(tbls%beam_final), ptr(tbls%beam_size),    &
      ptr(tbls%kinetic_profiles), ptr(tbls%ec_resonance), ptr(tbls%inputs_maps)]

  end subroutine init_tables


  function kinetic_profiles(params) result(tbl)
    ! Generates the plasma kinetic profiles

    use gray_params,  only : gray_parameters, EQ_VACUUM
    use equilibrium,  only : fq, frhotor
    use coreprofiles, only : density, temp
    use magsurf_data, only : npsi, vajphiav

    ! function arguments
    type(gray_parameters), intent(in) :: params
    type(table)                       :: tbl

    ! local variables
    integer :: i, ntail

    real(wp_) :: rho_t, J_phi, dens, ddens
    real(wp_) :: psi_n, rho_p, drho_p

    call tbl%init(title='kinetic-profiles', id=55, active=is_active(params, 25), &
                  labels=[character(64) :: 'psi_n', 'rho_t', 'n_e', 'T_e', 'q', 'J_φ'])

    if (.not. tbl%active) return
    if (params%equilibrium%iequil == EQ_VACUUM) return

    ! Δρ_p for the uniform ρ_p grid
    ! Note: we don't use ψ_n because J_phi has been
    ! sampled uniformly in ρ_p (see flux_average)
    drho_p = 1.0_wp_ / (npsi - 1)

    ! extra points to reach ψ=ψ_bnd
    ntail = ceiling((sqrt(params%profiles%psnbnd) - 1) / drho_p)

    do i = 0, npsi + ntail
      rho_p = i * drho_p
      rho_t = frhotor(rho_p)
      psi_n = rho_p**2
      if (psi_n < 1) then
        J_phi = vajphiav(i+1) * 1.e-6_wp_
      else
        J_phi = 0
      end if
      call density(psi_n, dens, ddens)
      call tbl%append([psi_n, rho_t, dens, temp(psi_n), fq(psi_n), J_phi])
    end do

  end function kinetic_profiles


  function ec_resonance(params, B_res) result(tbl)
    ! Generates the EC resonance surface table

    use const_and_precisions, only : comp_eps
    use gray_params,          only : gray_parameters, EQ_VACUUM
    use equilibrium,          only : rmnm, rmxm, zmnm, zmxm, bfield
    use magsurf_data,         only : npsi
    use types,                only : item
    use cniteq,               only : cniteq_f

    ! function arguments
    type(gray_parameters), intent(in) :: params
    real(wp_),             intent(in) :: B_res
    type(table)                       :: tbl

    call tbl%init(title='ec-resonance', id=70, active=is_active(params, 70), &
                  labels=[character(64) :: 'i', 'B_tot', 'R', 'z'])

    if (.not. tbl%active) return
    if (params%equilibrium%iequil == EQ_VACUUM) return

    block
      ! local variables
      integer   :: j, k, n
      integer   :: n_conts, n_points(10), n_tot
      real(wp_) :: dr, dz, B_min, B_max
      real(wp_) :: B, B_R, B_z, B_phi, B_tot(npsi, npsi)
      real(wp_), dimension(2002) :: R_cont, z_cont
      real(wp_), dimension(npsi) :: R, z

      ! Build a regular (R, z) grid
      dr = (rmxm - rmnm - comp_eps)/(npsi - 1)
      dz = (zmxm - zmnm)/(npsi - 1)
      do j=1,npsi
        R(j) = comp_eps + rmnm + dr*(j - 1)
        z(j) = zmnm + dz*(j - 1)
      end do

      ! B_tot on psi grid
      B_max = -1.0e30_wp_
      B_min = +1.0e30_wp_
      do k = 1, npsi
        do j = 1, npsi
          call bfield(R(j), z(k), B_R, B_z, B_phi)
          B_tot(j,k) = sqrt(B_R**2 + B_z**2 + B_phi**2)
          if(B_tot(j,k) >= B_max) B_max = B_tot(j,k)
          if(B_tot(j,k) <= B_min) B_min = B_tot(j,k)
        end do
      end do

      ! compute Btot=Bres/n with n=1,5
      do n = 1, 5
        B = B_res/n
        if (B >= B_min .and. B <= B_max) then
          n_conts = size(n_points)
          n_tot   = size(R_cont)
          call cniteq_f(R, z, B_tot, size(R), size(z), B, &
                        n_conts, n_points, n_tot, R_cont, z_cont)
          do j = 1, size(R_cont)
            call tbl%append([wrap(j), wrap(B), wrap(R_cont(j)), wrap(z_cont(j))])
          end do
        end if
      end do
    end block

  end function ec_resonance


  function inputs_maps(params, B_res, X_norm) result(tbl)
    ! Generates 2D maps of several input quantities

    use const_and_precisions, only : comp_eps, cm, degree
    use gray_params,          only : gray_parameters, EQ_VACUUM
    use equilibrium,          only : rmnm, rmxm, zmnm, zmxm, pol_flux, bfield
    use coreprofiles,         only : density, temp
    use magsurf_data,         only : npsi

    ! function arguments
    type(gray_parameters), intent(in) :: params
    real(wp_),             intent(in) :: B_res  ! resonant magnetic field, e/m_eω
    real(wp_),             intent(in) :: X_norm ! X normalisation, e²/ε₀m_eω²
    type(table)                       :: tbl

    ! local variables
    integer :: j, k
    real(wp_) :: R0, Npl0
    real(wp_) :: dR, dz, B_R, B_phi, B_z, B
    real(wp_) :: psi_n, ne, dne, Te, X, Y, Npl
    real(wp_), dimension(npsi) :: R, z

    call tbl%init(title='inputs-maps', id=72, active=is_active(params, 72), &
                  labels=[character(64) :: 'R', 'z', 'ψ_n', 'B_R', 'B_z', 'B', &
                                           'n_e', 'T_e', 'X', 'Y', 'N_∥'])

    if (.not. tbl%active) return
    if (params%equilibrium%iequil == EQ_VACUUM) return

    R0   = norm2(params%antenna%pos(1:2)) * cm  ! initial value of R
    Npl0 = sin(params%antenna%beta*degree)      ! initial value of N∥

    ! Build a regular (R, z) grid
    dR = (rmxm - rmnm - comp_eps)/(npsi - 1)
    dz = (zmxm - zmnm)/(npsi - 1)
    do j = 1, npsi
      R(j) = comp_eps + rmnm + dR*(j - 1)
      z(j) = zmnm + dz*(j - 1)
    end do

    do j = 1, npsi
      Npl = Npl0 * R0/r(j)
      do k = 1, npsi
        call pol_flux(r(j), z(k), psi_n)
        call bfield(r(j), z(k), B_R, B_z, B_phi)
        call density(psi_n, ne, dne)
        B  = sqrt(B_R**2 + B_phi**2 + B_z**2)
        X  = X_norm*ne
        Y  = B/B_res
        Te = temp(psi_n)
        call tbl%append([R(j), z(k), psi_n, B_R, B_phi, &
                         B_z, B, ne, Te, X, Y, Npl])
      end do
    end do

  end function inputs_maps


  subroutine find_flux_surfaces(qvals, params, tbl)
    ! Finds the ψ for a set of values of q and stores the
    ! associated surfaces to the flux surface table

    use gray_params,  only : gray_parameters
    use equilibrium,  only : fq, frhotor, rmaxis, zmaxis, zbsup, zbinf
    use magsurf_data, only : contours_psi, npoints
    use logger,       only : log_info
    use minpack,      only : hybrj1
    use types,        only : table

    ! subroutine arguments
    real(wp_),             intent(in)    :: qvals(:)
    type(gray_parameters), intent(in)    :: params
    type(table),           intent(inout) :: tbl

    ! local variables
    integer :: i
    real(wp_) :: rho_t, rho_p, psi_n
    character(256) :: msg ! for log messages formatting

    call log_info('constant ψ surfaces for:', &
                  mod='gray_tables', proc='find_flux_surfaces')
    do i = 1, size(qvals)
      ! Find value of ψ_n for the current q
      block
        real(wp_) :: sol(1), fvec(1), fjac(1,1), wa(7)
        integer   :: info

        ! Note: since we are mostly interested in q=3/2,2 we start
        ! searching near the boundary in case q is not monotonic.
        sol = [0.8_wp_]  ! first guess

        ! Solve fq(ψ_n) = qvals(i) for ψ_n
        call hybrj1(equation, n=1, x=sol, fvec=fvec, fjac=fjac, &
                    ldfjac=1, tol=1e-3_wp_, info=info, wa=wa, lwa=7)
        ! Solution
        psi_n = sol(1)

        ! Handle spurious or no solutions
        if (info /= 1 .or. psi_n < 0 .or. psi_n > 1) cycle
      end block

      ! Compute and print the ψ_n(R,z) = ψ_n₀ contour
      block
        real(wp_), dimension(npoints) :: R_cont, z_cont
        real(wp_) :: R_hi, R_lo, z_hi, z_lo

        ! Guesses for the high,low horzizontal-tangent points
        R_hi = rmaxis;
        z_hi = (zbsup + zmaxis)/2
        R_lo = rmaxis
        z_lo = (zbinf + zmaxis)/2

        call contours_psi(params, psi_n, R_cont, z_cont, R_hi, z_hi, R_lo, z_lo)
        call store_flux_surface(tbl, psi_n, R_cont, z_cont)
      end block

      ! Log the values found for ψ_n, ρ_p, ρ_t
      rho_p = sqrt(psi_n)
      rho_t = frhotor(rho_p)
      write (msg, '(4(x,a,"=",g0.3))') &
        'q', qvals(i), 'ψ_n', psi_n, 'ρ_p', rho_p, 'ρ_t', rho_t
      call log_info(msg, mod='gray_tables', proc='find_flux_surfaces')
    end do

  contains

    subroutine equation(n, x, f, df, ldf, flag)
      ! The equation to solve: f(x) = q(x) - q₀ = 0
      use const_and_precisions, only : comp_eps

      ! subroutine arguments
      integer,   intent(in)    :: n, ldf, flag
      real(wp_), intent(in)    :: x(n)
      real(wp_), intent(inout) :: f(n), df(ldf,n)

      ! optimal step size
      real(wp_), parameter :: e = comp_eps**(1/3.0_wp_)

      if (flag == 1) then
        ! return f(x)
        f(1) = fq(x(1)) - qvals(i)
      else
        ! return f'(x), computed numerically
        df(1,1) = (fq(x(1) + e) - fq(x(1) - e)) / (2*e)
      end if
    end subroutine

  end subroutine find_flux_surfaces


  subroutine store_flux_surface(tbl, psi_n, R, z)
    ! Helper to add a contour to the flux surfaces table

    use const_and_precisions, only : wp_, comp_tiny
    use types,                only : table, wrap

    ! subroutine arguments
    real(wp_),   intent(in)    :: psi_n
    real(wp_),   intent(in)    :: R(:)
    real(wp_),   intent(in)    :: z(:)
    type(table), intent(inout) :: tbl

    ! local variables
    integer :: i

    if (.not. tbl%active) return
    if (psi_n < comp_tiny) return

    do i = 1, size(R)
      call tbl%append([wrap(i), wrap(psi_n), wrap(R(i)), wrap(z(i))])
    end do
  end subroutine store_flux_surface


  subroutine store_ray_data(params, tables, &
                            i, jk, s, P0, pos, psi_n, B, b_n, k0, &
                            N_pl, N_pr, N, N_pr_im, n_e, T_e,     &
                            alpha, tau, dIds, nhm, nhf, iokhawa,  &
                            index_rt, lambda_r, lambda_i, X, Y, grad_X)
    ! Stores some ray variables and local quantities

    use const_and_precisions, only : degree, cm
    use equilibrium,          only : frhotor
    use gray_params,          only : gray_parameters, gray_tables
    use beamdata,             only : unfold_index

    ! subroutine arguments

    ! tables where to store the data
    type(gray_parameters), intent(in)    :: params
    type(gray_tables),     intent(inout) :: tables

    integer,   intent(in) :: i, jk               ! step, ray index
    real(wp_), intent(in) :: s                   ! arclength
    real(wp_), intent(in) :: P0                  ! initial power
    real(wp_), intent(in) :: pos(3)              ! position vector
    real(wp_), intent(in) :: B, b_n(3)           ! magnetic field and unit vector
    real(wp_), intent(in) :: N(3)                ! refractive index vector, N
    real(wp_), intent(in) :: psi_n               ! normalised poloidal flux, ψ_n
    real(wp_), intent(in) :: N_pl, N_pr          ! N_∥, N_⊥
    real(wp_), intent(in) :: N_pr_im             ! Im(N_⊥) from warm dispersion
    real(wp_), intent(in) :: k0                  ! vacuum wavenumber (k₀=ω/c)
    real(wp_), intent(in) :: n_e, T_e            ! electron density and temperature
    real(wp_), intent(in) :: alpha, tau          ! absorption coeff., optical depth
    real(wp_), intent(in) :: dIds                ! ECCD current density
    integer,   intent(in) :: nhm, nhf            ! EC harmonic numbers
    integer,   intent(in) :: iokhawa             !
    integer,   intent(in) :: index_rt            !
    real(wp_), intent(in) :: lambda_r, lambda_i  ! quasioptical dispersion relation
    real(wp_), intent(in) :: X, Y                ! X=(ω_pe/ω)², Y=ω_ce/ω
    real(wp_), intent(in) :: grad_X(3)           ! gradient of X

    ! local variables
    real(wp_) :: s_m, pos_m(3), R_m, phi_deg  ! coordinates in SI units
    real(wp_) :: rho_t, k_im, Pt, dIds_n
    integer :: jkv(2)

    s_m   = s * cm
    pos_m = pos * cm
    R_m   = hypot(pos_m(1), pos_m(2))

    ! Store central ray data
    if (jk == 1 .and. tables%central_ray%active) then
      phi_deg = atan2(pos_m(2), pos_m(1)) / degree
      if(psi_n >= 0 .and. psi_n <= 1) then
        rho_t = frhotor(sqrt(psi_n))
      else
        rho_t = 1.0_wp_
      end if
      k_im   = N_pr_im * k0 / cm
      Pt     = P0 * exp(-tau)
      dIds_n = (1/P0 * dIds) / cm

      call tables%central_ray%append([ &
         wrap(s_m), wrap(R_m), wrap(pos_m(3)), wrap(phi_deg), wrap(psi_n),      &
         wrap(rho_t), wrap(n_e), wrap(T_e), wrap(B), wrap(b_n(1)),              &
         wrap(b_n(2)), wrap(b_n(3)), wrap(N_pr), wrap(N_pl), wrap(N(1)),        &
         wrap(N(2)), wrap(N(3)), wrap(k_im), wrap(alpha), wrap(tau), wrap(Pt),  &
         wrap(dIds_n), wrap(nhm), wrap(nhf), wrap(iokhawa), wrap(index_rt),     &
         wrap(lambda_r), wrap(X), wrap(Y), wrap(grad_X(1)), wrap(grad_X(2)),    &
         wrap(grad_X(3))])
    end if

    ! Store outer rays data
    if (mod(i, params%output%istpl) == 0) then
      jkv = unfold_index(params%raytracing, jk)  ! (j,k)
      if(jkv(1) == params%raytracing%nrayr .and. tables%outer_rays%active) then
        call tables%outer_rays%append([ &
            wrap(i), wrap(jkv(2)), wrap(s_m), wrap(pos_m(1)), wrap(pos_m(2)),  &
            wrap(R_m), wrap(pos_m(3)), wrap(psi_n), wrap(tau), wrap(N_pl),     &
            wrap(alpha), wrap(index_rt)])
      end if
    end if

    ! Store dispersion relation data
    if (jk == params%raytracing%nray .and. tables%dispersion%active) then
      call tables%dispersion%append([s, lambda_r, lambda_i])
    end if

  end subroutine store_ray_data


  subroutine store_ec_profiles(tbl, rho_p, rho_t, J_phi, J_cd, &
                               dPdV, I_ins, P_ins, index_rt)
    ! Helper to add data to the ECRH&CD profiles table

    use types, only : table, wrap

    ! subroutine arguments
    type(table),             intent(inout) :: tbl
    real(wp_), dimension(:), intent(in)    :: rho_p, rho_t, J_phi, J_cd, &
                                              dPdV, I_ins, P_ins
    integer,                 intent(in)    :: index_rt

    integer :: i

    if (.not. tbl%active) return

    do i = 1, size(rho_p)
      call tbl%append([wrap(rho_p(i)), wrap(rho_t(i)), wrap(J_phi(i)), &
                       wrap(J_cd(i)),  wrap(dPdV(i)),  wrap(I_ins(i)), &
                       wrap(P_ins(i)), wrap(index_rt)])
    end do
  end subroutine store_ec_profiles


  subroutine store_beam_shape(params, tables, s, y, full)
    ! Computes the beam shape tables

    use const_and_precisions, only : zero, one, comp_huge
    use types,                only : table, wrap
    use gray_params,          only : raytracing_parameters, gray_tables
    use beamdata,             only : unfold_index

    ! subroutine arguments
    type(raytracing_parameters), intent(in)    :: params
    type(gray_tables), target,   intent(inout) :: tables
    real(wp_),                   intent(in)    :: s(:)     ! arclength
    real(wp_),                   intent(in)    :: y(:, :)  ! (x̅, N̅) vector

    ! whether to show all rays or only the outer ring
    logical, intent(in), optional :: full

    ! local variables
    logical   :: full_
    integer   :: jk, jk0, jkv(2)
    real(wp_) :: n(3), M(3,3), dx(3), dx_loc(3), c
    real(wp_) :: r, r_min, r_max
    type(table), pointer :: beam_shape

    if (.not. (tables%beam_shape%active .or. tables%beam_size%active)) return

    full_ = .false.
    if (present(full)) full_ = full

    if (full_) then
      beam_shape => tables%beam_final
    else
      beam_shape => tables%beam_shape
    end if

    ! Convert to the local beam basis
    !
    ! 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̅.
    !
    n = y(4:6, 1) / norm2(y(4:6, 1))
    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 = reshape([ one,  zero,  zero, &
                   zero,  n(3),  zero, &
                   zero,  zero,  n(3)  ], [3, 3])
    end if

    ! start from the central ray or the last ring
    jk0 = 1 + merge(0, params%nray - params%nrayth, full_)

    r_min = comp_huge
    r_max = 0
    do jk = jk0, params%nray
      dx     = y(1:3, jk) - y(1:3, 1)    ! distance from the central ray
      dx_loc = matmul(M, dx)             ! local ray position
      r      = norm2(dx_loc(1:2))        ! cross-section radius
      jkv    = unfold_index(params, jk)  ! (j, k)

      if (full_ .or. jk /= 1) &
        call beam_shape%append([ &
          wrap(s(jk)), wrap(jkv(1)), wrap(jkv(2)), &
          wrap(dx_loc(1)), wrap(dx_loc(2)), wrap(dx_loc(3)), wrap(r)])

      ! Compute min/max radius
      if (r>=r_max .and. jkv(1)==params%nrayr) r_max = r
      if (r<=r_min .and. jkv(1)==params%nrayr) r_min = r
    end do

    call tables%beam_size%append([s(1), r_min, r_max])

  end subroutine store_beam_shape

end module gray_tables