!Physical and mathematical constants
module constantParameters
  implicit none

  public

  integer,  parameter:: dp    = kind(0.d0)          ! Precision
  real(dp), parameter:: PI    = 4.0_dp*ATAN(1.0_dp) ! Number pi
  real(dp), parameter:: qe    = 1.60217662e-19_dp   ! Elementary charge
  real(dp), parameter:: kb    = 1.38064852e-23_dp   ! Boltzmann constants SI
  real(dp), parameter:: eV2J  = qe                  ! Electron volt to Joule conversion
  real(dp), parameter:: eps_0 = 8.8542e-12_dp       ! Epsilon_0

  real(dp), parameter:: eV_to_K   = 11604.5_dp ! Convert eV to K
  real(dp), parameter:: cm3_to_m3 = 1.0e6_dp   ! Convert cm^-3 to m^-3

end module constantParameters

module referenceValues
  use constantParameters, only: dp
  implicit none
  
  real(dp):: L_ref, t_ref, n_ref, u_ref, Temp_ref ! Reference values
  real(dp):: phi_ref                              ! Reference values

end module referenceValues

module output
  use constantParameters, only: dp
  implicit none

  public
  private:: dataRef_id, dataBC_id, dataF_id, dataPhi_id, dataCum_id
  private:: formatFloat, formatSep, formatTime

  integer:: everyOutput, everyWrite
  character(:), allocatable:: pathOutput
  integer, parameter:: dataRef_id = 10
  integer, parameter:: dataBC_id  = 20
  integer, parameter:: dataF_id   = 30
  integer, parameter:: dataPhi_id = 40
  integer, parameter:: dataCum_id = 50
  character(len=7), parameter:: formatFloat = 'ES0.6e3'
  character(len=3), parameter:: formatSep   = '","'
  character(len=7):: formatTime

  contains
    subroutine setTimeFormat(nt)
      integer, intent(in):: nt
      integer:: l

      l=max(1,ceiling(log10(real(nt))))
      write(formatTime, '(A2,I0,".",I0,A1)') '(I',l,l,')'

    end subroutine setTimeFormat

    subroutine createPath()
      character(8)  :: date_now
      character(10) :: time_now

      call date_and_time(date_now, time_now)

      !Compose the folder name
      pathOutput = date_now(1:4) // '-' // date_now(5:6) // '-' // date_now(7:8) // '_' // &
                   time_now(1:2) // '.' // time_now(3:4) // '.' // time_now(5:6) // '/'

      call system('mkdir ' // pathOutput)

    end subroutine createPath

    subroutine writeOutputF(t, dt, nr, r, nv, v, f)
      use referenceValues, only: L_ref, n_ref, u_ref, t_ref

      integer, intent(in):: t
      integer,  intent(in):: nr, nv
      real(dp), intent(in):: dt
      real(dp), intent(in):: r(1:nr)
      real(dp), intent(in):: v(1:nv)
      real(dp), intent(in):: f(1:nr,1:nv)
      character(len=30) :: myfmt
      character(:), allocatable:: filename
      integer:: i
      character(len=10):: timeString

      write (timeString, formatTime) t
      filename = 'time_' // trim(timeString) // '_f_i.csv'
      write (*, '(A, A)') 'Writing: ', filename

      open(unit=dataF_id, file=pathOutput // filename)
      write(dataF_id, '(A)') "t (s)"
      write(dataF_id, '('//formatFloat//')') t*dt*t_ref
      write(myfmt, "(I0)") nr
      myfmt =  '(A,' // trim(myfmt) // '(' // formatSep // ',' // formatFloat // '))'
      write(dataF_id, myfmt) "v (m/s) / r (m)", r*L_ref
      write(myfmt, "(I0)") nr
      myfmt =  '(' // formatFloat // ',' // trim(myfmt) // '(' // formatSep // ',' // formatFloat // '))'
      do i = 1, nv
        write(dataF_id, myfmt) v(i)*u_ref, f(:,i)*n_ref/u_ref

      end do

      close(unit=dataF_id)

    end subroutine writeOutputF

    subroutine writeOutputPhi(t, dt, nr, r, phi, E, n_e)
      use constantParameters, only: eV_to_K
      use referenceValues,    only: L_ref, phi_ref, t_ref, n_ref

      integer, intent(in):: t
      integer,  intent(in):: nr
      real(dp), intent(in):: dt
      real(dp), intent(in):: r(1:nr)
      real(dp), intent(in):: phi(1:nr)
      real(dp), intent(in):: E(1:nr)
      real(dp), intent(in):: n_e(1:nr)
      character(:), allocatable:: filename
      integer:: i
      character(len=10):: timeString

      write (timeString, formatTime) t
      filename = 'time_' // trim(timeString)//'_phi.csv'
      write (*, '(A, A)') 'Writing: ', filename

      open(unit=dataPhi_id, file=pathOutput//filename)
      write(dataPhi_id, '(A)') "t (s)"
      write(dataPhi_id, '('//formatFloat//')') t*dt*t_ref
      write(dataPhi_id, '(A,3('//formatSep//',A))') "r (m)","phi (V)","E (V m^-1)","n_e (m^-3)"
      do i = 1, nr
        write(dataPhi_id, '('//formatFloat//',3('//formatSep //','//formatFloat//'))') &
                          r(i)*L_ref,                                                  &
                          phi(i)*phi_ref,                                              &
                          E(i)*phi_ref/L_ref,                                          &
                          n_e(i)*n_ref           

      end do

      close(unit=dataPhi_id)

    end subroutine writeOutputPhi

    subroutine writeOutputMom(t, dt, nr, r, n_i, u_i, T_i, Z)
      use constantParameters, only: eV_to_K
      use referenceValues,    only: L_ref, t_ref, n_ref, u_ref, Temp_ref

      integer, intent(in):: t
      integer,  intent(in):: nr
      real(dp), intent(in):: dt
      real(dp), intent(in):: r(1:nr)
      real(dp), intent(in):: n_i(1:nr)
      real(dp), intent(in):: u_i(1:nr)
      real(dp), intent(in):: T_i(1:nr)
      real(dp), intent(in):: Z(1:nr)
      character(:), allocatable:: filename
      integer:: i
      character(len=10):: timeString

      write (timeString, formatTime) t
      filename = 'time_' // trim(timeString)//'_mom_i.csv'
      write (*, '(A, A)') 'Writing: ', filename

      open(unit=dataPhi_id, file=pathOutput//filename)
      write(dataPhi_id, '(A)') "t (s)"
      write(dataPhi_id, '('//formatFloat//')') t*dt*t_ref
      write(dataPhi_id, '(A,4('//formatSep//',A))') "r (m)","n_i (m^-3)","u_i (m s^-1)", "T_i (eV)","Zave"
      do i = 1, nr
        write(dataPhi_id, '('//formatFloat//',4('//formatSep //','//formatFloat//'))') &
                          r(i)*L_ref,                                                  &
                          n_i(i)*n_ref,                                                &
                          u_i(i)*u_ref,                                                &
                          T_i(i)*Temp_ref/ev_to_K,                                     &
                          Z(i)                    

      end do

      close(unit=dataPhi_id)

    end subroutine writeOutputMom

    subroutine writeOutputBoundary(t, dt, n, u, Temp, Z)
      use constantParameters, only: eV_to_K
      use referenceValues,    only: t_ref, n_ref, u_ref, Temp_ref

      integer, intent(in):: t
      real(dp), intent(in):: dt
      real(dp), intent(in):: n, u, Temp, Z
      character(len=6), parameter:: filename = 'bc.csv'
      logical:: res

      inquire(file=pathOutput // filename, exist=res)
      if (.not. res) then
        write (*, '(A, A)') 'Writing: ', filename
        open(unit=dataBC_id, file=pathOutput // filename, action='write', position='append')
        write(dataBC_id, '(A,4(' // formatSep // ',A))') 't (s)', 'n (m^-3)', 'u (m s^-1)', 'T (eV)', 'Z'
        close(dataBC_id)

      end if

      open(unit=dataBC_id, file=pathOutput // filename, action='write', position='append')
      write(dataBC_id, '(' // formatFloat // ',4('// formatSep // ',' // formatFloat // '))') &
                        t*dt*t_ref,  n*n_ref, u*u_ref, Temp*Temp_ref/eV_to_K, Z

      close(dataBC_id)

    end subroutine writeOutputBoundary

    subroutine writeOutputRef()
      use referenceValues, only: t_ref, L_ref, n_ref, u_ref, Temp_ref, phi_ref

      character(len=7), parameter:: filename = 'ref.csv'

      write (*, '(A, A)') 'Writing: ', filename
      open(unit=dataRef_id, file=pathOutput // filename)
      write(dataRef_id, '(A,5(' // formatSep // ',A))') 't_ref (s)', 'L_ref (m)', 'n_ref (m^-3)', &
                                                        'u_ref (m s^-1)', 'T_ref (K)', 'phi_ref (V)'
      write(dataRef_id, '(' // formatFloat // ',5('// formatSep // ',' // formatFloat // '))') &
                        t_ref, L_ref, n_ref, &
                        u_ref, Temp_ref, phi_ref

      close(dataRef_id)

    end subroutine writeOutputRef

    subroutine writeOutputFCum(t, dt, r, nv, v, f)
      use referenceValues, only: L_ref, n_ref, u_ref, t_ref

      integer, intent(in):: t
      real(dp), intent(in):: dt
      integer,  intent(in):: nv
      real(dp), intent(in):: r
      real(dp), intent(in):: v(1:nv)
      real(dp), intent(in):: f(1:nv)
      character(len=30) :: myfmt
      character(:), allocatable:: filename
      integer:: i
      character(len=10):: timeString

      write (timeString, formatTime) t
      filename = 'time_' // trim(timeString) // '_fCum_i.csv'
      write (*, '(A, A)') 'Writing: ', filename

      open(unit=dataCum_id, file=pathOutput // filename)
      write(dataCum_id, '(A)') "t (s)"
      write(dataCum_id, '('//formatFloat//')') t*dt*t_ref
      write(myfmt, "(I0)") 1
      myfmt =  '(A,' // trim(myfmt) // '(' // formatSep // ',' // formatFloat // '))'
      write(dataCum_id, myfmt) "v (m/s) / r (m)", r*L_ref
      write(myfmt, "(I0)") 1
      myfmt =  '(' // formatFloat // ',' // trim(myfmt) // '(' // formatSep // ',' // formatFloat // '))'
      do i = 1, nv
        write(dataCum_id, myfmt) v(i)*u_ref, f(i)*n_ref/u_ref

      end do

      close(unit=dataCum_id)

    end subroutine writeOutputFCum

end module output

module eos
  use constantParameters, only: dp
  implicit none

  private
  public:: T_to_Z

  contains
    pure function T_to_Z(T) result(Z)
      use constantParameters, only: eV_to_K
      use referenceValues,    only: Temp_ref
      implicit none

      real(dp), intent(in):: T
      real(dp):: Z

      ! Z = (Temp_ref * T / eV_to_K)**0.6
      ! Z = max(Z, 1.0_dp)
      ! Z = min(Z, 22.0_dp)
      Z = 12.0_dp

    end function T_to_Z

end module eos

program plasmaExpansion
  use constantParameters, only: dp, kb, qe, eps_0, ev_to_K, cm3_to_m3, PI
  use output
  use referenceValues
  use eos,                only: T_to_Z
  use omp_lib
  implicit none

  real(dp), parameter:: m_i = 1.9712258e-25_dp ! Tin atom mass in kg
  real(dp), parameter:: m_e = 9.1093837e-31_dp ! Electron mass in kg
  real(dp), parameter:: gam = 1.0_dp           ! Adiabatic coefficient

  real(dp):: r0, rf
  real(dp), allocatable, dimension(:):: r
  real(dp):: v0, vf
  real(dp), allocatable, dimension(:):: v
  real(dp):: t0, tf
  real(dp):: dr, dv, dt
  integer::  nr, nv, nt
  integer::   i,  j,  t
  integer:: j0 ! First integer of positive velocity

  real(dp):: Temp_bc ! Temperature
  real(dp):: Temp0, TempF
  real(dp):: n_ecr   ! Electron critical density for the laser
  real(dp):: c_s     ! Ion sound speed
  real(dp):: u_bc    ! Injection velocity
  real(dp):: u_bc0, u_bcF
  real(dp):: n_bc    ! Injection density
  real(dp):: n_bc0, n_bcF
  integer:: t_bc0, t_bcF

  real(dp), allocatable, dimension(:,:):: f_i, f_i_old
  real(dp), allocatable, dimension(:):: f0   ! Boundary at r = x_0
  real(dp), allocatable, dimension(:):: n_i
  real(dp), allocatable, dimension(:):: u_i
  real(dp), allocatable, dimension(:):: E_i
  real(dp), allocatable, dimension(:):: T_i
  real(dp), allocatable, dimension(:):: n_e
  real(dp), allocatable, dimension(:):: Zave
  real(dp), allocatable, dimension(:):: diag, diag_low, diag_high
  real(dp), allocatable, dimension(:,:):: A
  real(dp), allocatable, dimension(:):: Res
  real(dp), allocatable, dimension(:):: b
  integer:: info

  real(dp), allocatable, dimension(:):: phi, phi_old, E
  real(dp):: phiConv
  real(dp):: phi0, phiF
  integer:: k

  real(dp), allocatable, dimension(:):: fCum_i
  real(dp):: rCum
  integer:: rCum_index

  ! Set number of threads
  call omp_set_num_threads(8)

  ! Set reference numbers (in SI units)
  Temp_ref = 60.0_dp * eV_to_K
  n_ref    = 1.0e19_dp * cm3_to_m3
  t_ref    = sqrt(eps_0 * m_i / (n_ref * 1.0_dp * qe**2)) ! 1.0_dp represents Z = 1 for reference values
  u_ref    = sqrt(kb * Temp_ref / m_i)
  L_ref    = u_ref * t_ref
  phi_ref  = kb * Temp_ref / qe

  ! Set position to calculate cumulative sum of f (non-dimensional units)
  rCum = 5.0e-3 / L_ref

  ! Set input parameters (remember these have to be in non-dimensional units)
  Temp0    = 60.0_dp * eV_to_K / Temp_ref
  TempF    = 60.0_dp * eV_to_K / Temp_ref
  n_ecr    = 1.0e19_dp * cm3_to_m3 / n_ref
  c_s      = sqrt(T_to_Z(Temp0) * gam * Temp0)
  u_bc0    = c_s!sqrt(Temp0)
  u_bcF    = u_bc0!sqrt(TempF)
  n_bc0    = n_ecr / T_to_Z(Temp0)
  n_bcF    = n_bc0!n_ecr*1.0e-1 / T_to_Z(Temp0)

  ! Set domain boundaries (non-dimensional units)
  r0 = 200.0e-6_dp / L_ref
  rf =   6.0e-3_dp / L_ref
  dr = 1.0e3_dp
  nr = nint((rf - r0) / dr) + 1
  dr =      (rf - r0) / float(nr-1)
  allocate(r(1:nr))
  do i = 1, nr
    r(i) = dr * float(i-1) + r0
  end do

  ! Index for cumulative sum
  rCum_index = minloc(abs(r - rCum), 1)

  v0 =-1.0e1_dp*c_s
  vf = 2.0e1_dp*c_s
  dv = 1.0e-1_dp
  nv = nint((vf - v0) / dv) + 1
  dv =      (vf - v0) / float(nv-1)
  allocate(v(1:nv))
  do j = 1, nv
    v(j) = dv * float(j-1) + v0
  end do
  ! Shift v mesh so it passes by 0
  v = v - (minval(abs(v)))
  j0 = minloc(abs(v), 1)
  if (v(j0) < 0.0_dp) then
    j0 = j0 + 1
  end if

  t0 = 0.0_dp
  tf = 1.0e-6_dp / t_ref
  ! tf = 1.0e1_dp * (rf - r0) / c_s
  dt = 1.0e-1_dp*dr/(10.0_dp*c_s)
  nt = nint((tf - t0) / dt) 
  dt =      (tf - t0) / float(nt)

  t_bc0 = nint(100.0e-9_dp / t_ref / dt)
  t_bcF = nint(105.0e-9_dp / t_ref / dt)

  everyOutput = nint(1.0e-9_dp/t_ref/dt)
  if (everyOutput  == 0) then
    everyOutput = 1

  end if

  everyWrite = everyOutput/10
  if (everyWrite  == 0) then
    everyWrite = 1

  end if

  write(*, '(A,ES0.4e3)') 'CFL: ', dt*vf/dr

  ! Allocate vectors
  allocate(f_i(1:nr,1:nv), f_i_old(1:nr,1:nv))
  allocate(n_i(1:nr))
  allocate(u_i(1:nr), E_i(1:nr), T_i(1:nr))
  allocate(Zave(1:nr))
  allocate(n_e(1:nr))
  allocate(phi(1:nr), phi_old(1:nr), E(1:nr))
  allocate(fCum_i(1:nv))
  f_i     = 0.0_dp
  f_i_old = 0.0_dp
  n_i     = 0.0_dp
  u_i     = 0.0_dp
  E_i     = 0.0_dp
  T_i     = 0.0_dp
  n_e     = 1.0_dp
  Zave    = 0.0_dp
  phi     = 0.0_dp
  phi_old = 0.0_dp
  E       = 0.0_dp
  fCum_i  = 0.0_dp

  ! Allocate matrix for Poisson equation
  allocate(diag(1:nr), diag_low(1:nr-1), diag_high(1:nr-1))
  allocate(b(1:nr))
  diag      = 0.0_dp
  diag_low  = 0.0_dp
  diag_high = 0.0_dp
  b         = 0.0_dp
  diag           = -2.0_dp / dr**2
  diag_low       =  1.0_dp / dr**2 - 1.0_dp / (r(2:nr)   * 2.0_dp * dr)
  diag_high      =  1.0_dp / dr**2 + 1.0_dp / (r(1:nr-1) * 2.0_dp * dr)
  diag(1)        =  1.0_dp ! Dirichlet
  diag_high(1)   =  0.0_dp ! Dirichlet
  ! diag_high(1)   =  2.0_dp / dr**2  ! Neumann
  ! diag(nr)       =  1.0_dp ! Dirichlet
  ! diag_low(nr-1) =  0.0_dp ! Dirichlet
  diag_low(nr-1) =  2.0_dp / dr**2  ! Neumann

  allocate(A(1:nr,1:nr))
  A = 0.0_dp
  A(1,1) = diag(1)
  A(1,2) = diag_high(1)
  do i = 2, nr - 1
    A(i, i-1) = diag_low(i-1)
    A(i, i)   = diag(i)
    A(i, i+1) = diag_high(i)
  end do
  A(nr,nr-1) = diag_low(nr-1)
  A(nr,nr)   = diag(nr)

  allocate(Res(1:nr))
  Res = 0.0_dp

  ! Set boundary values
  phi0 = 0.0_dp / phi_ref ! Dirichlet
  ! phi0 = phi(1) ! Neumann
  ! phiF = 0.0_dp ! Dirichlet
  allocate(f0(j0:nv))
  f0 = 0.0_dp

  ! Output initial values
  call createPath()
  call setTimeFormat(nt)
  t = 0
  call writeOutputRef()
  call writeOutputF(t, dt, nr, r, nv, v, f_i_old)
  call writeOutputPhi(t, dt, nr, r, phi, E, n_e)
  call writeOutputMom(t, dt, nr, r, n_i, u_i, T_i, Zave)

  ! Main loop
  Temp_bc = Temp0
  u_bc    = u_bc0
  n_bc    = n_bc0
  do t = 1, nt
    if      (t >  t_bc0 .and. t <= t_bcF) then
      Temp_bc = (TempF - Temp0) / float(t_bcF - t_bc0)*(t - t_bc0) + Temp0
      u_bc    = (u_bcF - u_bc0) / float(t_bcF - t_bc0)*(t - t_bc0) + u_bc0
      n_bc    = (n_bcF - n_bc0) / float(t_bcF - t_bc0)*(t - t_bc0) + n_bc0
    else if (t >  t_bcF) then
      Temp_bc = TempF
      u_bc    = u_bcF
      n_bc    = n_bcF
    end if
    call writeOutputBoundary(t, dt, n_bc, u_bc, Temp_bc, T_to_Z(Temp_bc))
    f0(j0:nv) = n_bc / sqrt(PI*Temp_bc) * exp(-(v(j0:nv) - u_bc)**2 / Temp_bc)

    ! Boundary conditions
    ! r = r0, v>0
    f_i_old(1,j0:nv) = f0
    f_i(1,j0:nv) = f_i_old(1,j0:nv)
    T_i(1) = Temp_bc
    Zave(1) = T_to_Z(Temp_bc)
    ! r = rf, v<0
    f_i_old(nr,1:j0-1) = 0.0_dp
    f_i(nr,1:j0-1) = f_i_old(nr,1:j0-1)
    ! set edge velocities to 0
    f_i_old(:,1)  = 0.0_dp
    f_i_old(:,nv) = 0.0_dp
    ! Advect in the r direction
    !$omp parallel do
    do i = 1, nr
      ! Advect negative velocity
      if (i < nr) then
        f_i(i,1:j0-1) = f_i_old(i,1:j0-1) - v(1:j0-1)*dt/dr/r(i)**2*(r(i+1)**2*f_i_old(i+1,1:j0-1) - &
                                                                     r(i  )**2*f_i_old(i  ,1:j0-1))
      end if
      ! Advect positive velocity
      if (i >  1) then
        f_i(i,j0:nv)  = f_i_old(i, j0:nv) - v( j0:nv)*dt/dr/r(i)**2*(r(i  )**2*f_i_old(i  , j0:nv) - &
                                                                     r(i-1)**2*f_i_old(i-1, j0:nv))
      end if

      n_i(i) = sum(f_i(i,:)*dv)
      if (n_i(i) > 1.0e-10_dp) then
        u_i(i)  = sum(v(:)   *f_i(i,:)*dv) / n_i(i)
        E_i(i)  = sum(v(:)**2*f_i(i,:)*dv) / n_i(i)
        T_i(i)  = 2.0_dp*E_i(i) - 2.0_dp*u_i(i)**2
        Zave(i) = T_to_Z(T_i(i))

      else
        u_i(i)  = 0.0_dp
        T_i(i)  = 0.0_dp
        Zave(i) = 0.0_dp
      end if

    end do
    !$omp end parallel do

    ! Solve Poission (maximum number of iterations, break if convergence is reached before)
    do k = 1, 1000
      ! Store previous value
      phi_old = phi

      ! Diagonal matrix for Newton integration scheme
      diag           = -2.0_dp / dr**2 - n_e
      diag_low       =  1.0_dp / dr**2 - 1.0_dp / (r(2:nr)   * 2.0_dp * dr)
      diag_high      =  1.0_dp / dr**2 + 1.0_dp / (r(1:nr-1) * 2.0_dp * dr)
      diag(1)        =  1.0_dp ! Dirichlet
      diag_high(1)   =  0.0_dp ! Dirichlet
      ! diag_high(1)   =  2.0_dp / dr**2 - n_e(1) ! Neumann
      ! diag(nr)       =  1.0_dp ! Dirichlet
      ! diag_low(nr-1) =  0.0_dp ! Dirichlet
      diag_low(nr-1) =  2.0_dp / dr**2 - n_e(nr) ! Neumann
 
      ! Calculate charge density
      b = -(Zave*n_i - n_e)
      ! Apply boundary conditions
      b(1)  = phi0 ! Dirichlet
      ! b(nr) = phiF ! Dirichlet

      ! Calculate residual
      !$omp parallel workshare
      Res = -(MATMUL(A, phi_old) - b)
      !$omp end parallel workshare

      ! Iterate system
      call dgtsv(nr, 1, diag_low, diag, diag_high, Res, nr, info)
      phi = phi_old + Res
      ! phi0=phi(1)

      ! Calculate distribution of electrons
      n_e = Zave(1) * n_i(1) * exp((phi - phi0) / T_i(1))

      ! Check if the solution has converged
      phiConv = maxval(abs(Res),1)
      if (phiConv < 1.0e-3_dp) then
        exit

      end if

      ! ! Calculate new potential to ensure 0 current at the edge
      ! if (n_i(nr) > 1.0e-10_dp) then
      !   phiF = phi0 + T_i(1) * log((2.0_dp*sqrt(pi)*Zave(nr)*n_i(nr)*u_i(nr)) / (Zave(1)*n_i(1)*sqrt(m_i*T_i(1)/m_e)))
      !
      ! else
      !   phiF = phi(nr-5)
      !
      ! end if

    end do

    ! Calculate electric field
    E(1) = - (phi(2) - phi(1)) / dr ! Dirichlet
    ! E(1) = 0.0_dp ! Neumann
    !$omp parallel do
    do i = 2, nr-1
      E(i) = - 0.5_dp*(phi(i+1) - phi(i-1)) / dr

    end do
    !$omp end parallel do
    ! E(nr) = - (phi(nr) - phi(nr-1)) / dr ! Dirichlet
    E(nr) = 0.0_dp ! Neumann

    ! Update intermediate f
    f_i_old = f_i

    ! Advect in the v direction
    ! i = 1, v<0
    i = 1
    if (E(i) >= 0.0_dp) then
      f_i(i,2:j0-2) = f_i_old(i,2:j0-2) - Zave(i)*E(i)*dt/dv*(f_i_old(i,2:j0-2) - f_i_old(i,1:j0-3))
    else
      f_i(i,2:j0-2) = f_i_old(i,2:j0-2) - Zave(i)*E(i)*dt/dv*(f_i_old(i,3:j0-1) - f_i_old(i,2:j0-2))

    end if
    ! i = 2, nr-1; all v
    !$omp parallel do
    do i = 2, nr-1
      if (E(i) >= 0.0_dp) then
        f_i(i,2:nv-1) = f_i_old(i,2:nv-1) - Zave(i)*E(i)*dt/dv*(f_i_old(i,2:nv-1) - f_i_old(i,1:nv-2))
      else
        f_i(i,2:nv-1) = f_i_old(i,2:nv-1) - Zave(i)*E(i)*dt/dv*(f_i_old(i,3:nv)   - f_i_old(i,2:nv-1))

      end if

    end do
    !$omp end parallel do
    ! i = nr, v>=0
    i = nr
    if (E(i) >= 0.0_dp) then
      f_i(i,j0+1:nv-1) = f_i_old(i,j0+1:nv-1) - Zave(i)*E(i)*dt/dv*(f_i_old(i,j0+1:nv-1) - f_i_old(i,j0:nv-2))
    else
      f_i(i,j0+1:nv-1) = f_i_old(i,j0+1:nv-1) - Zave(i)*E(i)*dt/dv*(f_i_old(i,j0+2:nv)   - f_i_old(i,j0+1:nv-1))

    end if

    ! Reset values for next iteration
    f_i_old = f_i
    fCum_i  = fCum_i + f_i_old(rCum_index,:)

    ! Write output
    if (mod(t,everyOutput) == 0 .or. t == nt) then
      call writeOutputF(t, dt, nr, r, nv, v, f_i_old)
      call writeOutputPhi(t, dt, nr, r, phi, E, n_e)
      call writeOutputMom(t, dt, nr, r, n_i, u_i, T_i, Zave)
      call writeOutputFCum(t, dt, r(rCum_index), nv, v, fCum_i)

    end if

    ! Write progress
    if (mod(t,everyWrite) == 0) then
      write (*, '(I10, A, I10)' ) t, '/', nt
      write (*, '(A, ES0.4e3,","ES0.4e3)') 'phi max,min: ', maxval(phi)*phi_ref, minval(phi)*phi_ref

    end if

  end do

end program plasmaExpansion