! This modules handles the loading, interpolation and evaluation of the
! plasma profiles (density, temperature, effective charge)
!
! Two kinds of profiles are supported: analytical (suffix `_an` in the
! subroutine names) or numerical (suffix `_spline`). For the latter, the
! the data is interpolated using splines.
module coreprofiles
  use const_and_precisions, only : wp_, zero, one

  implicit none

  ! Parameters of the plasma profiles splines
  type spline_parameters
    integer :: ndata   ! Number of data points
    integer :: nknots  ! Number of spline knots
    ! Density spline (ψ, knots, B-spline coefficients)
    real(wp_), dimension(:), allocatable :: knots, coeffs
    ! Temperature and effective charge arrays (ψ, T(ψ), Zeff(ψ))
    real(wp_), dimension(:),    allocatable :: psi
    real(wp_), dimension(:, :), allocatable :: temp, zeff
  end type

  ! Parameters of the C² polynomial tail of the density spline
  type density_tail
    real(wp_) :: start   ! ψ₀, start of the tail
    real(wp_) :: end     ! ψ₁, end of the end
    real(wp_) :: value   ! s(ψ₀), value at the start
    real(wp_) :: deriv1  ! s'(ψ₀), first derivative at the start
    real(wp_) :: deriv2  ! s"(ψ₀), second derivative at the start
  end type

  ! Parameters of the analytical profiles model
  type analytic_model
    real(wp_) :: dens0     ! Density scaling factor
    real(wp_) :: n1, n2    ! Density exponents
    real(wp_) :: te0, te1  ! Temperature at ψ=0, ψ=1
    real(wp_) :: t1, t2    ! Temperature exponents
    real(wp_) :: zeff      ! Effective charge
  end type

  ! Global variable storing the state of the module
  type(spline_parameters), save :: spline
  type(density_tail),      save :: tail
  type(analytic_model),    save :: model

  private
  public read_profiles, read_profiles_an       ! Reading data files
  public scale_profiles                        ! Applying rescaling
  public density, temp, fzeff                  ! Accessing interpolated data
  public set_profiles_spline, set_profiles_an  ! Initialising internal state
  public unset_profiles_spline                 ! Deinitialising internal state

contains

  subroutine density(psin, dens, ddens)
    ! Computes the density its first derivative as a function of
    ! normalised poloidal flux.
    !
    ! Note: density has units of 10¹⁹ m⁻³.
    use gray_params, only : iprof
    use dierckx,     only : splev, splder
    use logger,      only : log_error

    implicit none

    ! subroutine arguments
    real(wp_), intent(in)  :: psin         ! normalised poloidal flux
    real(wp_), intent(out) :: dens, ddens  ! density and first derivative

    ! local variables
    integer :: ier                      ! dierck error code
    real(wp_) :: f(1)                   ! dierck output (must be an array)
    real(wp_) :: wrkfd(spline%ndata+4)  ! dierck working space array
    character(256) :: msg               ! for log messages formatting

    ! Initialise both to zero
    dens  = zero
    ddens = zero

    ! Outside the tail end both density and its
    ! derivatives are identically zero
    if (psin >= tail%end .or. psin < 0) return

    if (iprof == 0) then
      ! Use the analytical model
      !
      !   n(ψ) = dens0⋅(1 - ψ^aln1)^aln2
      !
      if (psin > 1) return
      dens  =  model%dens0 * (1 - psin**model%n1)**model%n2
      ddens = -model%dens0 * model%n1*model%n2 * psin**(model%n1 - 1)  &
                           * (1 - psin**model%n1)**(model%n2 - 1)
    else
      ! Use the numerical data

      if (psin < tail%start) then
        ! Use the interpolating spline when in range

        ! Evaluate the spline
        ier = 0
        call splev(spline%knots, spline%nknots, spline%coeffs, &
                   3, [psin], f, 1, ier)
        dens = f(1)

        ! Evaluate the spline 1st derivative
        ier = 0
        call splder(spline%knots, spline%nknots, spline%coeffs, &
                    3, 1, [psin], f, 1, wrkfd, ier)
        ddens = f(1)
        if (abs(dens) < 1.0e-10_wp_) dens = zero
      else
        ! Use a C² polynomial extension outside (ψ > ψ₀)

        ! The tail consists of the product p(ψ)⋅h(t), where:
        !
        ! - p(ψ) is the 2nd order Taylor polynomial of the spline,
        !   centered at ψ₀. See set_profiles_spline for details.
        !
        ! - h(t) is a "smoothing" polynomial in the variable
        !   t = (ψ - ψ₀)/(ψ₁ - ψ₀), defined as:
        !
        !     h(t) = (1 - t)³⋅(1 + 3t + 6t²)
        !
        !   with the following properties:
        !
        !     h(0) = 1   h'(0)=0   h"(0)=0
        !     h(1) = 0   h'(1)=0   h"(1)=0
        block
          real(wp_) :: dpsi, t, p, dp, h, dh
          dpsi = psin - tail%start  ! Δψ = (ψ - ψ₀)

          ! Taylor polynomial p(ψ) and its derivative
          p  = tail%value + dpsi*tail%deriv1 + dpsi**2*tail%deriv2/2
          dp = tail%deriv1 + dpsi*tail%deriv2

          ! Smoothing polynomial h(t) and its derivative
          t  = dpsi/(tail%end - tail%start)
          h  = (1 - t)**3 * (1 + 3*t + 6*t**2)
          dh = -30*(1 - t)**2 * t**2 / (tail%end - tail%start)

          dens  = p*h
          ddens = dp*h + p*dh
        end block
      end if

      if (dens < 0) then
        write (msg, '("negative density:", 2(x,a,"=",g0.3))') &
          'ne', dens, 'ψ', psin
        call log_error(msg, mod='coreprofiles', proc='density')
      end if

    end if
  end subroutine density


  function temp(psin)
    ! Computes the temperature as a function of the
    ! normalised poloidal flux.
    !
    ! Note: temperature has units of keV.
    use gray_params,  only : iprof
    use utils,        only : locate
    use simplespline, only : spli

    implicit none

    ! subroutine arguments
    real(wp_), intent(in) :: psin
    real(wp_) :: temp

    ! local variables
    integer :: k
    real(wp_) :: proft, dps

    temp = zero
    if (psin >= 1 .or. psin < 0) return
    if (iprof == 0) then
      ! Use the analytical model
      !
      !   T(ψ) = (te0 - te1)⋅(1 - ψ^t1)^t2 + te1
      !
      proft = (1 - psin**model%t1)**model%t2
      temp = (model%te0 - model%te1)*proft + model%te1
    else
      ! Use the interpolated numerical data
      call locate(spline%psi, spline%ndata, psin, k)
      k = max(1, min(k, spline%ndata - 1))
      dps = psin - spline%psi(k)
      temp = spli(spline%temp, spline%ndata, k, dps)
    endif
  end function temp


  function fzeff(psin)
    ! Computes the effective charge Zeff as a
    ! function of the normalised poloidal flux.
    use gray_params,  only : iprof
    use utils,        only : locate
    use simplespline, only : spli

    implicit none

    ! subroutine arguments
    real(wp_), intent(in) :: psin
    real(wp_) :: fzeff

    ! local variables
    integer :: k
    real(wp_) :: dps

    fzeff = one
    if (psin >= 1 .or. psin < 0) return
    if (iprof == 0) then
      ! Use the analytical model (just a constant)
      fzeff = model%zeff
    else
      ! Use the interpolated numerical data
      call locate(spline%psi, spline%ndata, psin, k)
      k = max(1, min(k, spline%ndata - 1))
      dps = psin - spline%psi(k)
      fzeff = spli(spline%zeff, spline%ndata, k, dps)
    endif
  end function fzeff


  subroutine read_profiles(filenm, data, unit)
    ! Reads the radial plasma profiles from `file` and store them
    ! into `data`. If given, the file is opened in the `unit` number.
    ! Format notes:
    ! 1. The file is formatted as a table with the following columns:
    !   radial coordinate, temperature, density, effective charge.
    ! 2. The first line is a header specifying the number of rows.
    use utils,       only : get_free_unit
    use gray_params, only : profiles_data
    use logger,      only : log_error

    implicit none

    ! subroutine arguments
    character(len=*),    intent(in)  :: filenm
    type(profiles_data), intent(out) :: data
    integer, optional,   intent(in)  :: unit

    ! local variables
    integer :: u, i, nrows
    integer :: err

    ! Free the arrays when already allocated
    if (allocated(data%psrad)) deallocate(data%psrad)
    if (allocated(data%terad)) deallocate(data%terad)
    if (allocated(data%derad)) deallocate(data%derad)
    if (allocated(data%zfc))   deallocate(data%zfc)

    u = get_free_unit(unit)

    ! Read number of rows and allocate the arrays
    open(file=filenm, status='old', action='read', unit=u, iostat=err)
    if (err /= 0) then
      call log_error('opening profiles file ('//trim(filenm)//') failed!', &
                     mod='coreprofiles', proc="read_profiles")
      call exit(1)
    end if

    read(u, *) nrows
    allocate(data%psrad(nrows), data%terad(nrows), &
             data%derad(nrows), data%zfc(nrows))

    ! Read the table rows
    do i=1,nrows
      read(u, *) data%psrad(i), data%terad(i), data%derad(i), data%zfc(i)
    end do
    close(u)

    ! ??
    data%psrad(1) = max(data%psrad(1), zero)

  end subroutine read_profiles


  subroutine read_profiles_an(filenm, data, unit)
    ! Reads the plasma profiles `data` in the analytical format
    ! from params%filenm.
    ! If given, the file is opened in the `unit` number.
    !
    ! The file should be formatted as follows:
    !
    !   1 dens0 n1 n2
    !   2 te0 te1 t1 t2
    !   3 zeff
    !
    ! See `density`, `temp`, `fzeff` subroutines for the meaning
    ! of the parameters (i.e. the formulae for n,T,Zeff).
    use gray_params, only : profiles_data
    use utils,       only : get_free_unit
    use logger,      only : log_error

    implicit none

    ! subroutine arguments
    character(len=*),    intent(in)  :: filenm
    type(profiles_data), intent(out) :: data
    integer, optional,   intent(in)  :: unit

    ! local variables
    integer :: u
    integer :: err

    u = get_free_unit(unit)

    if (allocated(data%terad)) deallocate(data%terad)
    if (allocated(data%derad)) deallocate(data%derad)
    if (allocated(data%zfc))   deallocate(data%zfc)
    allocate(data%terad(4), data%derad(3), data%zfc(1))

    open(file=filenm, status='old', action='read', unit=u, iostat=err)
    if (err /= 0) then
      call log_error('opening profiles file ('//trim(filenm)//') failed!', &
                     mod='coreprofiles', proc='read_profiles_an')
      call exit(1)
    end if

    read (u,*) data%derad(1:3)  ! dens0, n1, n2
    read (u,*) data%terad(1:4)  ! te0, te1, t1, t2
    read (u,*) data%zfc(1)      ! zeff
    close(u)
  end subroutine read_profiles_an


  subroutine scale_profiles(params, factb, data)
    ! Rescales the temperature and density profiles by `params%factte`
    ! and `params%factne`, respectively, according to the model
    ! specified by `params%iscal`.
    ! See the GRAY user manual for the explanation.
    use gray_params, only : profiles_parameters, profiles_data

    implicit none

    ! subroutine arguments
    type(profiles_parameters), intent(in)    :: params
    real(wp_),                 intent(in)    :: factb
    type(profiles_data),       intent(inout) :: data

    ! local variables
    real(wp_) :: aat, aan, ffact
    integer   :: last_te, last_ne

    if (params%iscal==0) then
      aat = 2.0_wp_/3.0_wp_
      aan = 4.0_wp_/3.0_wp_
    else
      aat = one
      aan = one
    end if

    if (params%iscal==2) then
      ffact = one
    else
      ffact = factb
    end if

    if (params%iprof==0) then
      last_te = 2
      last_ne = 1
    else
      last_te = size(data%terad)
      last_ne = size(data%derad)
    end if

    data%terad(1:last_te) = data%terad(1:last_te) * ffact**aat * params%factte
    data%derad(1:last_ne) = data%derad(1:last_ne) * ffact**aan * params%factne
  end subroutine scale_profiles


  subroutine set_profiles_spline(params, data, launch_pos)
    ! Computes splines for the plasma profiles data and stores them
    ! in their respective global variables, see the top of this file.
    !
    ! When `launch_pos` (cartesian launch coordinates in cm) is present,
    ! the subroutine will also check that the wave launcher is strictly
    ! outside the reconstructed plasma density boundary.
    use simplespline, only : difcs
    use dierckx,      only : curfit, splev, splder
    use gray_params,  only : profiles_parameters, profiles_data
    use logger,       only : log_debug, log_info, log_warning, log_error

    implicit none

    ! subroutine arguments
    type(profiles_parameters), intent(inout) :: params
    type(profiles_data),       intent(inout) :: data
    real(wp_), optional,       intent(in)    :: launch_pos(3)

    ! curfit parameters
    integer, parameter :: iopt = 0  ! smoothing spline mode
    integer, parameter :: kspl = 3  ! order of spline (cubic)

    ! local variables
    integer :: n, npest, ier
    real(wp_) :: xb, xe, fp, ssplne_loc

    ! working space arrays for the dierckx functions
    integer :: lwrkf
    real(wp_), dimension(:), allocatable :: wf, wrkf
    integer,   dimension(:), allocatable :: iwrkf

    ! for log messages formatting
    character(256) :: msg

    n = size(data%psrad)
    npest = n + 4
    lwrkf = n*4 + npest*16
    allocate(wrkf(lwrkf), iwrkf(npest), wf(n))

    ssplne_loc=params%sspld

    ! If necessary, reallocate the spline arrays
    if (.not. allocated(spline%psi)) then
      allocate(spline%psi(n), spline%temp(n, 4), spline%zeff(n, 4))
    else
      if (size(spline%psi) < n) then
        deallocate(spline%psi, spline%temp, spline%zeff)
        allocate(spline%psi(n), spline%temp(n, 4), spline%zeff(n, 4))
      end if
    end if
    if (.not. allocated(spline%coeffs)) then
      allocate(spline%knots(npest), spline%coeffs(npest))
    else
      if (size(spline%coeffs) < npest) then
        deallocate(spline%knots, spline%coeffs)
        allocate(spline%knots(npest), spline%coeffs(npest))
      end if
    end if

    ! Spline interpolation of temperature and effective charge
    call difcs(data%psrad, data%terad, n, iopt, spline%temp, ier)
    call difcs(data%psrad, data%zfc,   n, iopt, spline%zeff, ier)
    spline%psi = data%psrad
    spline%ndata = n

    ! Spline interpolation of density
    xb = zero
    xe = data%psrad(n)
    wf(:) = one
    call curfit(iopt, n, data%psrad, data%derad, wf, xb, xe, kspl, &
                ssplne_loc, npest, spline%nknots, spline%knots,    &
                spline%coeffs, fp,  wrkf, lwrkf, iwrkf, ier)
    ! if ier=-1 data are re-fitted using sspl=0
    if (ier == -1) then
        call log_warning('curfit failed with error -1: re-fitting with '// &
                         's=0', mod='coreprofiles', proc='density')
        ssplne_loc = zero
        call curfit(iopt, n, data%psrad, data%derad, wf, xb, xe, kspl, &
                    ssplne_loc, npest, spline%nknots, spline%knots,    &
                    spline%coeffs, fp, wrkf, lwrkf, iwrkf, ier)
    end if

    ! Computation of the polynomial tail parameters
    !
    ! Note: The density is the only quantity that needs to be evaluated
    ! at the edge. The spline thus has to be extended to transition
    ! smoothly from the last profile point to 0 outside the plasma.
    block
      real(wp_), dimension(1) :: s0, s1, s2  ! spline, 1st, 2nd derivative
      real(wp_), dimension(1) :: delta4      ! discriminant Δ/4 of q(x)
      real(wp_), dimension(1) :: x0, x1      ! vertex of q(x), solution

      ! Compute the coefficients of a 2nd order Taylor polinomial to
      ! extend the spline beyond the last point:
      !
      !   p(ψ) = s(ψ₀) + (ψ - ψ₀)s'(ψ₀) + ½(ψ - ψ₀)²s"(ψ₀)
      !
      !  where s(ψ) is the spline and ψ₀ the last point.
      !
      call splev(spline%knots, spline%nknots, spline%coeffs, kspl, &
                 data%psrad(n:n), s0, 1, ier)
      call splder(spline%knots, spline%nknots, spline%coeffs, kspl, 1, &
                  data%psrad(n:n), s1, 1, wrkf(1:spline%nknots), ier)
      call splder(spline%knots, spline%nknots, spline%coeffs, kspl, 2, &
                  data%psrad(n:n), s2, 1, wrkf(1:spline%nknots), ier)

      ! Determine where to end the tail (to ensure the density remains
      ! positive) from the zeros of the Taylor polynomial p(ψ)
      !
      ! Define x=(ψ - ψ₀), then p(ψ)=0 is rewritten as
      !
      !   q(x) = x² + 2s'/s" x + 2s/s" = 0
      !
      ! The discriminant is Δ/4 = (s'/s")² - 2(s/s") and
      ! the solutions are x = x₀ ± √(Δ/4), with x₀ = -s'/s".
      !
      x0 = -s1 / s2                    ! vertex of parabola y=q(x)
      delta4 = (s1 / s2)**2 - 2*s0/s2  ! Δ/4 of q(x)

      if (delta4(1) > 0) then
        ! Pick the smallest positive zero (implying >ψ₀)
        x1 = x0 + sign(sqrt(delta4), sqrt(delta4) - x0)
      else
        ! There are no zeros, use the parabola vertex
        x1 = x0
        call log_debug('spline extension has no zeros', &
                       mod='coreprofiles', proc='set_profiles_spline')
      end if

      ! Store the tail parameters
      tail%start  = data%psrad(n)
      tail%end    = tail%start + x1(1)
      tail%value  = s0(1)
      tail%deriv1 = s1(1)
      tail%deriv2 = s2(1)
    end block

    ! Make sure the wave launcher does not fall inside the tail
    ! Note: if it does, the initial wave conditions become
    ! invalid as they are given assuming a vacuum (N=1)
    if (present(launch_pos)) then
      block
        use equilibrium, only : equinum_psi
        real(wp_) :: R, Z, psi
        real(wp_), parameter :: cm = 1.0e-2_wp_

        ! Convert from cartesian to cylindrical coordinates
        R = hypot(launch_pos(1), launch_pos(2)) * cm  ! R = √(x²+y²)
        z = launch_pos(3) * cm

        ! Get the poloidal flux at the launcher
        ! Note: this returns -1 when the data is not available
        call equinum_psi(R, z, psi)

        if (psi > tail%start .and. psi < tail%end) then
          ! Fall back to the midpoint of ψ₀ and the launcher ψ
          tail%end = (tail%start + psi)/2
            call log_warning('downscaled tail to not reach the wave launcher', &
                             mod='coreprofiles', proc='set_profiles_spline')
        end if

        if (psi > 0 .and. psi < tail%start) then
          ! This must be a user error, stop here
          write (msg, '(a, a, g0.3, a, g0.3)') &
            'wave launcher is inside the plasma! ', &
            'launcher: ψ=', psi, ' boundary: ψ=', tail%end
          call log_error(msg, mod='coreprofiles', proc='set_profiles_spline')
          call exit(2)
        end if
      end block
    end if

    ! Set the density boundary ψ
    ! Note: this is used to detect entrance in the plasma
    params%psnbnd = tail%end

    write (msg, '(a,g0.4)') 'density boundary: ψ=', tail%end
    call log_info(msg, mod='coreprofiles', proc='set_profiles_spline')

    deallocate(iwrkf, wrkf, wf)
  end subroutine set_profiles_spline


  subroutine unset_profiles_spline
    ! Unsets the splines global variables, see the top of this file.

    implicit none

    if (allocated(spline%psi))     deallocate(spline%psi)
    if (allocated(spline%temp))    deallocate(spline%temp)
    if (allocated(spline%zeff))    deallocate(spline%zeff)
    if (allocated(spline%knots))   deallocate(spline%knots)
    if (allocated(spline%coeffs))  deallocate(spline%coeffs)
  end subroutine unset_profiles_spline


  subroutine set_profiles_an(params, data)
    ! Stores the analytical profiles data in their respective
    ! global variables, see the top of this file.
    use gray_params, only : profiles_parameters, profiles_data

    implicit none

    ! subroutine arguments
    type(profiles_parameters), intent(inout) :: params
    type(profiles_data),       intent(in)    :: data

    model%te0   = data%terad(1)
    model%te1   = data%terad(2)
    model%t1    = data%terad(3)
    model%t2    = data%terad(4)
    model%dens0 = data%derad(1)
    model%n1    = data%derad(2)
    model%n2    = data%derad(3)
    model%zeff  = data%zfc(1)

    ! Define the plasma boundary to be exactly ψ=1
    ! Note: this is used to detect entrance in the plasma
    params%psnbnd = one

  end subroutine set_profiles_an

end module coreprofiles