module gray_params

  use const_and_precisions, only : wp_
  use types,                only : table, contour

  implicit none

  integer, parameter :: lenfnm = 256

  ! Polarisation mode
  enum, bind(c)
    enumerator :: polar_enum = -1
    enumerator :: MODE_O = 1 ! ordinary mode (O)
    enumerator :: MODE_X = 2 ! extraordinary mode (X)
  end enum

  ! Beam parameters format
  enum, bind(c)
    enumerator :: beam_enum = -1
    enumerator :: BEAM_0D = 0 ! fixed beam parameters
    enumerator :: BEAM_1D = 1 ! 1D steering angle table
    enumerator :: BEAM_2D = 2 ! 2D steering angles table
  end enum

  ! Position of the X point
  enum, bind(c)
    enumerator :: x_point_enum = -1
    enumerator :: X_IS_MISSING = 0  ! no X point
    enumerator :: X_AT_TOP     = +1 ! at the top of the plasma
    enumerator :: X_AT_BOTTOM  = -1 ! at the bottom of the plasma
  end enum

  ! MHD equilibrium kind
  enum, bind(c)
    enumerator :: equil_enum = -1
    enumerator :: EQ_VACUUM        = 0 ! vacuum (i.e. no plasma at all)
    enumerator :: EQ_ANALYTICAL    = 1 ! analytical model
    enumerator :: EQ_EQDSK_FULL    = 2 ! G-EQDSK format - data valid on the whole domain
    enumerator :: EQ_EQDSK_PARTIAL = 3 ! G-EQDSK format - data valid only inside the LCFS
  end enum

  ! Temperature/density scaling model
  enum, bind(c)
    enumerator :: scaling_enum = -1
    enumerator :: SCALE_COLLISION = 0 ! preserve collisionality
    enumerator :: SCALE_GREENWALD = 1 ! preserve Greenwald fraction
    enumerator :: SCALE_OFF       = 2 ! don't rescale at all
  end enum

  ! Radial profile coordinate
  enum, bind(c)
    enumerator :: rho_enum = -1
    enumerator :: RHO_TOR = 0 ! ρ_t = √Φ (where Φ is the normalised toroidal flux)
    enumerator :: RHO_POL = 1 ! ρ_p = √ψ (where ψ is the normalised poloidal flux)
    enumerator :: RHO_PSI = 2 ! normalised poloidal flux ψ
  end enum

  ! Plasma profiles kind
  enum, bind(c)
    enumerator :: profiles_enum = -1
    enumerator :: PROF_ANALYTIC = 0 ! analytical model
    enumerator :: PROF_NUMERIC  = 1 ! numerical data (ρ, n_e, T_e, table)
  end enum

  ! Choice of the integration variable
  enum, bind(c)
    enumerator :: step_enum = -1
    enumerator :: STEP_ARCLEN = 0  ! arc length (s)
    enumerator :: STEP_TIME   = 1  ! time (actually c⋅t)
    enumerator :: STEP_PHASE  = 2  ! phase (actually real part of eikonal S_r=k₀⋅φ)
  end enum

  ! Absorption model
  enum, bind(c)
    enumerator :: absorption_enum = -1
    enumerator :: ABSORP_OFF      = 0 ! no absorption at all
    enumerator :: ABSORP_WEAK     = 1 ! weakly relativistic
    enumerator :: ABSORP_FULL     = 2 ! fully relativistic (faster variant)
    enumerator :: ABSORP_FULL_ALT = 3 ! fully relativistic (slower variant)
  end enum

  ! Current drive model
  enum, bind(c)
    enumerator :: current_drive_enum = -1
    enumerator :: CD_OFF         = 0 ! no current drive at all
    enumerator :: CD_COHEN       = 1 ! Cohen
    enumerator :: CD_NO_TRAP     = 2 ! no trapping
    enumerator :: CD_NEOCLASSIC  = 3 ! Neoclassical
  end enum

  ! Antenna/wave launcher parameters
  type antenna_parameters
    ! From gray_params.data:
    real(wp_)                 :: alpha, beta      ! Launching angles
    real(wp_)                 :: power = 1.0_wp_  ! Initial power
    real(wp_)                 :: psi = 0, chi = 0 ! Initial polarisation ellipse parameters
    integer(kind(polar_enum)) :: iox              ! Initial wave mode
    integer(kind(beam_enum))  :: ibeam            ! Beam parameters format
    character(len=lenfnm)     :: filenm           ! beamdata.txt filename

    ! From beamdata.txt:
    real(wp_) :: fghz   ! EC frequency
    real(wp_) :: pos(3) ! Launcher position (tokamak frame)
    real(wp_) :: w(2)   ! Beam waists w(z) for the amplitude           (local frame)
    real(wp_) :: ri(2)  ! Beam inverse radii 1/R(z) for the phase      (local frame)
    real(wp_) :: phi(2) ! Axes orientation angles for amplitude, phase (local frame)
  end type

  ! MHD equilibrium parameters
  type equilibrium_parameters
    real(wp_)                   :: ssplps = 0.005_wp_      ! Spline tension for ψ(R,Z)
    real(wp_)                   :: ssplf  = 0.01_wp_       ! Spline tension for F(ψ)
    real(wp_)                   :: factb  = 1.0_wp_        ! Rescaling factor for B
    integer                     :: sgnb = 0, sgni = 0      ! Sign of B, I
    integer(kind(x_point_enum)) :: ixp    = X_IS_MISSING   ! Position of X point
    integer(kind(equil_enum))   :: iequil = EQ_EQDSK_FULL  ! Equilibrium kind
    integer                     :: icocos = 3              ! COCOS index
    logical                     :: ipsinorm = .false.      ! Whether ψ is normalised
    logical                     :: idesc    = .true.       ! G-EQDSK quirks
    logical                     :: ifreefmt = .false.      !
    character(len=lenfnm)       :: filenm                  ! Equilibrium data filepath
  end type

  ! Kinetic plasma profiles
  type profiles_parameters
    real(wp_)                    :: psnbnd = 1.0_wp_    ! Plasma boundary (ψ: ne(ψ)=0)
    real(wp_)                    :: sspld  = 0.020_wp_  ! Density mollifier width
    real(wp_)                    :: factne = 1.0_wp_    ! Rescale factor for n_e
    real(wp_)                    :: factte = 1.0_wp_    ! Rescale factor for T_e
    integer(kind(scaling_enum))  :: iscal  = SCALE_OFF  ! Rescaling model for n_e,T_e
    integer(kind(rho_enum))      :: irho                ! Radial coordinate
    integer(kind(profiles_enum)) :: iprof               ! Profile kind
    character(len=lenfnm)        :: filenm              ! Profiles data filepath
  end type

  ! Raytracing parameters
  type raytracing_parameters
    real(wp_)                :: dst                  ! Integration step size
    real(wp_)                :: rwmax = 1            ! Normalized maximum radius of beam power
    integer                  :: nray, nrayr, nrayth  ! Total/radial/angular number of rays
    integer                  :: igrad = 0            ! Complex eikonal switch
    integer                  :: nstep = 12000        ! Max number of integration steps
    integer(kind(step_enum)) :: idst  = STEP_TIME    ! Choice of the integration variable
    integer                  :: ipass = 1            ! Number of plasma passes
    logical                  :: ipol  = .false.      ! Whether to compute wave polarisation (from chi, psi)
  end type

  ! EC resonant heating & Current Drive parameters
  type ecrh_cd_parameters
    integer(kind(absorption_enum))    :: iwarm = ABSORP_FULL    ! Choice of power absorption model
    integer(kind(current_drive_enum)) :: ieccd = CD_NEOCLASSIC  ! Choice of current drive model
    integer :: ilarm = 5                                        ! Larmor-radius expansion order
    integer :: imx = -20                                        ! Max iterations for solving the dispersion relation
  end type

  ! Output data parameters
  type output_parameters
    real(wp_), allocatable :: grid(:)         ! Radial grid, defaults to uniform ρ ∈ [0, 1]
    integer                :: ipec = RHO_TOR  ! Radial coordinate of the EC profiles
    integer                :: nrho = 501      ! Number of grid steps for EC profiles + 1
    integer                :: istpr = 5       ! Subsampling factor for beam cross-section (tables 8, 12)
    integer                :: istpl = 5       ! Subsampling factor for outer rays (table 33)
  end type

  ! Other parameters
  type misc_parameters
    real(wp_) :: rwall                        ! R of the limiter (fallback)
    integer, allocatable :: active_tables(:)  ! IDs of output tables to fill in
  end type

  ! All GRAY parameters
  type gray_parameters
    type(antenna_parameters)     :: antenna
    type(equilibrium_parameters) :: equilibrium
    type(profiles_parameters)    :: profiles
    type(raytracing_parameters)  :: raytracing
    type(ecrh_cd_parameters)     :: ecrh_cd
    type(output_parameters)      :: output
    type(misc_parameters)        :: misc
  end type

  ! List of mandatory GRAY parameters
  integer, parameter :: max_name_size = 40
  character(len=*), parameter :: mandatory(*) = [ &
    character(len=max_name_size) :: &
    "antenna.alpha",       &
    "antenna.beta",        &
    "antenna.iox",         &
    "antenna.ibeam",       &
    "antenna.filenm",      &
    "equilibrium.filenm",  &
    "profiles.irho",       &
    "profiles.iprof",      &
    "profiles.filenm",     &
    "raytracing.dst",      &
    "raytracing.nrayr",    &
    "raytracing.nrayth",   &
    "misc.rwall"           &
  ]

  ! Wrapper type for array of pointers
  type table_ptr
    type(table), pointer :: ptr => null()
  end type

  ! Extra outputs tables
  type gray_tables
    type(table) :: flux_averages, flux_surfaces
    type(table) :: summary, ec_profiles
    type(table) :: central_ray, outer_rays, dispersion
    type(table) :: beam_shape, beam_final, beam_size
    type(table) :: kinetic_profiles, ec_resonance, inputs_maps

    ! for iterating over the tables
    type(table_ptr) :: iter(13)
  end type

  ! GRAY final results
  type gray_results
    real(wp_)              :: pabs     ! Total absorbed power
    real(wp_)              :: icd      ! Total driven current
    real(wp_), allocatable :: dpdv(:)  ! Absorbed power density
    real(wp_), allocatable :: jcd(:)   ! Driven current density
    type(gray_tables)      :: tables   ! Extra outputs
  end type

contains

  function make_gray_header(params) result(header)
    ! Writes the Gray output header to a file

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

    ! local variables
    character(len=200)          :: line
    character(len=21*len(line)) :: header
    character(len=8)            :: date
    character(len=10)           :: time
#ifndef REVISION
    character(len=*), parameter :: REVISION="unknown"
#endif

    header = ""

    ! Date and time
    call date_and_time(date, time)
    write(line, '("# Run date/time: ",a4,2("/",a2),1x,2(a2,":"),a6)') &
                      date(1:4), date(5:6), date(7:8), &
                      time(1:2), time(3:4), time(5:10)
    call append(header, line)

    ! Git revision
    write(line, '("# GRAY Git revision: ",a)') REVISION
    call append(header, line)

    ! MHD equilibrium parameters
    if (params%equilibrium%iequil > 0) then
      write(line, '("# EQL input: ",a)') trim(params%equilibrium%filenm)
      call append(header, line)
      ! TODO add missing values
      write(line, '("# EQL B0 R0 aminor Rax zax:",5(1x,g0.5))') &
        0, 0, 0, 0, 0
      call append(header, line)
    else
      write(line, '("# EQL input: N/A (vacuum)")')
      call append(header, line)
      write(line, '("# EQL B0 R0 aminor Rax zax: N/A (vacuum)")')
      call append(header, line)
    end if

    write(line, '("# EQL iequil sgnb sgni factb:",3(1x,g0),1x,g0.3)') &
      params%equilibrium%iequil, params%equilibrium%sgnb, &
      params%equilibrium%sgni, params%equilibrium%factb
    call append(header, line)
    if (params%equilibrium%iequil > 1) then
      write(line, '("# EQL icocos ipsinorm idesc ifreefmt:",4(1x,g0))') &
        params%equilibrium%icocos, params%equilibrium%ipsinorm, &
        params%equilibrium%idesc, params%equilibrium%ifreefmt
      call append(header, line)
      write(line, '("# EQL ssplps ssplf ixp:",2(1x,g8.3e1),1x,g0)') &
        params%equilibrium%ssplps, params%equilibrium%ssplf, &
        params%equilibrium%ixp
      call append(header, line)
    else
      write(line, '("# EQL icocos ipsinorm idesc ifreefmt: N/A (analytical)")')
      call append(header, line)
      write(line, '("# EQL ssplps ssplf ixp: N/A (analytical)")')
      call append(header, line)
    end if

    ! Profiles parameters
    if (params%equilibrium%iequil > 0) then
      write(line, '("# PRF input: ",a)') trim(params%profiles%filenm)
      call append(header, line)
      write(line, '("# PRF iprof iscal factne factte:",4(1x,g0.4))') &
        params%profiles%iprof, params%profiles%iscal, &
        params%profiles%factne, params%profiles%factte
      call append(header, line)
      if (params%profiles%iprof > 0) then
        write(line, '("# PRF irho psnbnd sspld:",3(1x,g0.3))') &
          params%profiles%irho, params%profiles%psnbnd, params%profiles%sspld
        call append(header, line)
      else
        write(line, '("# PRF irho psnbnd sspld: N/A (analytical)")')
        call append(header, line)
      end if
      ! TODO: add missing values
      write(line, '("# PRF Te0 ne0 Zeff0:",3(1x,g0.4))') 0, 0, 0
      call append(header, line)
    else
      write(line, '("# PRF input: N/A (vacuum)")')
      call append(header, line)
      write(line, '("# PRF iprof iscal factne factte: N/A (vacuum)")')
      call append(header, line)
      write(line, '("# PRF irho psnbnd sspld: N/A (vacuum)")')
      call append(header, line)
      write(line, '("# PRF Te0 ne0 Zeff0: N/A (vacuum)")')
      call append(header, line)
    end if

    ! Antenna parameters
    write(line, '("# ANT input: ",a)') trim(params%antenna%filenm)
    call append(header, line)
    write(line, '("# ANT ibeam iox psi chi:",4(1x,g0.4))') &
      params%antenna%ibeam, params%antenna%iox, &
      params%antenna%psi, params%antenna%chi
    call append(header, line)
    write(line, '("# ANT alpha beta power fghz:",4(1x,g0.3))') &
      params%antenna%alpha, params%antenna%beta, &
      params%antenna%power, params%antenna%fghz
    call append(header, line)
    write(line, '("# ANT x0 y0 z0:",3(1x,g0.3))') params%antenna%pos
    call append(header, line)
    write(line, '("# ANT wx wy Rcix Rciy psiw psir:",6(1x,g15.5e1))') &
      params%antenna%w, params%antenna%ri, params%antenna%phi
    call append(header, line)

    ! Other parameters
    write(line, '("# RFL rwall:",1x,e12.5)') params%misc%rwall
    call append(header, line)

    ! code parameters
    write(line, '("# COD igrad idst ipass ipol:",3(1x,i4),1x,l4)') &
      params%raytracing%igrad, params%raytracing%idst, &
      params%raytracing%ipass, params%raytracing%ipol
    call append(header, line)
    write(line, '("# COD nrayr nrayth nstep rwmax dst:",5(1x,g0.4))') &
      params%raytracing%nrayr, params%raytracing%nrayth, &
      params%raytracing%nstep, params%raytracing%rwmax, &
      params%raytracing%dst
    call append(header, line)
    write(line, '("# COD iwarm ilarm imx ieccd:",4(1x,i4))') &
      params%ecrh_cd%iwarm, params%ecrh_cd%ilarm, &
      params%ecrh_cd%imx, params%ecrh_cd%ieccd
    call append(header, line)
    write(line, '("# COD ipec nrho istpr istpl:",4(1x,i4))') &
      params%output%ipec, params%output%nrho, &
      params%output%istpr, params%output%istpl
    header = trim(header) // trim(line)

  contains

    subroutine append(header, line)
      ! Appends a line to the header
      character(len=*), intent(inout) :: header
      character(len=*), intent(in)    :: line
      header = trim(header) // trim(line) // new_line('a')
    end subroutine

  end function make_gray_header


  function update_parameter(params, name, value) result(err)
    ! Updates the value of a parameter, addressed by a string
    ! The return error is:
    !   ERR_SUCCESS on success;
    !   ERR_VALUE   on invalid parameter value;
    !   ERR_UNKNOWN on unknown parameter name.
    !
    ! Ex. update_parameter(params, 'raytracing.nrayr', '10')
    use ini_parser, only : ini_error, ERR_SUCCESS, ERR_VALUE, ERR_UNKNOWN

    ! function arguments
    type(gray_parameters), intent(inout) :: params
    character(*),          intent(in)    :: name, value
    integer(kind(ini_error)) :: err

    character(len=:), allocatable :: temp

    ! hope for the best...
    err = ERR_SUCCESS

#include "gray_params.inc"

  end function update_parameter


  subroutine read_gray_config(filename, params, err)
    ! Reads the GRAY parameters from the gray.ini configuration file
    use ini_parser, only : parse_ini, property_handler, ini_error, ERR_SUCCESS
    use logger,     only : log_error

    ! subroutine arguments
    character(len=*),         intent(in)  :: filename
    type(gray_parameters),    intent(out) :: params
    integer(kind(ini_error)), intent(out) :: err

    integer, parameter :: max_lines = 100
    integer :: count, i
    character(len=max_name_size) :: provided(max_lines), name

    ! parse the configuration
    count = 0
    call parse_ini(filename, ini_handler, err)
    if (err /= ERR_SUCCESS) return

    ! check if all mandatory parameters were given
    do i = 1, size(mandatory)
      name = mandatory(i)
      if (any(provided(1:count)==trim(name))) cycle
      call log_error('mandatory parameter `'//trim(name)//'` is missing!', &
                     mod='gray_params', proc='read_gray_config')
      err = 1
    end do

    ! set computed values
    params%raytracing%nray = 1 + (params%raytracing%nrayr-1) * params%raytracing%nrayth
    contains

      function ini_handler(section, name, value) result(err)
        ! This function handles a single INI property and updates
        ! the `params` structure

        ! function arguments
        character(*), intent(in) :: section, name, value
        integer(kind(ini_error)) :: err

        err = update_parameter(params, section // "." // name, value)

        ! store the parameter name
        count = count + 1
        provided(count) = section // "." // name
      end function ini_handler

  end subroutine read_gray_config


  subroutine read_gray_params(filename, params, err)
    ! Reads the GRAY parameters from the legacy gray_params.data file
    use logger, only : log_error

    ! subrouting arguments
    character(len=*),      intent(in)  :: filename
    type(gray_parameters), intent(out) :: params
    integer,               intent(out) :: err

    ! local variables
    integer :: u, temp(3)

    open(newunit=u, file=filename, status='old', action='read', iostat=err)
    if (err /= 0) then
      call log_error('opening gray_params file ('//filename//') failed!', &
                     mod='gray_params', proc='read_gray_params')
      return
    end if

    read(u, *) params%raytracing%nrayr, params%raytracing%nrayth, &
               params%raytracing%rwmax
    read(u, *) params%raytracing%igrad, params%raytracing%ipass, temp(1)

    ! convert integers to logicals
    params%raytracing%ipol = temp(1) > 0

    read(u, *) params%raytracing%dst, params%raytracing%nstep, &
               params%raytracing%idst

    read(u, *) params%ecrh_cd%iwarm, params%ecrh_cd%ilarm, params%ecrh_cd%imx
    read(u, *) params%ecrh_cd%ieccd

    ! restrict to 0-3
    params%ecrh_cd%ieccd = min(params%ecrh_cd%ieccd, 3)

    read(u, *) params%antenna%alpha, params%antenna%beta
    read(u, *) params%antenna%power
    read(u, *) params%antenna%iox, params%antenna%psi, params%antenna%chi
    read(u, *) params%antenna%ibeam
    read(u, *) params%antenna%filenm

    read(u, *) params%equilibrium%iequil
    read(u, *) params%equilibrium%filenm
    read(u, *) params%equilibrium%icocos, temp(1:3)

    ! convert integers to logicals
    params%equilibrium%ipsinorm = temp(1) > 0
    params%equilibrium%idesc    = temp(2) > 0
    params%equilibrium%ifreefmt = temp(3) > 0

    read(u, *) params%equilibrium%ixp, params%equilibrium%ssplps, &
               params%equilibrium%ssplf
    read(u, *) params%equilibrium%sgnb, params%equilibrium%sgni, &
               params%equilibrium%factb

    read(u, *) params%misc%rwall

    read(u, *) params%profiles%iprof, params%profiles%irho
    read(u, *) params%profiles%filenm
    read(u, *) params%profiles%psnbnd, params%profiles%sspld
    read(u, *) params%profiles%factte, params%profiles%factne, &
               params%profiles%iscal

    read(u, *) params%output%ipec, params%output%nrho
    read(u, *) params%output%istpr, params%output%istpl

    ! remap to 0,1 (as in irho)
    params%output%ipec = mod(params%output%ipec, 2)

    close(u)

    ! set computed values
    params%raytracing%nray = 1 + (params%raytracing%nrayr-1) * params%raytracing%nrayth
  end subroutine read_gray_params

end module gray_params