! This file is part of GTC version 4.4 
! GTC version 4.4 is released under the 3-Clause BSD license:

! Copyright (c) 2002,2010,2016, GTC Team (team leader: Zhihong Lin, zhihongl@uci.edu)
! All rights reserved.

! Redistribution and use in source and binary forms, with or without 
! modification, are permitted provided that the following conditions are met:

! 1. Redistributions of source code must retain the above copyright notice, 
!    this list of conditions and the following disclaimer.

! 2. Redistributions in binary form must reproduce the above copyright notice, 
!    this list of conditions and the following disclaimer in the documentation 
!    and/or other materials provided with the distribution.

! 3. Neither the name of the GTC Team nor the names of its contributors may be 
!    used to endorse or promote products derived from this software without 
!    specific prior written permission.
! ==============================================================================


!> \file
!> \brief load VMEC equilibrium file 
!> \author
!> Lei Shi
!> \date
!> Feb. 26, 2018

!######################################################################

module vmec
  use system_env, only: GTC_ABORT
  use precision, only: lk
  use utility, only: mstat, std_err
  use global_parameters, only: gtcout

  implicit none
  private
  public open_vmec_file, close_vmec_file, read_vmec_dims, read_vmec_eq, calc_and_broadcast_vmec, VMEC_FILE
  ! Default VMEC file name
  character(len=*), parameter :: VMEC_FILE = './spdata.dat'
  ! Default file ID for opening VMEC file
  integer, parameter :: isp=110

  ! total toroidal harmonic number contained in VMEC file
  integer :: ndim_total
  ! coordinates direction convention parameter
  integer :: sgn

  ! VMEC raw toroidal harmonic data
  real(lk),dimension(:,:,:),allocatable :: bcn,bsn,xcn,xsn,zcn,zsn,fcn,fsn
  integer,dimension(:),allocatable :: ntor

  ! VMEC temporary spline arrays
  real(lk),dimension(:,:,:,:),allocatable :: sptmp3d

  save

  contains

    !> \brief open the specified vmec file
    !> @param[in] vmecfile
    !> string, path to the VMEC file
    subroutine open_vmec_file(vmecfile)
      character(len=*), intent(in) :: vmecfile
      character(len=1024) :: errmsg
      open(isp, file=vmecfile, status='old', ACTION='READ', IOSTAT=mstat)
      if (mstat /= 0) then
        call std_err("Failed to read VMEC file:"//trim(vmecfile))
        write(errmsg, "(A7, I5)") "IOSTAT:", mstat
        call std_err(errmsg)
        call GTC_ABORT
      endif
    end subroutine open_vmec_file

    !> \brief close the opened vmec file
    subroutine close_vmec_file
      character(len=1024) :: errmsg
      close(isp, IOSTAT=mstat)
      if (mstat /= 0) then
        call std_err("Failed to close VMEC file:")
        write(errmsg, "(A7, I5)") "IOSTAT:", mstat
        call std_err(errmsg)
        call GTC_ABORT
      endif
    end subroutine close_vmec_file

    !> \brief read in the dimension info from VMEC data file
    subroutine read_vmec_dims
      use equilibrium, only: lsp, lst, ndim, psiw, ped

      ! temp character array holding commented lines
      character(len=99) cdum
      ! local variables
      integer :: nfp


      ! read pass the header line
      read(isp,*) cdum

      ! start reading dimension parameters
      read(isp,102) lsp,lst
      read(isp,101) psiw
      read(isp,102) ndim_total,nfp

      ! convert the VMEC coordinate convention to GTC convention
      sgn=1
      if(psiw<0)sgn=-1 ! poloidal flux determines the direction of axes
      psiw=sgn*psiw
      ped=psiw
      lst=lst+1
      ! First check if ndim is specified by user, or using the read-in
      ! total mode number
      if (ndim==0) then
       ndim = ndim_total
      ! then check if the used n mode number exceeds the total read-in mode
      ! number
      elseif (ndim > ndim_total) then
        write(gtcout,*) "ndim exceeds ndim_total, only", ndim_total,&
          "toroidal harmonics are read and used."
        ndim = ndim_total
      endif
      write(gtcout,*) "VMEC DIMENSION READ:"
      write(gtcout,*) "lsp=",lsp,"lst=",lst, "ndim_total=", ndim_total,&
        "psiw=", psiw

101 format(1p4e18.10)
102 format(6i4)
    endsubroutine read_vmec_dims

    subroutine allocate_vmec_tmparray
      use utility, only: check_allocation
      use equilibrium, only: lst, lsp

      allocate(bcn(lsp,lst,ndim_total),bsn(lsp,lst,ndim_total),xcn(lsp,lst,ndim_total),&
        xsn(lsp,lst,ndim_total),zcn(lsp,lst,ndim_total),zsn(lsp,lst,ndim_total),&
        ntor(ndim_total),fcn(lsp,lst,ndim_total),fsn(lsp,lst,ndim_total), STAT=mstat)
      call check_allocation(mstat, "bcn and others in eqdata.F90")

    endsubroutine allocate_vmec_tmparray

    subroutine deallocate_vmec_tmparray
      deallocate(bcn,bsn,xcn,xsn,zcn,zsn,fcn,fsn,ntor)
    endsubroutine deallocate_vmec_tmparray


    !> \brief read in equilibrium data from the provided VMEC data file

    subroutine read_vmec_eq
      use global_parameters
      use equilibrium
      use spline_function, only: construct_spline1d
      implicit none

      integer i,j,n
      real(lk) :: b00,x00,dum,temparray(lst)

      ! allocate memory for temporary arrays
      call allocate_vmec_tmparray

    ! read toroidal mode numbers
      read(isp,102) ntor(1:ndim_total)
    ! read toroidal harmonics of B,X,Z,dPHI=PHI-phi_boozer
      do n = 1, ndim_total
         do i = 1, lsp
               read(isp,101) temparray(:)
               bcn(i,1:lst,n) = temparray(:)
               read(isp,101) temparray(:)
               bsn(i,1:lst,n) = temparray(:)
               read(isp,101) temparray(:)
               xcn(i,1:lst,n) = temparray(:)
               read(isp,101) temparray(:)
               xsn(i,1:lst,n) = temparray(:)
               read(isp,101) temparray(:)
               zcn(i,1:lst,n) = temparray(:)
               read(isp,101) temparray(:)
               zsn(i,1:lst,n) = temparray(:)
               read(isp,101) temparray(:)
               fcn(i,1:lst,n) = temparray(:)
               read(isp,101) temparray(:)
               fsn(i,1:lst,n) = temparray(:)
         enddo
      enddo
      do i=1,lsp
        read(isp,101)qpsi(1,i) !q
        read(isp,101)gpsi(1,i) !g-current
        read(isp,101)cpsi(1,i) !i-current (plasma current)
        read(isp,101)rpsi(1,i) !minor radius
        read(isp,101)torpsi(1,i) !toroidal flux
      enddo
      close(isp)

    ! Coordinates conversion from VMEC to GTC
    ! Notes can be found in Doc/Developers_Note/notes_on_3d_spline_of_toroidal3d_in_GTC.pdf

    ! VMEC (psi, theta, zeta) is left-handed
    ! first change it to righthanded by switching z->-z
      fcn=-fcn
      fsn=-fsn
      torpsi=-torpsi
      cpsi=-cpsi
      qpsi=-qpsi

    ! VMEC to GTC coordinates conversion depends on the field configuration
    ! sgn (direction of I_P/B_theta)
    ! torpsi (direction of B_T) 

    ! fielddir | B_T | I_P
    !==========|=====|====
    !     0    | out | out
    !     1    |  in | out
    !     2    |  in | in
    !     3    | out | in
      if((sgn>0).and.(torpsi(1,lsp)>0))fielddir=0
      if((sgn>0).and.(torpsi(1,lsp)<0))fielddir=1 
      if((sgn<0).and.(torpsi(1,lsp)<0))fielddir=2
      if((sgn<0).and.(torpsi(1,lsp)>0))fielddir=3

      write(gtcout,*)"In zeta out system: sgn(psip)=",sgn,"torpsi(1,lsp)=",torpsi(1,lsp),&
                     "gpsi(1,1)=",gpsi(1,1),'fielddir=',fielddir

    ! determine signs of toroidal direction dependent quantities according to GTC zeta direction
    ! fielddir=0,1 : GTC zeta is along VMEC right-handed zeta, no addition change needed
    ! fielddir=2,3 : GTC zeta opposite to VMEC right-handed zeta, sign flip is needed
      if((fielddir==2) .or. (fielddir==3)) then
        fcn=-fcn
        fsn=-fsn
        torpsi=-torpsi
        gpsi=-gpsi
        cpsi=-cpsi
      endif

    ! In fielddir=2,3, GTC theta direction is also opposite to VMEC theta, i.e. theta_gtc = 2pi-theta_vmec
    ! This conversion is done in spline_3d_dep in function.F90, where function values are evaluated on equilibrium grids. 
    ! Notes can be found in Doc/Developers_Note/notes_on_3d_spline_of_toroidal3d_in_GTC.pdf

    ! Coordinate conversion finished.

    ! Normalization with GTC units
      x00=sum(xcn(1,:,1))/lst
      b00=sum(bcn(1,:,1))/lst

      write(gtcout,*)"X(0,0)=",x00,"B(0,0)=",b00,'ndim=',ndim
      flush(gtcout)

    ! make sure all quantities from spdata are normalized to gtc units R_0=B_0=1
    ! poloidal flux values get scaled by B_0*R_0^2

      psiw=psiw/(b00*x00*x00)
      ped=ped/(b00*x00*x00)

      spdtheta=2.0_lk*pi/real(lst-1,lk)
      spdpsi=ped/real(lsp-1,lk)

    ! toroidal flux also scaled by B_0*R_0^2
      torpsi=torpsi/(b00*x00*x00)

    ! current normalized by R_0*B_0
      gpsi=gpsi/(b00*x00)
      cpsi=cpsi/(b00*x00)

    ! Jacobian normalized by R_0*B_0
    !  gsp=gsp*b00/x00

    ! normalize B-field by on-axis value
      bcn=bcn/b00
      bsn=bsn/b00
    ! normalize spatial vars normalized by major radius
      xcn=xcn/x00
      xsn=xsn/x00

      zcn=zcn/x00
      zsn=zsn/x00

      rpsi=rpsi/x00

      rpsi(1,1)=0.0_lk

    ! construct 1D poloidal splines
      call construct_spline1d(lsp,spdpsi,qpsi,0)
      call construct_spline1d(lsp,spdpsi,gpsi,0)
      call construct_spline1d(lsp,spdpsi,cpsi,0)
      call construct_spline1d(lsp,spdpsi,torpsi,0)
      call construct_spline1d(lsp,spdpsi,rpsi,1)

      ! rd is the 2D current spline array used only in diagnostic
      rd = 0.0_lk
      do i=1, lsp
        rd(1:3,i,1)=cpsi(1:3,i)
        do j=2, lst
          rd(1:3, i, j) = rd(1:3,i,1)
        enddo
      enddo

101 format(1p4e18.10)
102 format(6i4)

    endsubroutine read_vmec_eq


    !> \brief calculate the 3D equilibrium quantities, generates the 3D spline
    !> coefficients, and broadcast the proper part to all MPI processes
    !> This subroutine populates the global arrays:
    subroutine calc_and_broadcast_vmec
      use system_env
      use utility, only: mstat, check_allocation 
      use precision, only: mpi_Rsize
      use equilibrium, only: lsp, lst, lszeta, nzsp_sec, spdim, bsp, xsp, zsp,&
        fsp
      use global_parameters, only: mype, toroidal_domain_location

      integer :: icount

      ! allocate temporary spline array
      allocate(sptmp3d(spdim,lsp,lst,lszeta), STAT=mstat)
      call check_allocation(mstat, "sptmp3d in VMEC_Interface.F90")

      icount=spdim*lsp*lst*lszeta
      !construct 3D spline
      if(mype==0) call generate_vmec_spline(bcn,bsn)
      call MPI_BCAST(sptmp3d,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
      ! assign each toroidal_domain_location with its own values of spline coeffs
      bsp(:,:,:,:)=sptmp3d(:,:,:,toroidal_domain_location*nzsp_sec+1:(toroidal_domain_location+1)*nzsp_sec)

      if(mype==0) call generate_vmec_spline(xcn,xsn)
      call MPI_BCAST(sptmp3d,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
      xsp(:,:,:,:)=sptmp3d(:,:,:,toroidal_domain_location*nzsp_sec+1:(toroidal_domain_location+1)*nzsp_sec)

      if(mype==0) call generate_vmec_spline(zcn,zsn)
      call MPI_BCAST(sptmp3d,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
      zsp(:,:,:,:)=sptmp3d(:,:,:,toroidal_domain_location*nzsp_sec+1:(toroidal_domain_location+1)*nzsp_sec)

      if(mype==0) call generate_vmec_spline(fcn,fsn)
      call MPI_BCAST(sptmp3d,icount,mpi_Rsize,0,MPI_COMM_WORLD,ierror)
      fsp(:,:,:,:)=sptmp3d(:,:,:,toroidal_domain_location*nzsp_sec+1:(toroidal_domain_location+1)*nzsp_sec)

      deallocate(sptmp3d)
      if(mype==0) call deallocate_vmec_tmparray

    endsubroutine calc_and_broadcast_vmec

    !> \brief generate 3D spline coefficients on the whole torus based on
    !> the given toroidal harmonics

    !> @param[in] ycn 
    !> coefficients in front of the cosine terms
    !> @param[in] ysn
    !> coefficients in front of the sine terms

    !> output: sptmp3d
    subroutine generate_vmec_spline(ycn, ysn)
      use precision
      use global_parameters
      use equilibrium, only: lsp, lst, lszeta, ndim, spdpsi, spdtheta, spdzeta
      use spline_function, only: construct_spline3d
      implicit none

      integer i,j,k,n, bcx, bcy, bcz

      real(lk),intent(in) :: ycn(lsp,lst,ndim),ysn(lsp,lst,ndim)
      real(lk) :: zdum, cosnz, sinnz

      sptmp3d(1,:,:,:) = 0.0_lk

      ! Determine the boundary condition based on the grid numbers
      ! x direction is always psi, needs BC No. 1
      bcx = 1
      ! lst is the total data points in poloidal direction
      if (mod(lst,2)==0) then
         bcy=2
      else
         bcy=0
         write(gtcout,*) "SPLINE WARNING: total poloidal spline grid number is odd, periodic spline not available. Non-periodic spline is used.", "lst = ", lst 
      endif
      ! lszeta is the total poloidal planes, thus is one less than the toroidal data points
      if (mod(lszeta-1,2)==0) then
         bcz=0
         write(gtcout,*) "SPLINE WARNING: total toroidal section number is even, periodic spline in toroidal direction is not available. Non-periodic spline is used.", "lszeta = ", lszeta
      else
         bcz=2
      endif
      
    ! Now generate the 3D data and create spline coefficients
    ! for each toroidal grid the value is calculated as sum over cos and sin harmonics
    ! negative sign in zeta due to VMEC convention, sgn flipps directions of theta and zeta
    ! to ensure poloidal flux is positive
      if(sgn>0)then
        do k = 1, lszeta
          zdum= -real(k-1,lk)*spdzeta 
          do n = 1, ndim
            cosnz = cos(real(ntor(n),lk)*zdum)
            sinnz = sin(real(ntor(n),lk)*zdum)
            do j = 1, lst
                do i = 1, lsp
                  sptmp3d(1,i,j,k)= sptmp3d(1,i,j,k)+ycn(i,j,n)*cosnz + ysn(i,j,n)*sinnz
                enddo
            enddo
          enddo
        enddo
      else !VMEC original z-direction
        do k = 1, lszeta
          zdum= real(k-1,lk)*spdzeta ! changed sign of zeta
          do n = 1, ndim
            cosnz = cos(real(ntor(n),lk)*zdum)
            sinnz = sin(real(ntor(n),lk)*zdum)
            do j = 1, lst
              do i = 1, lsp
                sptmp3d(1,i,j,k)= sptmp3d(1,i,j,k)+ycn(i,lst-j+1,n)*cosnz + ysn(i,lst-j+1,n)*sinnz
              enddo
            enddo
          enddo
        enddo
      endif

      call construct_spline3d(lsp, lst, lszeta, spdpsi, spdtheta, spdzeta,&
      sptmp3d, bcx, bcy, bcz)

    end subroutine generate_vmec_spline

end module vmec
