!> Interface with LAMYRIDGE FRC equilibrium
!> \author
!> Jian Bao
!> \date (created)
!> Feb 12, 2018 
!> \date (modified)
!> Oct 19, 2018
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!>  LAMYRIDGE_interface module contains:
!>   Read in LAMYRIDGE equilibrium using taehdf5.F90 from Tri Alpha Energy Inc.
!>   (taehdf5.F90 requires libraries: cray-netcdf-hdf5parallel and cray-hdf5-parallel)
!>   Public Subroutines:
!>   LR_eqdata (read in equilibrium quantites from lamyridge)
!>   LR_output(only for unit test)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!#define _LR


module LAMYRIDGE_interface
use precision, only: lk
use taehdf5, only: h5read_data
implicit none

public LR_eqdata, LR_wrapper_out, LR_common_allocation, &
       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
    
private

    !  character(len=*), parameter :: LAMYRIDGE_file="symmetric_LRstrip.h5"
character(len=*), parameter :: LAMYRIDGE_file="symmetric_c2prof_eta=2.0_R0Lp=5_Te=80eV_flat-exponential.h5"
     ! character(len=*), parameter :: LAMYRIDGE_file="dump000.h5"

integer lsr_LR, lsz_LR
real(lk) eden0_LR, etemp0_LR, b0_LR, r0_LR
real(lk) spdR_LR, spdZ_LR, psi_sep_LR, zmax_LR, rmax_LR, psifrcmin
real(lk),dimension(:),allocatable :: r_Eq_LR, z_Eq_LR
real(lk),dimension(:,:),allocatable :: ni_Eq_LR, ne_Eq_LR, ti_Eq_LR, te_Eq_LR, &
                                         psirz_Eq_LR, B0mag_Eq_LR, B0R_Eq_LR, B0Z_Eq_LR
save
contains

!=======================================================================
! Subroutine LR_eqdata is used for reading in LAMYRIDGE equilibrium from taehdf.F90 in cylindrical coordinates
subroutine LR_eqdata

    call LAMYRIDGE_read_dims(LAMYRIDGE_file)

    call LR_common_allocation

    call LAMYRIDGE_read_eq(LAMYRIDGE_file) !read FRC equilibrium in RZ coordinates from LamyRidge hdf5 file

    call LAMYRIDGE_norm  ! normalization 

end subroutine LR_eqdata
!======== output lamyridge equilibrium quatities for unit test =========
subroutine LR_wrapper_out

    integer,parameter :: ieq=120

101 format(' ',9g16.8)
    open(ieq,file='LR_wrapper.out',status='replace')

    write(ieq,101)lsr_LR,lsz_LR
    write(ieq,101)r_Eq_LR,z_Eq_LR,psirz_Eq_LR(:,:),B0Mag_Eq_LR(:,:),B0R_Eq_LR(:,:),B0Z_Eq_LR(:,:),ni_Eq_LR(:,:),ti_Eq_LR(:,:),ne_Eq_LR(:,:),te_Eq_LR(:,:)
    close(ieq)

end subroutine LR_wrapper_out

!==== lsr and lsz are the grid numbers in radial and Z directions ======
subroutine LAMYRIDGE_read_dims(lrfile)

    character(len=*), intent(in):: lrfile

    call h5read_data(lrfile,lsr_LR,'mesh/nr')
    call h5read_data(lrfile,lsz_LR,'mesh/nz')

end subroutine LAMYRIDGE_read_dims

!==================== allocate FRC equilibrium arrays ==================
subroutine LR_common_allocation
   
    allocate(r_Eq_LR(lsr_LR),z_Eq_LR(lsz_LR),ni_Eq_LR(lsr_LR,lsz_LR),ne_Eq_LR(lsr_LR,lsz_LR),&
    ti_Eq_LR(lsr_LR,lsz_LR),te_Eq_LR(lsr_LR,lsz_LR), psirz_Eq_LR(lsr_LR,lsz_LR),&
    B0mag_Eq_LR(lsr_LR,lsz_LR),B0R_Eq_LR(lsr_LR,lsz_LR),B0Z_Eq_LR(lsr_LR,lsz_LR))

    r_Eq_LR=0.0_lk
    z_Eq_LR=0.0_lk
    psirz_Eq_LR=0.0_lk
    ni_Eq_LR=0.0_lk
    ne_Eq_LR=0.0_lk
    ti_Eq_LR=0.0_lk
    te_Eq_LR=0.0_lk

    B0R_Eq_LR=0.0_lk
    B0Z_Eq_LR=0.0_lk
    B0Mag_Eq_LR=0.0_lk

end subroutine LR_common_allocation

!==== read in density, temperature, poloidal magnetic flux from LAMYRIDGE equilibrium====
subroutine LAMYRIDGE_read_eq(lrfile)

    real(8),dimension(:,:),allocatable :: tmp
    real(8),dimension(:),allocatable :: rtmp,ztmp
    character(len=*), intent(in):: lrfile

    allocate(rtmp(lsr_LR),ztmp(lsz_LR),tmp(lsr_LR,lsz_LR))

! Read in TAE hdf5 files, on-point values only, no spline coefficients
    call h5read_data(lrfile,rtmp,'mesh/r')
    r_Eq_LR(:)=real(rtmp(:),lk)

    call h5read_data(lrfile,ztmp,'mesh/z')
    z_Eq_LR(:)=real(ztmp(:),lk)

    call h5read_data(lrfile,tmp,'vacuum/psi')
    psirz_Eq_LR(:,:)=real(tmp(:,:),lk)
  
    call h5read_data(lrfile,tmp,'state/en01')
  !    call h5read_data(lrfile,tmp,'state/en')
    ni_Eq_LR(:,:)=real(tmp(:,:),lk)!*3.2373_lk
  
    call h5read_data(lrfile,tmp,'state/en')
    ne_Eq_LR(:,:)=real(tmp(:,:),lk)!*3.2373_lk
  
    call h5read_data(lrfile,tmp,'state/tme')
  !    call h5read_data(lrfile,tmp,'state/te')
    te_Eq_LR(:,:)=real(tmp(:,:),lk)!*0.5_lk
  
    call h5read_data(lrfile,tmp,'state/tm01')
    ti_Eq_LR(:,:)=real(tmp(:,:),lk)*1.1_lk!*2.85_lk
  
  !    call h5read_data(lrfile,tmp,'state/tm')
  !    ti_Eq_LR(:,:)=real(tmp(:,:),lk)-te_Eq_LR(:,:)
  
    call h5read_data(lrfile,tmp,'state/bbr')
    B0R_Eq_LR(:,:)=real(tmp(:,:),lk)
  
    call h5read_data(lrfile,tmp,'state/bbz')
    B0Z_Eq_LR(:,:)=real(tmp(:,:),lk)
  
    B0Mag_Eq_LR(:,:)=sqrt(B0R_Eq_LR(:,:)*B0R_Eq_LR(:,:)+B0Z_Eq_LR(:,:)*B0Z_Eq_LR(:,:))

    deallocate(rtmp,ztmp,tmp)


end subroutine LAMYRIDGE_read_eq
!====================== LR normalization =========================
subroutine LAMYRIDGE_norm

    integer :: LReq=117, z0_index, r_sep_index(1), r_axis_index(1), i, rmax_index(1)
    real(lk) :: B0_norm,ne_norm,te_norm,x_norm, psiout_min

    real(lk),dimension(:),allocatable :: psitmp_1d 
!=================================================================
!             Calculation of normalized quantities
!=================================================================
    z0_index = lsz_LR/2 + 1
    r_axis_index = minloc(psirz_Eq_LR(:,z0_index))

    allocate(psitmp_1d(lsr_LR - r_axis_index(1)))
    psitmp_1d = psirz_Eq_LR(r_axis_index(1) + 1:lsr_LR, z0_index)
    r_sep_index = minloc(abs(psitmp_1d))
    r_sep_index(1) = r_sep_index(1) + r_axis_index(1)

    r0_LR = r_Eq_LR(r_axis_index(1))*1.0e2
    b0_LR = B0Mag_Eq_LR(r_sep_index(1),z0_index)*1.0e4
    eden0_LR = maxval(ne_Eq_LR(:,:))*1.0e-6
    etemp0_LR = maxval(te_Eq_LR(:,:))

!=================================================================
!                         Normalization
!=================================================================
!Tesla to gauss, m^-3 to cm^-3, m to cm
    B0_norm=b0_LR*1.0e-4
    ne_norm=eden0_LR*1.0e6
    te_norm=etemp0_LR
    x_norm=r0_LR*1.0e-2

    r_Eq_LR = r_Eq_LR/x_norm
    z_Eq_LR = z_Eq_LR/x_norm
    ne_Eq_LR = ne_Eq_LR/ne_norm
    ni_Eq_LR = ni_Eq_LR/ne_norm
    te_Eq_LR = te_Eq_LR/te_norm
    ti_Eq_LR = ti_Eq_LR/te_norm
    psirz_Eq_LR = psirz_Eq_LR/(B0_norm*x_norm*x_norm)
    B0R_Eq_LR = B0R_Eq_LR/B0_norm
    B0Z_Eq_LR = B0Z_Eq_LR/B0_norm
    B0Mag_Eq_LR = B0Mag_Eq_LR/B0_norm


    psifrcmin = abs(minval(psirz_Eq_LR(:,:)))
    psi_sep_LR =0.0_lk

    spdR_LR = maxval(r_Eq_LR)/real(lsr_LR-1)

    spdZ_LR = (maxval(z_Eq_LR)-minval(z_Eq_LR))/real(lsz_LR-1)

    zmax_LR = maxval(z_Eq_LR)

    rmax_LR = maxval(r_Eq_LR)
101 format(' ',9g16.8)
    open(LReq,file='LReq_parameters.out',status='replace')
        write(LReq,*)"psifrcmin",psifrcmin
        write(LReq,*)'FRC normalization: B0,ne,te,x'
        write(LReq,101)B0_norm,ne_norm,te_norm,x_norm
        write(LReq,*)'FRC_LRfile:ne,Te,ni,Ti:'
        write(LReq,101)maxval(ne_Eq_LR),maxval(te_Eq_LR),maxval(ni_Eq_LR),maxval(ti_Eq_LR)
        write(LReq,*)"etemp0_LR, eden0_LR, b0_LR, r0_LR="
        write(LReq,101)etemp0_LR,eden0_LR,b0_LR,r0_LR
        write(LReq,*)"lsr_LR,lsz_LR,spdR_LR,spdZ_LR="
        write(LReq,101)lsr_LR,lsz_LR,spdR_LR,spdZ_LR
        write(LReq,*)"rmaxval,rminval,zmaxval,zminval="
        write(LReq,101)maxval(r_Eq_LR),minval(r_Eq_LR),maxval(z_Eq_LR),minval(z_Eq_LR)
        write(LReq,101)"zmax_LR=",zmax_LR
        write(LReq,101)"rmax_LR=",rmax_LR

        write(LReq,101) "b0_LR =", b0_LR, "Gauss"
        write(LReq,101) "r0_LR =", r0_LR, "cm"
        write(LReq,101) "z0_index =", z0_index
        write(LReq,101) "r_sep_index =", r_sep_index(1)
        write(LReq,101) "r_axis_index =", r_axis_index(1)
        write(LReq,101) "eden0_LR =",eden0_LR, "cm^-3"
        write(LReq,101) "etemp0_LR =",etemp0_LR, "eV"
        
        psiout_min=minval(psirz_Eq_LR(lsr_LR,:))        
        write(LReq,*) "minimal psirz_Eq(lsr_LR,:)=",psiout_min
        rmax_index=minloc(abs(psiout_min - psirz_Eq_LR(:,z0_index)))
        write(LReq,*) "maximal radial index=",rmax_index(1)
        write(LReq,*) "radial position of minimal psirz_Eq_out-midplane=",r_Eq_LR(rmax_index(1))
        write(LReq,*) "normalized psirz_Eq_out-midplane ="
        do i = r_axis_index(1), lsr_LR
            write(LReq,*)i,psirz_Eq_LR(i,z0_index)
        enddo
    close(LReq)


end subroutine LAMYRIDGE_norm
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
end module LAMYRIDGE_interface


!=================================================================
!    check lamyridge equilibrium interface
!    edison/cori: ftn -o lrtest module.F90 taehdf5.F90 LAMYRIDGE_interface.F90
!=================================================================
#ifdef _LR
program unittest_LR_interface
use LAMYRIDGE_interface,only:LR_eqdata,LR_wrapper_out
implicit none

    call LR_eqdata
    call LR_wrapper_out

end program
#endif

