!> Interface with EFIT equilibrium
!> \author
!> original written by Animesh Kuley, modified by Jian Bao
!> \data
!> Oct 17, 2018
!========================================================================
!>  EFIT_interface module contains:
!>  Read in EFIT
!>
!>  Public Subroutines:
!>  EFIT_eqdata, EFIT_output
!========================================================================
!#define _EFIT

module EFIT_interface
use precision, only: lk
implicit none


public EFIT_eqdata, EFIT_wrapper_out, &
       lsr_EFIT, lsz_EFIT, spdR_EFIT, spdZ_EFIT, psi_sep_EFIT, &
       r0_EFIT, b0_EFIT, &
       r_Eq_EFIT, z_Eq_EFIT, &
       psirz_Eq_EFIT, B0mag_Eq_EFIT, B0R_Eq_EFIT, B0Z_Eq_EFIT
private

character(len=*), parameter :: efit_file="g158103.03050"

integer lsr_EFIT, lsz_EFIT, nbbbs_EFIT, limitr_EFIT
real(lk) rdim_EFIT, zdim_EFIT, rcentr_EFIT, rleft_EFIT, zmid_EFIT, &
         zmaxis_EFIT, psi_mag_EFIT, psi_sep_EFIT, &
         current_EFIT, psimx(2), xax(2), zax(2), &
         psisep_dummy_EFIT, xsep_EFIT, zsep_EFIT, &
         spdR_EFIT, spdZ_EFIT, r0_EFIT, b0_EFIT
real(lk),dimension(:),allocatable :: fpol_EFIT, pres_EFIT, ffprime_EFIT, pprime_EFIT, qrpsi_EFIT, &
                                     rbbbs_EFIT, zbbbs_EFIT, rlim_EFIT, zlim_EFIT, &
                                     r_Eq_EFIT, z_Eq_EFIT
real(lk),dimension(:,:),allocatable :: psirz_Eq_EFIT, B0mag_Eq_EFIT, B0R_Eq_EFIT, B0Z_Eq_EFIT
save
contains
!========================================================================
subroutine EFIT_eqdata
    integer isp
    isp=110

    call EFIT_read_dims(efit_file,isp)

    call EFIT_read_eqdata(isp)

    call EFIT_normalize

    call EFIT_magnetic_field

end subroutine EFIT_eqdata
!========================================================================
subroutine EFIT_wrapper_out

    integer,parameter :: ieq=120
    integer i, j
    real(lk) r(lsr_EFIT),conv

103 format(5e16.9)
104 format(2i5)
    open(ieq,file='EFIT_wrapper.out',status='replace')
  
! Equilibrium grid numbers
    write(ieq,104)lsr_EFIT
    write(ieq,104)lsz_EFIT
! Major radius   
    write(ieq,103)r0_EFIT
! Z of magnetic axis in m
    write(ieq,103)zmaxis_EFIT
! Poloidal flux at magnetic axis in Weber/rad
    write(ieq,103)psi_mag_EFIT

! Poloidal flux at limiter (last closed flux surface) in Weber/rad
    write(ieq,103)psi_sep_EFIT

! magnetic field strength on the axis
    write(ieq,103)b0_EFIT

! Minimum R of rectangular computational box
    write(ieq,103)rleft_EFIT
  
! Horizontal dimension of computational box
    write(ieq,103)rdim_EFIT
  
! Vertical dimension of computational box
    write(ieq,103)zdim_EFIT
  
! Number of boundary points
    write(ieq,104)nbbbs_EFIT
  
! Number of limiter points
    write(ieq,104)limitr_EFIT
  
! Poloidal current function
    do i=1,lsr_EFIT
        write(ieq,103)fpol_EFIT(i)
    enddo
  
! pressure profile
    do i=1,lsr_EFIT
        write(ieq,103)pres_EFIT(i)
    enddo

! poloidal current funtion radial derivative with respect to psi
    do i=1,lsr_EFIT
        write(ieq,103)ffprime_EFIT(i)
    enddo

! plasma pressure radial derivative with respect to psi
    do i=1,lsr_EFIT
        write(ieq,103)pprime_EFIT(i)
    enddo

! safety factor radial derivative with respect to psi
    do i=1,lsr_EFIT
        write(ieq,103)qrpsi_EFIT(i)
    enddo

! R of boundary points
    do i=1,nbbbs_EFIT
        write(ieq,103)rbbbs_EFIT(i)
    enddo
! Z of boundary points
    do i=1,nbbbs_EFIT
        write(ieq,103)zbbbs_EFIT(i)
    enddo
  
! R of surrounding limiter contour
    do i=1,limitr_EFIT
        write(ieq,103)rlim_EFIT(i)
    enddo
  
! Z of surrounding limiter contour
    do i=1,limitr_EFIT
        write(ieq,103)zlim_EFIT(i)
    enddo
  
! Polial flux function
    do i=1,lsr_EFIT
        do j=1,lsz_EFIT
            write(ieq,103)psirz_Eq_EFIT(i,j)
        enddo
    enddo
  
    close(ieq)
end subroutine EFIT_wrapper_out
!========================================================================
subroutine EFIT_read_dims(efile,isp)
    character(len=*), intent(in) :: efile
    integer, intent(in) :: isp
    character(len=52) :: cdum

    open(unit=isp,file=efile,action='read')
 
     !Dummy variable
     ! no. of grids in the R and Z directions
    read(isp,'(A,2i4)')cdum,lsr_EFIT,lsz_EFIT
     !R dimension (width) of the grid (poloidal c.s.) in metre
     !Z dimension (height)of the grid in metre
     !R (in m) of the vacuum toroidal B-field i.e. R of the geometric centre of the c.s.
     !R coord of leftmost end of the grid
     !z coord at the centre of the grid (i.e. at rcentr); taken as 0      
    read(isp,103)rdim_EFIT,zdim_EFIT,rcentr_EFIT,rleft_EFIT,zmid_EFIT
     !R of magnetic axis in m
     !Z of magnetic axis in m
     !Poloidal flux at magnetic axis in Weber/rad
     !Poloidal flux at limiter (last closed flux surface) in Weber/rad
     !Vacuum toroidal B-field in Tesla at rcentr
    read(isp,103)r0_EFIT,zmaxis_EFIT,psi_mag_EFIT,psi_sep_EFIT,b0_EFIT
     !Plasma current in ampere
     !Don't know what the other parameters are??
    read(isp,103)current_EFIT,psimx(1),psimx(2),xax(1),xax(2)
     !Few more parameters
    read(isp,103)zax(1),zax(2),psisep_dummy_EFIT,xsep_EFIT,zsep_EFIT
     
     
     !To print the dimensions to a file
     
    print *,"No.of horizontal grid points (lsr_EFIT) =",lsr_EFIT
    print *,"No.of vertical grid points (lsz_EFIT) =",lsz_EFIT
    print *,"Physical R dimension of the grid (rdim_EFIT) =",rdim_EFIT,"m"
    print *,"Physical` Z dimension of the grid (zdim_EFIT) =",zdim_EFIT,"m"
    print *,"R coord of the geometric centre (rcentr_EFIT) =",rcentr_EFIT,"m"
    print *,"R at the leftmost end of the grid (rleft_EFIT) =",rleft_EFIT,"m"
    print *,"Z coord at the centre of the grid (zmid_EFIT) =",zmid_EFIT
    print *,"R of magnetic axis (r0_EFIT) =",r0_EFIT,"m"
    print *,"Z of magnetic axis (zmaxis_EFIT) =",zmaxis_EFIT,"m"
    print *,"Pol. flux (Wb/rad) at the magneic axis (psi_mag_EFIT) =",psi_mag_EFIT
    print *,"Pol. flux(Wb/rad) flux at limiter (psi_sep_EFIT) =",psi_sep_EFIT
    print *,"Vacuum toroidal B-field (Tesla) at rcentr (b0_EFIT) =",b0_EFIT
    print *,"Plasma current in ampere (current_EFIT) =",current_EFIT

    print *, "cdum=",cdum,"lsr_EFIT=",lsr_EFIT,"lsz_EFIT=",lsz_EFIT
    print *, "psisep_dummy_EFIT=",psisep_dummy_EFIT
    print *, "xsep_EFIT=",xsep_EFIT
    print *, "zsep_EFIT=",zsep_EFIT
103 format(5e16.9)
end subroutine EFIT_read_dims
!========================================================================
subroutine EFIT_read_eqdata(isp)

    integer, intent(in) :: isp
    integer i, j

    allocate(fpol_EFIT(lsr_EFIT), pres_EFIT(lsr_EFIT), ffprime_EFIT(lsr_EFIT), &
             pprime_EFIT(lsr_EFIT), psirz_Eq_EFIT(lsr_EFIT,lsz_EFIT), qrpsi_EFIT(lsr_EFIT))


    !Poloidal current function (R*B_tor) in m-T on uniform flux grid
    read(isp,103)(fpol_EFIT(i), i = 1, lsr_EFIT)
    !Plasma pressure in J/m^3
    read(isp,103)(pres_EFIT(i), i = 1, lsr_EFIT)
    !FF'(psi) in (mT)/(Weber/rad) on uniform flux grid
    read(isp,103)(ffprime_EFIT(i), i = 1, lsr_EFIT)
    !p'(psi) in (mT)/(Weber/rad) on uniform flux grid
    read(isp,103)(pprime_EFIT(i), i = 1, lsr_EFIT)
    !Poloidal flux in Wb/rad on the rectangular grid
    read(isp,103)((psirz_Eq_EFIT(i,j), i = 1, lsr_EFIT), j = 1, lsz_EFIT)
    !q values on the flux surface from the axis to the boundary
    read(isp,103)(qrpsi_EFIT(i),i = 1, lsr_EFIT)


    !No. of boundary and limiter points
    read(isp,104)nbbbs_EFIT,limitr_EFIT

    allocate(rbbbs_EFIT(nbbbs_EFIT),zbbbs_EFIT(nbbbs_EFIT),rlim_EFIT(limitr_EFIT),zlim_EFIT(limitr_EFIT))
    !R and Z coords of boundary points in m
    read(isp,103)(rbbbs_EFIT(i),zbbbs_EFIT(i),i=1,nbbbs_EFIT)
    !R and Z coords of limiter points in m
    read(isp,103)(rlim_EFIT(i),zlim_EFIT(i),i=1,limitr_EFIT)

103 format(5e16.9)
104 format(2i5)

    close(isp)

end subroutine EFIT_read_eqdata
!========================================================================
!For normalizing the read in data
subroutine EFIT_normalize
    integer i    

    !The poloidal flux at the magnetic axis is subtracted from psirz
    !This sets the poloidal flux at the magnetic axis to zero
    psirz_Eq_EFIT(:,:)=psirz_Eq_EFIT(:,:)-psi_mag_EFIT
    !Data for equil. pol. flux plot
    !Pol. flux at limiter subtracted by psi_mag
    psi_sep_EFIT = psi_sep_EFIT-psi_mag_EFIT
    psi_sep_EFIT = abs(psi_sep_EFIT)

    !Change unit of flux from Weber to Tesla
    psirz_Eq_EFIT(:,:)=psirz_Eq_EFIT(:,:)/(abs(b0_EFIT)*r0_EFIT**2)
    psi_sep_EFIT=psi_sep_EFIT/(abs(b0_EFIT)*r0_EFIT**2)
    psi_mag_EFIT=abs(psi_mag_EFIT)/(abs(b0_EFIT)*r0_EFIT**2)

    !Leftmost R normalized by r0
    rleft_EFIT = rleft_EFIT/r0_EFIT
    !Horizontal dimension normalized by r0
    rdim_EFIT = rdim_EFIT/r0_EFIT
    !Vertical dimension normalized by r0
    zdim_EFIT = zdim_EFIT/r0_EFIT
    !Rlim, Zlim, R and Z at boundary, normalized by r0
    rlim_EFIT(:)=rlim_EFIT(:)/r0_EFIT
    zlim_EFIT(:)=zlim_EFIT(:)/r0_EFIT
    rbbbs_EFIT(:)=rbbbs_EFIT(:)/r0_EFIT
    zbbbs_EFIT(:)=zbbbs_EFIT(:)/r0_EFIT
    
    !Change unit of pol. current function from m-T to Tesla
    fpol_EFIT(:)=abs(fpol_EFIT(:))/r0_EFIT
    fpol_EFIT(:)=fpol_EFIT(:)/abs(b0_EFIT) !Normalizing with the B-field on B-axis
    !absolute B-filed on B-axis
    b0_EFIT=abs(b0_EFIT)

!======================
    spdR_EFIT = rdim_EFIT/(real(lsr_EFIT) - 1.0_lk)
    spdZ_EFIT = zdim_EFIT/(real(lsz_EFIT) - 1.0_lk)

    allocate(r_Eq_EFIT(lsr_EFIT), z_Eq_EFIT(lsz_EFIT))   
 
    do i = 1, lsr_EFIT
        r_Eq_EFIT(i) = rleft_EFIT + spdR_EFIT*(real(i) - 1.0_lk)
    enddo

    do i = 1, lsz_EFIT
        z_Eq_EFIT(i) = - 0.5_lk*zdim_EFIT + spdZ_EFIT*(real(i) - 1.0_lk)
    enddo


    b0_EFIT = b0_EFIT*1e4    ! gauss
    r0_EFIT = r0_EFIT*1e2    ! cm

    print *, "b0_EFIT = ", b0_EFIT
    print *, "r0_EFIT = ", r0_EFIT

end subroutine EFIT_normalize
!========================================================================
subroutine EFIT_magnetic_field

    allocate(B0mag_Eq_EFIT(lsr_EFIT, lsz_EFIT), B0R_Eq_EFIT(lsr_EFIT, lsz_EFIT), &
             B0Z_Eq_EFIT(lsr_EFIT, lsz_EFIT))
! Todo
    B0mag_Eq_EFIT = 0.0_lk
    B0R_Eq_EFIT = 0.0_lk
    B0Z_Eq_EFIT = 0.0_lk
   
end subroutine EFIT_magnetic_field
!========================================================================
end module EFIT_interface
!===============================================================
!      check EFIT equilibrium interface
!      edison/cori: ftn -o efit_test module.F90 EFIT_interface.F90
!===============================================================
#ifdef _EFIT
program unittest_EFIT_interface
    use EFIT_interface, only: EFIT_eqdata, EFIT_wrapper_out
    implicit none

    call EFIT_eqdata
    call EFIT_wrapper_out

end program
#endif
