!Equilibrium
module equilibrium
use precision, only: lk
implicit none 

public eqdata, &
       lsr, lsz, spdR, spdZ, psi_sep, &
       r_Eq, z_Eq, &
       ne_Eq, ni_Eq, ti_Eq, te_Eq, &
       psirz_Eq, B0MAG_Eq, B0R_Eq, B0Z_Eq, &
       zmax, rmax
private



integer  lsr, lsz, sdim
real(lk) psi_sep, spdR, spdZ
real(lk) zmax, rmax
real(lk),dimension(:),allocatable :: r_Eq, z_Eq
real(lk),dimension(:,:,:),allocatable :: B0mag_Eq, B0R_Eq, B0Z_Eq, &
                                         ne_Eq, ni_Eq, ti_Eq, te_Eq, &
                                         psirz_Eq
save
#ifndef GPU_UM
    !$acc declare create(psirz_Eq)
#endif

contains
!=======================================================================
subroutine eqdata
    use precision
    use global_parameters, only: gtcout, ulength, utime, rho0, betae, mype, &
                                 eden0, etemp0, r0, b0
    use read_in_parameters, only: eden0_input, etemp0_input,r0_input, b0_input, &
                                  tstep, aion, gamma0
    use spline_function, only: construct_spline2d
    implicit none
  
    integer i,icount,ierror
    real(lk) spdum(4)
    namelist /physical_parameters/ r0, b0, etemp0, eden0, ulength, utime, rho0, betae, zmax,rmax
  
! quadratic spline
    sdim = 9
! equilibrium and profile are represented using 2D-spline on (psi, theta) or 1D-spline on (psi)
    if(mype==0)then
#ifdef _FRC
        call FRC_equilibrium
#elif _TOKAMAK
        call Tokamak_equilibrium   
#endif
 
        call construct_spline2d(lsr,lsz,spdR,spdZ,ni_Eq,0,0)
        call construct_spline2d(lsr,lsz,spdR,spdZ,ne_Eq,0,0)
        call construct_spline2d(lsr,lsz,spdR,spdZ,ti_Eq,0,0)
        call construct_spline2d(lsr,lsz,spdR,spdZ,te_Eq,0,0)
        call construct_spline2d(lsr,lsz,spdR,spdZ,B0R_Eq,0,0)
        call construct_spline2d(lsr,lsz,spdR,spdZ,B0Z_Eq,0,0)
        call construct_spline2d(lsr,lsz,spdR,spdZ,B0Mag_Eq,0,0)
        call construct_spline2d(lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
    endif

    icount=1
    call MPI_BCAST(lsr,icount, mpi_Integer, 0, MPI_COMM_WORLD, ierror)
    call MPI_BCAST(lsz,icount, mpi_Integer, 0, MPI_COMM_WORLD, ierror)
    call MPI_BCAST(spdR,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(spdZ,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(psi_sep,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
  
    if(mype/=0)then
! allocate memory for equilibrium spline array
        allocate(r_Eq(lsr),z_Eq(lsz))
        allocate(ni_Eq(sdim,lsr,lsz),ne_Eq(sdim,lsr,lsz),ti_Eq(sdim,lsr,lsz),te_Eq(sdim,lsr,lsz))
        allocate(psirz_Eq(sdim,lsr,lsz),B0mag_Eq(sdim,lsr,lsz),B0R_Eq(sdim,lsr,lsz),B0Z_Eq(sdim,lsr,lsz))
    endif
  
  
! broadcast radial and poloildal spline fit of B field and (X,Z) coordinates for axisymmetric equil.
! broadcast fielddir so every plane gets the updated fielddir
    icount = lsr
    call MPI_BCAST(r_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    icount = lsz
    call MPI_BCAST(z_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
  
    icount = sdim*lsr*lsz
    call MPI_BCAST(psirz_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(B0mag_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(B0R_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(B0Z_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(ne_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(ni_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(te_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(ti_Eq,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
#ifndef GPU_UM
    !$acc update device(psirz_Eq)
#endif
  
#ifdef _FRC
    icount=1
    call MPI_BCAST(zmax,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    call MPI_BCAST(rmax,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
#endif


! broadcast equilibrium and proile data to all MPI processes
    icount=4
    spdum(1)=r0
    spdum(2)=b0
    spdum(3)=etemp0
    spdum(4)=eden0
    call MPI_BCAST(spdum,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
    r0=spdum(1)
    b0=spdum(2)
    etemp0=spdum(3)
    eden0=spdum(4)
  
! GTC unit of length and time: major radius R_0=1, proton cyclotron frequency omega_i=1
    ulength=r0 !read in major radius: length unit (unit=cm)
    utime=1.0_lk/(9580.0_lk*b0) ! time unit (unit=second) = inverse proton cyclotron frequency
! rho0=rho_s=c_s/omega_0 in GTC unit, c_s=sqrt(T_e0/m_p), omega_0= eB_0/m_p, on-axis T_e0, B_0
    rho0=102.0_lk*sqrt(etemp0)/(b0*ulength)
! tstep unit uses ion sound speed instead of proton sound speed
    tstep=sqrt(aion)*tstep/rho0
    gamma0=gamma0*rho0/sqrt(aion)
! electron beta
    betae=4.03e-11_lk*eden0*etemp0/(b0*b0)
  
    ti_Eq=ti_Eq*rho0*rho0
    te_Eq=te_Eq*rho0*rho0
  
    write(gtcout,physical_parameters)
    call FLUSH(gtcout)
  
    call eqplot

101 format(1p4e18.10)
102 format(6i4)
103 format(5e16.9)

end subroutine eqdata
!=======================================================================
#ifdef _FRC
subroutine FRC_equilibrium
    use global_parameters, only: size_scaling, etemp0, eden0, r0, b0
    use read_in_parameters, only: qion
    use LAMYRIDGE_interface, only: LR_eqdata, LR_wrapper_out, &
                                   lsr_LR, lsz_LR, &
                                   spdR_LR, spdZ_LR, psi_sep_LR, zmax_LR, rmax_LR, &
                                   r_Eq_LR, z_Eq_LR, &
                                   ni_Eq_LR, ne_Eq_LR, ti_Eq_LR, te_Eq_LR, &
                                   psirz_Eq_LR, B0mag_Eq_LR, B0R_Eq_LR, B0Z_Eq_LR, &
                                   etemp0_LR, eden0_LR, r0_LR, b0_LR

    call LR_eqdata
    call LR_wrapper_out
  
    lsr = lsr_LR
    lsz = lsz_LR

    spdR = spdR_LR
    spdZ = spdZ_LR
    psi_sep = psi_sep_LR

    zmax = zmax_LR
    rmax = rmax_LR
  
    etemp0 = etemp0_LR
    eden0 = eden0_LR
! size_scaling = 1 to realistic machine simulation, which can be changed in module.F90
    r0 = r0_LR*size_scaling
    b0 = b0_LR

    allocate(r_Eq(lsr), z_Eq(lsz), &
             ni_Eq(sdim,lsr,lsz), ne_Eq(sdim,lsr,lsz), ti_Eq(sdim,lsr,lsz), te_Eq(sdim,lsr,lsz), &
             psirz_Eq(sdim,lsr,lsz), B0mag_Eq(sdim,lsr,lsz), B0R_Eq(sdim,lsr,lsz), B0Z_Eq(sdim,lsr,lsz))
   
    r_Eq=r_Eq_LR
    z_Eq=z_Eq_LR
  
    ne_Eq(1,:,:) = ne_Eq_LR
    ni_Eq(1,:,:) = ne_Eq(1,:,:)/qion
    ti_Eq(1,:,:) = ti_Eq_LR
    te_Eq(1,:,:) = te_Eq_LR
  
    psirz_Eq(1,:,:) = psirz_Eq_LR
    B0mag_Eq(1,:,:) = B0mag_Eq_LR
    B0R_Eq(1,:,:) = B0R_Eq_LR
    B0Z_Eq(1,:,:) = B0Z_Eq_LR
  
end subroutine FRC_equilibrium
!=======================================================================
#elif _TOKAMAK
!=======================================================================
subroutine Tokamak_equilibrium
    use global_parameters, only: etemp0, eden0, r0, b0
    use read_in_parameters, only: qion
    use EFIT_interface, only: EFIT_eqdata, EFIT_wrapper_out, &
                              lsr_EFIT, lsz_EFIT, &
                              spdR_EFIT, spdZ_EFIT, &
                              r_Eq_EFIT, z_Eq_EFIT, &
                              psirz_Eq_EFIT, psi_sep_EFIT, &
                              r0_EFIT, b0_EFIT, & 
                              B0mag_Eq_EFIT, B0R_Eq_EFIT, B0Z_Eq_EFIT
    call EFIT_eqdata
    call EFIT_wrapper_out

    lsr = lsr_EFIT
    lsz = lsz_EFIT
    
    spdR = spdR_EFIT
    spdZ = spdZ_EFIT
    psi_sep = psi_sep_EFIT
  
    eden0 = 1.0e13       ! cm^-3
    etemp0 = 2000.0_lk   ! eV
    r0 = r0_EFIT
    b0 = b0_EFIT
   
    allocate(r_Eq(lsr), z_Eq(lsz), &
             ni_Eq(sdim,lsr,lsz), ne_Eq(sdim,lsr,lsz), ti_Eq(sdim,lsr,lsz), te_Eq(sdim,lsr,lsz), &
             psirz_Eq(sdim,lsr,lsz), B0mag_Eq(sdim,lsr,lsz), B0R_Eq(sdim,lsr,lsz), B0Z_Eq(sdim,lsr,lsz))

    r_Eq = r_Eq_EFIT
    z_Eq = z_Eq_EFIT
  
    ne_Eq(1,:,:) = 0.0_lk 
    ni_Eq(1,:,:) = ne_Eq(1,:,:)/qion
    ti_Eq(1,:,:) = 0.0_lk 
    te_Eq(1,:,:) = 0.0_lk 

    psirz_Eq(1,:,:) = psirz_Eq_EFIT
    B0mag_Eq(1,:,:) = B0mag_Eq_EFIT
    B0R_Eq(1,:,:) = B0R_Eq_EFIT
    B0Z_Eq(1,:,:) = B0Z_Eq_EFIT
end subroutine Tokamak_equilibrium
!=======================================================================
#endif
!=======================================================================
subroutine eqplot
    use global_parameters, only: mype
    implicit none

    integer ieq
 
    if(mype==0)then
        ieq=120 !open equilibiurm plot output file
        open(ieq,file='equilibrium.out',status='replace')
        write(ieq,101)lsr, lsz
        write(ieq,101)r_Eq, z_Eq, psirz_Eq(1,:,:), B0Mag_Eq(1,:,:), B0R_Eq(1,:,:), B0Z_Eq(1,:,:),&
                      ni_Eq(1,:,:), ti_Eq(1,:,:), ne_Eq(1,:,:), te_Eq(1,:,:)
        close(ieq)
    endif

101 format(' ',9g16.8)
end subroutine eqplot
!=======================================================================


end module equilibrium
