module trace_fieldlines
use precision, only: lk
implicit none

public trace, trace_sep, trace_sepsol, trace_perp, fmroot
private

save
contains
!==============================================================================
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!       Tracing the magnetic field line (first part is core region)
!                  (second part is separatrix)
!   iflag=1: core, iflag=2: positive Z sol, iflag=3: negative Z sol
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
subroutine trace(iflag,np,ngtmp,r00,z00,psi00,deltas,zlow,zhigh,jpout,rout,zout,tout)
    use global_parameters,only:gtcout,mype,pi
    use spline_function, only: spline2d
    use cylindrical_function, only: geotheta
    use equilibrium, only: psirz_Eq,&
                           lsr, lsz, spdR, spdZ, zmax
    implicit none

    integer,intent(in) :: iflag,np,ngtmp
    integer,intent(out) :: jpout(np)
    real(lk),intent(in) :: r00(np),z00(np),psi00(np),deltas, zlow, zhigh
    real(lk),intent(out) :: rout(np,ngtmp),zout(np,ngtmp),tout(np,ngtmp)


    integer :: i,j,ij
    real(lk) :: pdum,delta_pdum,sperr,brtmp,bztmp,btmp,spsi

    do i=1,np
        rout(i,1)=r00(i)
        zout(i,1)=z00(i)
        tout(i,1)=geotheta(rout(i,1),zout(i,1))
        do j=1,ngtmp-1
            brtmp=spline2d(rout(i,j),zout(i,j)+zmax,2,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
            brtmp=-brtmp/rout(i,j)
            bztmp=spline2d(rout(i,j),zout(i,j)+zmax,1,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
            bztmp=bztmp/rout(i,j)
            btmp=sqrt(brtmp*brtmp+bztmp*bztmp)
            brtmp=brtmp/btmp
            bztmp=bztmp/btmp
            if(iflag==3)then
                brtmp=-brtmp
                bztmp=-bztmp
            endif
            rout(i,j+1)=rout(i,j)+brtmp*deltas
            zout(i,j+1)=zout(i,j)+bztmp*deltas
  
            pdum=spline2d(rout(i,j+1),zout(i,j+1)+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
            delta_pdum=pdum-psi00(i)
            sperr=1.0e-8_lk
 !  if(mype==0)write(gtcout,*)"pdum=",pdum,"rout(i,j+1)=",rout(i,j+1),"zout(i,j+1)=",zout(i,j+1)
 !  if(mype==0)write(gtcout,*)"rout(i,j)=",rout(i,j),"zout(i,j)=",zout(i,j)
 !  if(mype==0)write(gtcout,*)"brtmp=",brtmp,"bztmp=",bztmp
            tout(i,j+1)=geotheta(rout(i,j+1),zout(i,j+1))
  ! Correction of the new position 
            if(abs(delta_pdum)>sperr)then
                call fmroot(0,rout(i,j),zout(i,j),psi00(i),deltas,brtmp,bztmp,delta_pdum,sperr,rout(i,j+1),zout(i,j+1),pdum,spsi)
            endif
  
            if(iflag==1)then
                if(zout(i,j) .lt. zout(i,1) .and. zout(i,j+1) .ge. zout(i,1) .or. j==ngtmp-1)then
                    jpout(i)=j
                    exit
                endif
            elseif(iflag==2)then
                if(zout(i,j+1) .gt. zhigh .or. j==ngtmp-1)then
                    jpout(i)=j
                    exit
                endif
            elseif(iflag==3)then
                if(zout(i,j+1) .lt. zlow .or. j==ngtmp-1)then
                    jpout(i)=j
                    exit
                endif
            endif

        enddo
    enddo


end subroutine trace
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!                                                                       !
!                            Tracing the separatrix                     !
!                                                                       !
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
subroutine trace_sep(ngtmp,rsep0,zsep0,deltas,jpout,rout,zout,tout,z1,z2,a)
    use global_parameters,only:gtcout,mype,pi
    use spline_function, only: spline2d
    use cylindrical_function, only: geotheta
    use equilibrium, only: psi_sep, lsr, lsz, spdR, spdZ, zmax,&
                           psirz_Eq
    implicit none

    integer,intent(in) :: ngtmp
    integer,intent(out) :: jpout
    real(lk),intent(in) :: rsep0,zsep0,deltas
    real(lk),intent(out) :: rout(ngtmp),zout(ngtmp),tout(ngtmp)
    real(lk),intent(out),optional ::z1,z2,a

    integer :: i,j,ij,mtest,s1,s2,numal
    real(lk) pdum,delta_pdum,sperr,brtmp,bztmp,btmp,spsi,&
             r1,r2
    real(lk),dimension(:),allocatable :: Rfatmp1,Zfatmp1,Rfatmp2,Zfatmp2

! Tracing the separatrix
    rout(1)=rsep0
    zout(1)=zsep0

    do j=1,ngtmp-1
        brtmp=spline2d(rout(j),zout(j)+zmax,2,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        brtmp=-brtmp/rout(j)
        bztmp=spline2d(rout(j),zout(j)+zmax,1,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        bztmp=bztmp/rout(j)
        btmp=sqrt(brtmp*brtmp+bztmp*bztmp)
        brtmp=brtmp/btmp
        bztmp=bztmp/btmp
 
        rout(j+1)=rout(j)+brtmp*deltas
        zout(j+1)=zout(j)+bztmp*deltas
 
        pdum=spline2d(rout(j+1),zout(j+1)+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        delta_pdum=pdum-psi_sep
 
        sperr=1.0e-8_lk
 
 ! Correction of the new position 
        if(abs(delta_pdum)>sperr)then
          call fmroot(1,rout(j),zout(j),psi_sep,deltas,brtmp,bztmp,delta_pdum,sperr,rout(j+1),zout(j+1),pdum,spsi)
        endif
        if(rout(j+1) .lt. 0.001_lk)then
           jpout=j
           exit
        endif

    enddo


    r1=rout(jpout)
    z1=zout(jpout)
    s1=jpout+1

    rout(s1)=rsep0
    zout(s1)=zsep0

    do j=s1,ngtmp-1
        brtmp=spline2d(rout(j),zout(j)+zmax,2,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        brtmp=-brtmp/rout(j)
        bztmp=spline2d(rout(j),zout(j)+zmax,1,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        bztmp=bztmp/rout(j)
        btmp=sqrt(brtmp*brtmp+bztmp*bztmp)
        brtmp=-brtmp/btmp
        bztmp=-bztmp/btmp
 
        rout(j+1)=rout(j)+brtmp*deltas
        zout(j+1)=zout(j)+bztmp*deltas
 
        pdum=spline2d(rout(j+1),zout(j+1)+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        delta_pdum=pdum-psi_sep
 
        sperr=1.0e-8_lk
 
 ! Correction of the new position 
        if(abs(delta_pdum)>sperr)then
          call fmroot(1,rout(j),zout(j),psi_sep,deltas,brtmp,bztmp,delta_pdum,sperr,rout(j+1),zout(j+1),pdum,spsi)
        endif
 
        if(rout(j+1) .lt. 0.001_lk)then
           jpout=j
           exit
        endif
    enddo

    r2=rout(jpout)
    z2=zout(jpout)
    s2=jpout+1

    a=(z1-z2)
    numal=int(abs(a)/deltas)
    rout(s2)=0.0_lk
    zout(s2)=z1
  ! the number of this block is  numal
    do i=s2,s2+numal-2
        rout(i+1)=0.0_lk
        zout(i+1)=zout(i)-deltas
    enddo
    jpout=jpout+numal


! Redefine the array index to make it clockwise
    allocate(Rfatmp1(s2-s1),Zfatmp1(s2-s1),Rfatmp2(numal),Zfatmp2(numal))
    do i=s1,s2-1
        Rfatmp1(i-s1+1)=rout(i)
        Zfatmp1(i-s1+1)=zout(i)
    enddo
    do i=s1,s2-1
        rout(i)=Rfatmp1(s2-i)
        zout(i)=Zfatmp1(s2-i)
    enddo
    do i=s1,s2-1
        Rfatmp1(i-s1+1)=rout(i)
        Zfatmp1(i-s1+1)=zout(i)
    enddo

    do i=s2,s2+numal-1
        Rfatmp2(i-s2+1)=rout(i)
        Zfatmp2(i-s2+1)=zout(i)
    enddo

    do i=s1,s1+numal-1
        rout(i)=Rfatmp2(i-s1+1)
        zout(i)=Zfatmp2(i-s1+1)
    enddo
    do i=s1+numal,jpout
        rout(i)=Rfatmp1(i-s1-numal+1)
        zout(i)=Zfatmp1(i-s1-numal+1)
    enddo

! Calculating the geometrix \theta value on each small grids on the separatrix
    tout(1)=geotheta(rout(1),zout(1))
    do j=1,jpout-1
        tout(j+1)=geotheta(rout(j+1),zout(j+1))
    enddo

end subroutine trace_sep
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!                                                                       !
!                        Tracing the sol separatrix                     !
!        iflag=1 is for right sol and iflag=2 is for left sol           ! 
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
subroutine trace_sepsol(iflag,ngtmp,rsep0,zsep0,deltas, zlow, zhigh, jpout,rout,zout)
    use global_parameters,only:gtcout,mype,pi
    use spline_function, only: spline2d
    use cylindrical_function, only: geotheta
    use equilibrium, only: psi_sep, lsr, lsz, spdR, spdZ, zmax, &
                           psirz_Eq
    implicit none

    integer,intent(in) :: iflag,ngtmp
    integer,intent(out) :: jpout
    real(lk),intent(in) :: rsep0,zsep0,deltas, zlow, zhigh
    real(lk),intent(out) :: rout(ngtmp),zout(ngtmp)

    integer :: i,j,ij,mtest,s1,numal
    real(lk) pdum,delta_pdum,sperr,brtmp,bztmp,btmp,spsi,&
             r1,z1,a

! Tracing the separatrix
    rout(1)=rsep0
    zout(1)=zsep0

    do j=1,ngtmp-1
        brtmp=spline2d(rout(j),zout(j)+zmax,2,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        brtmp=-brtmp/rout(j)
        bztmp=spline2d(rout(j),zout(j)+zmax,1,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        bztmp=bztmp/rout(j)
        btmp=sqrt(brtmp*brtmp+bztmp*bztmp)
        brtmp=brtmp/btmp
        bztmp=bztmp/btmp
        if(iflag==2)then
            brtmp=-brtmp
            bztmp=-bztmp
        endif
        rout(j+1)=rout(j)+brtmp*deltas
        zout(j+1)=zout(j)+bztmp*deltas
 
        pdum=spline2d(rout(j+1),zout(j+1)+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        delta_pdum=pdum-psi_sep
 
        sperr=1.0e-8_lk
 
 ! Correction of the new position 
        if(abs(delta_pdum)>sperr)then
            call fmroot(1,rout(j),zout(j),psi_sep,deltas,brtmp,bztmp,delta_pdum,sperr,rout(j+1),zout(j+1),pdum,spsi)
        endif
        if(rout(j+1) .lt. 0.001_lk)then
            jpout=j
            exit
        endif

    enddo


    r1=rout(jpout)
    z1=zout(jpout)
    s1=jpout+1
    if(iflag==1)then
        a=abs(zhigh)-abs(z1)
    elseif(iflag==2)then
        a=abs(zlow)-abs(z1)
    endif
    if(a<0.0_lk .and. mype==0)write(gtcout,*)"trace sol sep not available"
    numal=int(abs(a)/deltas)
    rout(s1)=0.0_lk
    zout(s1)=z1
  ! the number of this block is  numal
    do i=s1,s1+numal-2
        rout(i+1)=0.0_lk
        if(iflag==1)then
            zout(i+1)=zout(i)+deltas
        elseif(iflag==2)then
            zout(i+1)=zout(i)-deltas
        endif
    enddo
    jpout=jpout+numal



end subroutine trace_sepsol
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
!  iflag==1 increasing \psi direction
!  iflag==2 decreasing \psi direction
!
subroutine trace_perp(iflag,ng,r00,z00,psi00,deltas,jpout,rout,zout)
    use global_parameters,only:gtcout,mype,pi
    use spline_function, only: spline2d
    use cylindrical_function, only: geotheta
    use equilibrium, only: psirz_Eq,&
                           lsr, lsz, spdR, spdZ, zmax
    implicit none

    integer,intent(in) :: iflag,ng
    real(lk),intent(in) :: r00,z00,psi00,deltas
    integer,intent(out) :: jpout
    real(lk),intent(out) :: rout(ng),zout(ng)


    integer :: j
    real(lk) :: pdum,delta_pdum,sperr,brtmp,bztmp,btmp,bperp_r,bperp_z



     rout(1)=r00
     zout(1)=z00
     do j=1,ng-1
! not the real brtmp
         brtmp=spline2d(rout(j),zout(j)+zmax,2,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
         brtmp=-brtmp/rout(j)
         bztmp=spline2d(rout(j),zout(j)+zmax,1,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
         bztmp=bztmp/rout(j)
         btmp=sqrt(brtmp*brtmp+bztmp*bztmp)
         brtmp=brtmp/btmp
         bztmp=bztmp/btmp
         if(iflag==1)then    ! positive \psi direction
             bperp_r=bztmp
             bperp_z=-brtmp
         elseif(iflag==2)then  ! negetive \psi direction
             bperp_r=-bztmp
             bperp_z=brtmp
         else
             if(mype==0)write(gtcout,*)"wrong iflag in trace_perp"
             stop
         endif

         rout(j+1)=rout(j)+bperp_r*deltas
         zout(j+1)=zout(j)+bperp_z*deltas

         pdum=spline2d(rout(j+1),zout(j+1)+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
         if(iflag==1)then
             delta_pdum=pdum-psi00
         elseif(iflag==2)then
             delta_pdum=psi00-pdum
         else
             if(mype==0)write(gtcout,*)"wrong iflag in trace_perp"
             stop
         endif

         if(delta_pdum .gt. 0.0_lk)then
             jpout=j
             exit
         elseif(j==ng-1)then
             if(mype==0)write(gtcout,*)"grids for perpedicular trace are not enough"
             stop
         endif
     enddo


end subroutine trace_perp
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!###### MODULE SPLINE_FUNCTION END ######################################
! only valid when z=psirz
! only be used for creating the field aligned mesh from efit data
subroutine fmroot(iflag,x0,y0,zi,si,br,bz,dp,errdum,x1,y1,z1,sp)
    use global_parameters,only:gtcout,pi
    use spline_function,only: spline2d
    use equilibrium, only: psirz_Eq,&
                           lsr, lsz, spdR, spdZ, zmax
    implicit none
  
    integer,intent(in) :: iflag
    real(lk),intent(in) :: x0,y0,zi,si,br,bz,dp,errdum
    real(lk),intent(out) :: x1,y1,z1,sp
  
    integer :: s
    real(lk) :: thdum,deltath,thdum0,brdum0,bzdum0,pdum0,dpdum0,thlow,thhigh,tdum,pdum1,dpdum1,pdum,dpdum,brdum1,bzdum1,brdum,bzdum
  
  
    thdum=acos(br)
    if(br*bz<0.0_lk .and. br>0.0_lk)then
        thdum=2.0_lk*pi-thdum
    elseif(br*bz>0.0_lk .and. br<0.0_lk)then
       thdum=2.0_lk*pi-thdum
    endif
    deltath=0.2*pi
    thdum0=thdum+deltath
    brdum0=cos(thdum0)
    bzdum0=sin(thdum0)
    x1=x0+brdum0*si
    y1=y0+bzdum0*si
    pdum0=spline2d(x1,y1+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
    dpdum0=pdum0-zi
    sp=dp*dpdum0
    if(sp>0)then
        thdum0=thdum-deltath
        brdum0=cos(thdum0)
        bzdum0=sin(thdum0)
        x1=x0+brdum0*si
        y1=y0+bzdum0*si
        pdum0=spline2d(x1,y1+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        dpdum0=pdum0-zi
        sp=dp*dpdum0
        if(sp>0)then
            write(gtcout,*)"sp =",sp
            write(gtcout,*)"position(R,Z) =",x0,y0
            write(gtcout,*)"The root is not found for the field aligned mesh"
            if(iflag==0)then
                stop
            elseif(iflag==1)then
                return
            endif
        endif
    endif
  
    s=1 ! interation times
    thlow=thdum
    thhigh=thdum0
    dpdum=dp
    do while (abs(dpdum)>errdum .and. s<50)
        s=s+1
        tdum=0.5_lk*(thlow+thhigh)
        brdum1=cos(tdum)
        bzdum1=sin(tdum)
        x1=x0+brdum1*si
        y1=y0+bzdum1*si
        pdum1=spline2d(x1,y1+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
        dpdum1=pdum1-zi
        if(dpdum1*dpdum0<0)then
            thlow=tdum
            brdum=cos(thlow)
            bzdum=sin(thlow)
            x1=x0+brdum*si
            y1=y0+bzdum*si
            pdum=spline2d(x1,y1+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
            dpdum=pdum-zi
        else
            thhigh=tdum
            brdum0=cos(thhigh)
            bzdum0=sin(thhigh)
            x1=x0+brdum0*si
            y1=y0+bzdum0*si
            pdum0=spline2d(x1,y1+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
            dpdum0=pdum0-zi
        endif
!       if(abs(dpdum)<errdum)write(gtcout,*)"s =",s
     enddo
     z1=spline2d(x1,y1+zmax,0,lsr,lsz,spdR,spdZ,psirz_Eq,0,0)
end subroutine fmroot

end module trace_fieldlines
