module interp_eqprof
  use const_and_precisions, only : wp_
  implicit none
  
! equidata
  ! === 1D array Fpol(psi), q(psi), rhotor(psi) ====================
  !INTEGER, SAVE :: npsieq
  !REAL(wp_), DIMENSION(:), ALLOCATABLE, SAVE :: rhopnr
  
  ! === 1D array plasma boundary Rbnd_i, Zbnd_i ====================
  INTEGER, SAVE :: nbbbs
  REAL(wp_), DIMENSION(:), ALLOCATABLE, SAVE :: rbbbs,zbbbs

  REAL(wp_), SAVE :: btaxis,rmaxis,zmaxis
  REAL(wp_), SAVE :: rmnm,rmxm,zmnm,zmxm,dr,dz
  REAL(wp_), SAVE :: zbmin,zbmax
  REAL(wp_), SAVE :: phitedge
  REAL(wp_), SAVE :: rup,zup,rlw,zlw
  REAL(wp_), SAVE :: rcen,btrcen ! rcen used to compute btrcen from fpol, btrcen used only for Jcd_ASTRA def.

  INTEGER, PARAMETER :: kspl=3,ksplp=kspl+1
  ! === 2D spline psi(R,z), normalization and derivatives ==========
  INTEGER, SAVE :: nsr, nsz
  REAL(wp_), SAVE :: psia, psiant, psinop
  REAL(wp_), DIMENSION(:), ALLOCATABLE, SAVE :: tr,tz
  REAL(wp_), DIMENSION(:), ALLOCATABLE, SAVE :: cceq,   cceq01, cceq10, &
                                                cceq20, cceq02, cceq11
  ! === 1D spline Fpol(psi) ========================================
  ! INTEGER, SAVE :: npsiest
  INTEGER, SAVE :: nsf
  REAL(wp_), DIMENSION(:), ALLOCATABLE, SAVE :: tfp, cfp
  REAL(wp_), SAVE :: fpolas
  ! === 1D spline rhot(rhop), rhop(rhot), q(psi) ===================
  ! computed on psinr,rhopnr [,rhotnr] arrays
  INTEGER, SAVE :: nq
  REAL(wp_), DIMENSION(:), ALLOCATABLE, SAVE :: rhopq,psiq,q
  REAL(wp_), DIMENSION(:,:), ALLOCATABLE, SAVE :: crhot,cq

contains

  subroutine read_eqdsk(filenm,rv,zv,psin,psia,psinr,fpol,q,rvac,rax,zax, &
                        rbnd,zbnd,rlim,zlim,ipsinorm,idesc,ifreefmt,unit)
    use const_and_precisions, only : one
    use utils, only : get_free_unit
    implicit none
! arguments
    character(len=*), intent(in) :: filenm
    real(wp_), intent(out) :: psia,rvac,rax,zax
    real(wp_), dimension(:), allocatable, intent(out) :: rv,zv,psinr,fpol,q
    real(wp_), dimension(:), allocatable, intent(out) :: rbnd,zbnd,rlim,zlim
    real(wp_), dimension(:,:), allocatable, intent(out) :: psin
    integer, optional, intent(in) :: ipsinorm,idesc,ifreefmt,unit
! local variables
    integer, parameter :: indef=0,iddef=1,iffdef=0
    integer :: in,id,iffmt,u,idum,i,j,nr,nz,nbnd,nlim
    character(len=48) :: string
    real(wp_) :: dr,dz,dps,rleft,zmid,zleft,xdum,psiedge,psiaxis
    
! set default values if optional arguments are absent
    in=indef; id=iddef; iffmt=iffdef
    if(present(ipsinorm)) in=ipsinorm
    if(present(idesc)) id=idesc
    if(present(ifreefmt)) iffmt=ifreefmt
    if (present(unit)) then
      u=unit
    else
      u=get_free_unit()
    end if

! open G EQDSK file (see http://fusion.gat.com/efit/g_eqdsk.html)
    open(file=trim(filenm),status='old',unit=u)

! get size of main arrays and allocate them
    if (id==1) then
      read (u,'(a48,3i4)') string,idum,nr,nz
    else
      read (u,*) nr,nz
    end if
    if (allocated(rv)) deallocate(rv)
    if (allocated(zv)) deallocate(zv)
    if (allocated(psin)) deallocate(psin)
    if (allocated(psinr)) deallocate(psinr)
    if (allocated(fpol)) deallocate(fpol)
    if (allocated(q)) deallocate(q)
    allocate(rv(nr),zv(nz),psin(nr,nz),psinr(nr),fpol(nr),q(nr))

! store 0D data and main arrays
    if (iffmt==1) then
      read (u,*) dr,dz,rvac,rleft,zmid
      read (u,*) rax,zax,psiaxis,psiedge,xdum
      read (u,*) xdum,xdum,xdum,xdum,xdum
      read (u,*) xdum,xdum,xdum,xdum,xdum
      read (u,*) (fpol(i),i=1,nr)
      read (u,*) (xdum,i=1,nr)
      read (u,*) (xdum,i=1,nr)
      read (u,*) (xdum,i=1,nr)
      read (u,*) ((psin(i,j),i=1,nr),j=1,nz)
      read (u,*) (q(i),i=1,nr)
    else
      read (u,'(5e16.9)') dr,dz,rvac,rleft,zmid
      read (u,'(5e16.9)') rax,zax,psiaxis,psiedge,xdum
      read (u,'(5e16.9)') xdum,xdum,xdum,xdum,xdum
      read (u,'(5e16.9)') xdum,xdum,xdum,xdum,xdum
      read (u,'(5e16.9)') (fpol(i),i=1,nr)
      read (u,'(5e16.9)') (xdum,i=1,nr)
      read (u,'(5e16.9)') (xdum,i=1,nr)
      read (u,'(5e16.9)') (xdum,i=1,nr)
      read (u,'(5e16.9)') ((psin(i,j),i=1,nr),j=1,nz)
      read (u,'(5e16.9)') (q(i),i=1,nr)
    end if

! get size of boundary and limiter arrays and allocate them
    read (u,*) nbnd,nlim
    if (allocated(rbnd)) deallocate(rbnd)
    if (allocated(zbnd)) deallocate(zbnd)
    if (allocated(rlim)) deallocate(rlim)
    if (allocated(zlim)) deallocate(zlim)

! store boundary and limiter data
    if(nbnd>0) then
      allocate(rbnd(nbnd),zbnd(nbnd))
      if (iffmt==1) then
        read(u,*) (rbnd(i),zbnd(i),i=1,nbnd)
      else
        read(u,'(5e16.9)') (rbnd(i),zbnd(i),i=1,nbnd)
      end if
    end if
    if(nlim>0) then
      allocate(rlim(nlim),zlim(nlim))
      if (iffmt==1) then
        read(u,*) (rlim(i),zlim(i),i=1,nlim)
      else
        read(u,'(5e16.9)') (rlim(i),zlim(i),i=1,nlim)
      end if
    end if

! reading of G EQDSK file completed    
    close(u)
    
! build rv,zv,psinr arrays and normalize psin
    zleft=zmid-0.5_wp_*dz
    dr=dr/(nr-1)
    dz=dz/(nz-1)
    dps=one/(nr-1)
    do i=1,nr
      psinr(i)=(i-1)*dps
      rv(i)=rleft+(i-1)*dr
    end do
    do i=1,nz
      zv(i)=zleft+(i-1)*dz
    end do
    psia=psiedge-psiaxis
    if(in==0) psin=(psin-psiaxis)/psia
  end subroutine read_eqdsk
  
  subroutine change_cocos(psia,fpol,q,cocosin,cocosout,ierr)
    use const_and_precisions, only : zero,one,pi
    implicit none
! arguments
    real(wp_), intent(inout) :: psia
    real(wp_), dimension(:), intent(inout) :: fpol,q
    integer, intent(in) :: cocosin, cocosout
!    real(wp_), intent(out) :: isign,bsign
    integer, intent(out), optional :: ierr
! local variables
    real(wp_) :: isign,bsign
    integer :: exp2pi,exp2piout
    logical :: phiccw,psiincr,qpos,phiccwout,psiincrout,qposout

    call decode_cocos(cocosin,exp2pi,phiccw,psiincr,qpos)
    call decode_cocos(cocosout,exp2piout,phiccwout,psiincrout,qposout)

!   check sign consistency
    isign=sign(one,psia)
    if (.not.psiincr) isign=-isign
    bsign=sign(one,fpol(size(fpol)))
    if (qpos.neqv.isign*bsign*q(size(q))>zero) then
!     warning: sign inconsistency found among q, Ipla and Bref
      q=-q
      if(present(ierr)) ierr=1
    else
      if(present(ierr)) ierr=0
    end if

! convert cocosin to cocosout
    if (phiccw.neqv.phiccwout) then
!     opposite direction of toroidal angle phi in cocosin and cocosout
!      bsign=-bsign
!      isign=-isign
      fpol=-fpol
    end if
!   q has opposite sign for given sign of Bphi*Ip
    if (qpos .neqv. qposout) q=-q
!   psi and Ip signs don't change accordingly
    if ((phiccw.eqv.phiccwout) .neqv. (psiincr.eqv.psiincrout)) psia=-psia
!   convert Wb to Wb/rad or viceversa
    psia=psia*(2.0_wp_*pi)**(exp2piout-exp2pi)
  end subroutine change_cocos

  subroutine decode_cocos(cocos,exp2pi,phiccw,psiincr,qpos)
    implicit none
    integer, intent(in) :: cocos
    integer, intent(out) :: exp2pi
    logical, intent(out) :: phiccw,psiincr,qpos
    integer :: cmod10,cmod4

    cmod10=mod(cocos,10)
    cmod4=mod(cmod10,4)
!  cocos>10 psi in Wb, cocos<10 psi in Wb/rad
    exp2pi=(cocos-cmod10)/10
!  cocos mod 10 = 1,3,5,7: toroidal angle phi increasing CCW
    phiccw=(mod(cmod10,2)==1)
!  cocos mod 10 = 1,2,5,6: psi increasing with positive Ip
    psiincr=(cmod4==1 .or. cmod4==2)
!  cocos mod 10 = 1,2,7,8: q positive for positive Bphi*Ip
    qpos=(cmod10<3 .or. cmod10>6)
  end subroutine decode_cocos

  subroutine eq_scal(psia,fpol,isign,bsign,factb)
    use const_and_precisions, only : one
    implicit none
    real(wp_), intent(inout) :: psia
    integer, intent(inout) :: isign,bsign
    real(wp_), dimension(:), intent(inout) :: fpol
    real(wp_), intent(in) :: factb
    
    ! isign and bsign ignored on input if equal to 0
    ! used to assign the direction of Bphi and Ipla BEFORE scaling
    ! cocos=3 assumed: CCW direction is >0
    ! Bphi and Ipla scaled by the same factor factb to keep q unchanged
    ! factb<0 reverses the directions of Bphi and Ipla
    if(isign/=0) psia=sign(psia,real(-isign,wp_))
    if(bsign/=0 .and. bsign*fpol(size(fpol))<0) fpol=-fpol
    psia=psia*factb
    fpol=fpol*factb
    isign=int(sign(one,-psia))
    bsign=int(sign(one,fpol(size(fpol))))
  end subroutine eq_scal

  subroutine set_eqspl(rv,zv,psin,psinr,fpol,sspl,ssfp)
    use const_and_precisions, only : zero,one
    use dierckx, only : regrid,coeff_parder,curfit,splev
    implicit none
! arguments
    real(wp_), dimension(:), intent(in) :: rv,zv,psinr,fpol
    real(wp_), dimension(:,:), intent(in) :: psin
    real(wp_), intent(inout) :: sspl,ssfp
! local constants
    integer, parameter :: iopt=0
    integer :: liwrk,lwrk,lw10,lw01,lw20,lw02,lw11,lwrkf
    integer :: nr,nz,nrest,nzest,npsest,nrz,npsi
    real(wp_) :: fp
    real(wp_), dimension(1) :: fpoli
    real(wp_), dimension(:), allocatable :: fvpsi,wf,wrk
    integer, dimension(:), allocatable :: iwrk
! local variables
    integer :: ier
!
!   compute array sizes and prepare working space arrays
    nr=size(rv)
    nz=size(zv)
    nrz=nr*nz
    nrest=nr+ksplp
    nzest=nz+ksplp
    lwrk=4+nrest*nz+(nrest+nzest)*(2*kspl+5)+(nr+nz)*ksplp+max(nz,nrest)
    liwrk=nz+nr+nrest+nzest+kspl
!
    npsi=size(psinr)
    npsest=npsi+ksplp
    lwrkf=npsi*ksplp+npsest*(7+3*kspl)
!
    allocate(wrk(max(lwrk,lwrkf)),iwrk(max(liwrk,npsest)))
!
!   spline fitting/interpolation of psin(i,j) and derivatives
!
!   length in m !!!
!
    rmnm=rv(1)
    rmxm=rv(nr)
    zmnm=zv(1)
    zmxm=zv(nz)
!   allocate knots and spline coefficients arrays
    if (allocated(tr)) deallocate(tr)
    if (allocated(tz)) deallocate(tz)
    allocate(tr(nrest),tz(nzest),cceq(nrz))
!   allocate work arrays
!   reshape 2D psi array to 1D (transposed) array and compute spline coeffs
    allocate(fvpsi(nrz))
    fvpsi=reshape(transpose(psin),(/nrz/))
    call regrid(iopt,nr,rv,nz,zv,fvpsi,rmnm,rmxm,zmnm,zmxm,         &
                kspl,kspl,sspl,nrest,nzest,nsr,tr,nsz,tz,cceq,fp,   &
                wrk(1:lwrk),lwrk,iwrk(1:liwrk),liwrk,ier)
!   if ier=-1 data are re-fitted using sspl=0    
    if(ier==-1) then
      sspl=0.0_wp_
      call regrid(iopt,nr,rv,nz,zv,fvpsi,rmnm,rmxm,zmnm,zmxm,       &
                  kspl,kspl,sspl,nrest,nzest,nsr,tr,nsz,tz,cceq,fp, &
                  wrk(1:lwrk),lwrk,iwrk(1:liwrk),liwrk,ier)
    end if
    deallocate(fvpsi)
!   compute spline coefficients for psi partial derivatives
    lw10 = nr*(ksplp-1) + nz*ksplp     + nrz
    lw01 = nr*ksplp     + nz*(ksplp-1) + nrz
    lw20 = nr*(ksplp-2) + nz*ksplp     + nrz
    lw02 = nr*ksplp     + nz*(ksplp-2) + nrz
    lw11 = nr*(ksplp-1) + nz*(ksplp-1) + nrz
    if (allocated(cceq10)) deallocate(cceq10)
    if (allocated(cceq01)) deallocate(cceq01)
    if (allocated(cceq20)) deallocate(cceq20)
    if (allocated(cceq02)) deallocate(cceq02)
    if (allocated(cceq11)) deallocate(cceq11)
    allocate(cceq10(lw10),cceq01(lw01),cceq20(lw20),cceq02(lw02),cceq11(lw11))
    call coeff_parder(tr,nsr,tz,nsz,cceq,kspl,kspl,1,0,cceq10,lw10,ier)
    call coeff_parder(tr,nsr,tz,nsz,cceq,kspl,kspl,0,1,cceq01,lw01,ier)
    call coeff_parder(tr,nsr,tz,nsz,cceq,kspl,kspl,2,0,cceq20,lw20,ier)
    call coeff_parder(tr,nsr,tz,nsz,cceq,kspl,kspl,0,2,cceq02,lw02,ier)
    call coeff_parder(tr,nsr,tz,nsz,cceq,kspl,kspl,1,1,cceq11,lw11,ier)
!
!   spline interpolation of fpol(i)
!
!   allocate knots and spline coefficients arrays
    if (allocated(tfp)) deallocate(tfp)
    if (allocated(cfp)) deallocate(cfp)
    allocate(tfp(npsest),cfp(npsest))
    allocate(wf(npsi))
    wf(1:npsi-1)=one
    wf(npsi)=1.0e2_wp_
    call curfit(iopt,npsi,psinr,fpol,wf,zero,one,kspl,ssfp,npsest,nsf, &
                tfp,cfp,fp,wrk(1:lwrkf),lwrkf,iwrk(1:npsest),ier)
    call splev(tfp,nsf,cfp,kspl,psinr(npsi:npsi),fpoli,1,ier)
    fpolas=fpoli(1)
    btrcen=fpolas/rcen
!   free temporary arrays
    deallocate(wrk,iwrk,wf)
  end subroutine set_eqspl

  subroutine unset_eqspl
    implicit none
    if(allocated(tr)) deallocate(tr)
    if(allocated(tz)) deallocate(tz)
    if(allocated(tfp)) deallocate(tfp)
    if(allocated(cfp)) deallocate(cfp)
    if(allocated(cceq)) deallocate(cceq)
    if(allocated(cceq01)) deallocate(cceq01)
    if(allocated(cceq10)) deallocate(cceq10)
    if(allocated(cceq02)) deallocate(cceq02)
    if(allocated(cceq20)) deallocate(cceq20)
    if(allocated(cceq11)) deallocate(cceq11)
  end subroutine unset_eqspl

  subroutine dealloc_bnd
    implicit none
    if(allocated(rbbbs))   deallocate(rbbbs)
    if(allocated(zbbbs))   deallocate(zbbbs)
  end subroutine dealloc_bnd

  subroutine equinum_psi(rpsim,zpsim,psinv,dpsidr,dpsidz, &
        ddpsidrr,ddpsidzz,ddpsidrz)
    use dierckx, only : fpbisp
    implicit none
! local constants
    integer, parameter :: lwrk=8,liwrk=2
! arguments
    real(wp_), intent(in) :: rpsim,zpsim
    real(wp_), intent(out), optional :: psinv,dpsidr,dpsidz, &
        ddpsidrr,ddpsidzz,ddpsidrz

! local variables
    integer, dimension(liwrk) :: iwrk
    real(wp_), dimension(1) :: rrs,zzs,ffspl
    real(wp_), dimension(lwrk) :: wrk
!
!     here lengths are measured in meters
!
    if (rpsim.le.rmxm .and. rpsim.ge.rmnm .and. &
        zpsim.le.zmxm .and. zpsim.ge.zmnm) then

      if (present(psinv)) then
        rrs(1)=rpsim
        zzs(1)=zpsim
        call fpbisp(tr,nsr,tz,nsz,cceq,3,3,rrs,1,zzs,1,ffspl, &
                   wrk(1),wrk(5),iwrk(1),iwrk(2))
        psinv=(ffspl(1)-psinop)/psiant
      end if
      if (present(dpsidr)) then
        call sub_derpsi(rpsim,zpsim,1,0,dpsidr,cceq10)
      end if
      if (present(dpsidz)) then
        call sub_derpsi(rpsim,zpsim,0,1,dpsidz,cceq01)
      end if
      if (present(ddpsidrr)) then
        call sub_derpsi(rpsim,zpsim,2,0,ddpsidrr,cceq20)
      end if
      if (present(ddpsidzz)) then
        call sub_derpsi(rpsim,zpsim,0,2,ddpsidzz,cceq02)
      end if
      if (present(ddpsidrz)) then
        call sub_derpsi(rpsim,zpsim,1,1,ddpsidrz,cceq11)
      end if
    else
      if(present(psinv)) psinv=-1.0_wp_
      if(present(dpsidr)) dpsidr=0.0_wp_
      if(present(dpsidz)) dpsidz=0.0_wp_
      if(present(ddpsidrr)) ddpsidrr=0.0_wp_
      if(present(ddpsidzz)) ddpsidzz=0.0_wp_
      if(present(ddpsidrz)) ddpsidrz=0.0_wp_
    end if
  end subroutine equinum_psi

  subroutine sub_derpsi(rpsim,zpsim,nur,nuz,derpsi,cc)
    use dierckx, only : fpbisp
    implicit none
! arguments
    real(wp_), intent(in) :: rpsim,zpsim
    integer, intent(in) :: nur,nuz
    real(wp_), intent(out) :: derpsi
    real(wp_), dimension(:) :: cc
! local variables
    integer, dimension(1) :: iwrkr,iwrkz
    real(wp_), dimension(1) :: rrs,zzs,ffspl
    real(wp_), dimension(1,ksplp) :: wrkr
    real(wp_), dimension(1,ksplp) :: wrkz

    rrs(1)=rpsim
    zzs(1)=zpsim
    call fpbisp(tr(nur+1),nsr-2*nur,tz(nuz+1),nsz-2*nuz,cc,kspl-nur,kspl-nuz, &
         rrs,1,zzs,1,ffspl,wrkr,wrkz,iwrkr,iwrkz)
    derpsi=ffspl(1)*psia
  end subroutine sub_derpsi

  subroutine equinum_fpol(psinv,fpolv,dfpv)
    use dierckx, only : splev,splder
    implicit none
! arguments
    real(wp_), intent(in) :: psinv
    real(wp_), intent(out) :: fpolv
    real(wp_), intent(out), optional :: dfpv
! local variables
    integer :: ier
    real(wp_), dimension(1) :: rrs,ffspl
    real(wp_), dimension(nsf) :: wrkfd
!      
    if(psinv.le.1.0_wp_.and.psinv.gt.0.0_wp_) then
      rrs(1)=psinv
      call splev(tfp,nsf,cfp,3,rrs,ffspl,1,ier)
      fpolv=ffspl(1)
      if(present(dfpv)) then       
        call splder(tfp,nsf,cfp,3,1,rrs,ffspl,1,wrkfd,ier)
        dfpv=ffspl(1)/psia
      end if
    else
      fpolv=fpolas
      if (present(dfpv)) dfpv=0._wp_
    end if
  end subroutine equinum_fpol

!  subroutine btor(rpsim,zpsim,bphi)
!    implicit none
!! arguments
!    real(wp_), intent(in) :: rpsim,zpsim
!    real(wp_), intent(out) :: bphi
!! local variables
!    real(wp_) :: psinv,fpolv
!
!    call equinum_psi(rpsim,zpsim,psinv)
!    call equinum_fpol(psinv,fpolv)
!    bphi=fpolv/rpsim
!  end subroutine btor

!  subroutine bpol(rpsim,zpsim,brr,bzz)
!    implicit none
!! arguments
!    real(wp_), intent(in) :: rpsim,zpsim
!    real(wp_), intent(out) :: brr,bzz
!! local variables
!    real(wp_) :: dpsidr,dpsidz
!
!    call equinum_psi(rpsim,zpsim,dpsidr=dpsidr,dpsidz=dpsidz)
!    brr=-dpsidz/rpsim
!    bzz= dpsidr/rpsim
!  end subroutine bpol

  subroutine bfield(rpsim,zpsim,bphi,br,bz)
    implicit none
! arguments
    real(wp_), intent(in) :: rpsim,zpsim
    real(wp_), intent(out), optional :: bphi,br,bz
! local variables
    real(wp_) :: psin,fpol

    call equinum_psi(rpsim,zpsim,psinv=bphi,dpsidr=bz,dpsidz=br)
    if (present(bphi)) then
      psin=bphi
      call equinum_fpol(psin,fpol)
      bphi=fpol/rpsim
    end if
    if (present(br)) br=-br/rpsim
    if (present(bz)) bz= bz/rpsim
  end subroutine bfield

end module interp_eqprof