! 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.
! ==============================================================================

subroutine tag
  use utility, only: check_allocation, mstat
  use global_parameters
  use particle_array
  use particle_tracking
  implicit none

! Each MPI process has its own "ptracked" array to accumulate the tracked
! particles that may or may not reside in its subdomain.
! The vector "ntrackp" contains the number of tracked particles currently residing on a MPI process
  ntrackp=0
!!array to store ion tracked particles
  allocate(ptrackedi(nparami,max(mimax,1)),ptrackedi0(nparami,max(mimax,1)), STAT=mstat)
  call check_allocation(mstat, "ptrackedi in tracking.F90")
  ptrackedi=0.0
  ptrackedi0=0.0
  if(nhybrid>0)then
     !!array to store tracked electrons
     allocate(ptrackede(nparame,max(memax,1)), STAT=mstat) 
     call check_allocation(mstat, "ptrackede in tracking.F90")
     ptrackede=0.0
  endif
! If irun==0 and if particle tracking is "on", tag each particle with a unique number
  call tag_particles

end subroutine tag
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

subroutine partout

  use global_parameters
  implicit none
  call locate_tracked_particles
  if(ipoincare==0)  call write_tracked_particles
  if(ipoincare==1)  call write_poincare_plot
end subroutine partout

!========================================
!! originally written by s. either, modified by Y. Xiao April 2008
subroutine tag_particles

!========================================

  use global_parameters
  use particle_array
  use field_array
  use particle_tracking
  implicit none
  integer :: m,np
  real(lk) :: ainside,aoutside,thickratio,a0,a1
! We keep track of the particles by tagging them with a unique number.
! We add an extra element to the particle array, which holds the
! particle tag, i.e. just a number.
! The input parameter "nptrack" is the total number of particles that
! we track.

! We tag each particle with a unique number that will be carried along
! with it as it moves from one processor/domain to another. To facilitate
! the search for the tracked particles, we give the non-tracked particles
! a negative value.
! CAVEAT: We are storing the tag in a floating point array, which means that
!         for a large number of particles, the truncation error associated
!         with the precision may make some particles have the same tag.
!         This is particularly true in single precision where the maximum
!         number of particles that can be distinguished with tags differing
!         by unity is 2**24-1 = 16,777,215 (16.7 million). The most important
!         is to distinguish the tracked particles between them. The other
!         particles just need to have a negative tag. In order to minimize
!         the effect of the precision truncation, we retag the tracked
!         particles with positive numbers starting at 1.

  do m=1,mi
     zion(nparami-1,m)=-real(m,lk) !!!-real(m+mype*mi)
     zion(nparami,m)=real(mype+1,lk)
     zion0(nparami-1,m)=-real(m,lk)
     zion0(nparami,m)=real(mype+1,lk)
  enddo    
  if (nhybrid>0) then
     do m=1,me
        zelectron(nparame-1,m)=-real(m,lk) !!!-real(m+mype*mi)
        zelectron(nparame,m)=real(mype+1,lk)
        zelectron0(nparame-1,m)=-real(m,lk)
        zelectron0(nparame,m)=real(mype+1,lk)
     enddo

  endif

  a0=rg0
  a1=rg0+deltar*real(mpsi)
  thickratio=0.000002
  ainside=0.5*(1.0-thickratio)*a1+0.5*(1.0+thickratio)*a0
  ainside=ainside*ainside*0.5

  aoutside=0.5*(1.0+thickratio)*a1+0.5*(1.0-thickratio)*a0
  aoutside=aoutside*aoutside*0.5
!tag all particles when doing poincare
  if(ipoincare==1)then
    a0=psimesh(0)
    a1=psimesh(mpsi)
    ainside=0.0
    aoutside=1.0
    ainside=a0+(a1-a0)*ainside
    aoutside=a0+(a1-a0)*aoutside
  endif

  if (nhybrid>0) then
      np=0
      do m=1,me
        
         if(zelectron(1,m)>ainside .and. zelectron(1,m)<aoutside) then
            np=np+1
            zelectron(nparame-1,m)=real(np,lk)
            zelectron0(nparame-1,m)=real(np,lk)
         endif
      enddo
  endif

  if(mype==0)write(*,*)'np=',np,'ainside=',ainside,'aoutside=',aoutside
!!  if(mype==0)write(*,*)'zelectron(1,1:20)=',zelectron(1,1:20),'zelectron(nparame,1)=',zelectron(nparame,1)
!! we only condiser those particles originally satify a certain condition
! namely within an annulous ain<r<aout
! We divide the number of tracked particles equally between all processors
! as much as possible. If nptrack is not a multiple of numberpe, we add
! 1 particle to some of the the processors until we get to nptrack. We also
! start at mi/2 so that the lower numbers fall around r=0.5 (see subroutine 
! load for the initial particle distribution in r).
!  acentre=0.5*(a0+a1)
!  awidth=0.1*(a1-a0)*0.5
!  ainside= acentre-awidth
!  aoutside=acentre+awidth

 np=0
! this part is for tgaaing a single particle 
 ! if(mype==0)then
	! np=np+1
	! zion(nparami-1,1)=real(np,wp)
        ! zion0(nparami-1,1)=real(np,wp)
 ! endif

! use this part for selecting a group of particles

  do m=1,mi
!      r=sqrt(2.0*zion(1,m))
     if(zion(1,m)>ainside .and. zion(1,m)<aoutside) then
       np=np+1
       zion(nparami-1,m)=real(np,lk)
       zion0(nparami-1,m)=real(np,lk)
     endif
  enddo

!! if (nhybrid>0) then
!!      np=0
!!      do m=1,me
!!         r=sqrt(2.0*zelectron(1,m))
!!         if(r>ainside .and. r<aoutside) then
!!            np=np+1
!!            zelectron(nparami-1,m)=real(np,wp)
!!            zelectron0(nparami-1,m)=real(np,wp)
!!         endif
!!      enddo
!! endif

!  if(mype==0)write(78,*)'np=',np,'  npp=',npp,'  zion(nparami-1,mi/2)=',zion(nparami-1,mi/2),&
!                        '  zion(nparami-1,1)=',zion(nparami-1,1)

!  if(mype==0)then
!  ! On the master process (mype=0), we pick "nptrack" particles that will
!  ! be followed at every time step. To facilitate the search of those
!  ! particles among all the others in the zion arrays of each processor,
!  ! we give them a positive number from 1 to nptrack. The particles are
!  ! picked around r=0.5, which is mi/2.
!    do m=(mi-nptrack)/2,(mi+nptrack)/2-1
!       zion(nparami-1,m)=-zion(nparami-1,m)
!       zion0(nparami-1,m)=-zion(nparami-1,m)
!    enddo
!  endif

end subroutine tag_particles

!========================================

subroutine locate_tracked_particles

!========================================

  use global_parameters
  use particle_array
  use particle_tracking
  implicit none
  integer :: m,npp

! Check if tracked particles are located on this processor. Particles that
! are tracked at every time step have a positive zion(nparami-1,m) value. All the
! others have a negative value.
  ntrackp=0
  ptrackedi=0.0
  ptrackedi0=0.0
  npp=0

  do m=1,mi
     if(zion(nparami-1,m)>0.0)then
       npp=npp+1
       ptrackedi(1:nparami,npp)=zion(1:nparami,m)
       ptrackedi0(1:nparami,npp)=zion0(1:nparami,m)
     endif
  enddo
  ntrackp(1)=npp
  npp=0

  if (nhybrid>0) then
     ptrackede=0.0
     do m=1,me
        if(zelectron(nparame-1,m)>0.0)then
          npp=npp+1
          ptrackede(1:nparame,npp)=zelectron(1:nparame,m)
        endif
     enddo
  endif
  ntrackp(2)=npp
!!  if(mype==0)write(*,*)'me=',me,'np=',npp,'zelectron(1,1:6)=',zelectron(nparame-1,1:6),&
!!                        'zelectron(nparame,1)= ',zelectron(nparame,1)

end subroutine locate_tracked_particles

!========================================

subroutine write_tracked_particles

!========================================

  use global_parameters
  use particle_tracking
  implicit none
 
  integer :: j
  character(len=30) :: cdum

  
     write(cdum,'("trackp_dir/TRACKP.",i5.5)')mype
     cdum=trim(cdum)//char(0)
     if(irun==0 .and. istep==ndiag)then
        open(57,file=cdum,status='replace',position='append')
     else
        open(57,file=cdum,status='old',position='append')
     endif
     write(57,*)istep+irun*mstep
     write(57,*)ntrackp(1:nspecies)

     do j=1,ntrackp(1)
        write(57,*)ptrackedi(1:nparami,j)
     enddo
     if(nhybrid>0)then
        do j=1,ntrackp(2)
           write(57,*)ptrackede(1:nparame,j)
        enddo
     endif

     close(57)
end subroutine write_tracked_particles
!========================================

!>  write poincare data into file 'trackp_dir/poincare.00000'
!>  we have 7 colums of data in the file
!>  corresponding to 7 parameters of particles when passing the zeta=0 surface
!>  col.1:psi
!>  col.2:theta
!>  col.3:R
!>  col.4:Z
!>  col.5:safety_factor
!>  col.6:particle_ID
!>  col.7:time_step

subroutine write_poincare_plot
  use precision
  use utility, only: check_allocation, mstat
  use global_parameters, only: irun,istep,mstep,ndiag,pi,pi2
  use particle_tracking, only: ptrackedi,ptrackedi0,ntrackp
  use spline_function, only: spx,spz,spq
  implicit none
 
  integer :: i,j,isp
  integer,parameter :: ndata=7
  character(len=30) :: poincare
  character(len=12) data_name(ndata)
! coordinate transformation,using spx,spz
! spq is the function of  calculating particle's safety factor
! ptrackedii is the output argument of ion when doing poincare
! psii is a variable related with psi,and thetai is a variable related with theta
  real(lk) ptrackedii(8),psii,thetai,q
! large and small are two variables to save the particles' coordinates when doing linear interpolation
  real(lk),dimension(:,:),allocatable :: large,small
! define the name of poincare data 
  data data_name/'psi','theta','R','Z','q','particle_ID','time_step'/

  
     write(poincare,'("trackp_dir/poincare.",i5.5)')0
     poincare=trim(poincare)//char(0)
! write the name of poincare data into data file 'poincare.00000' at the first step
! if we use multi-MPI, there would only be one file at last 
     if(irun==0 .and. istep==ndiag)then
        open(58,file=poincare,status='replace',position='append')
        write(58,100)data_name
        close(58)
     endif
! ntrackp(1) is the number of particle(ion) we have tag in the corresponding MPI  
     allocate(large(3,ntrackp(1)), small(3,ntrackp(1)), STAT=mstat)
     call check_allocation(mstat, "large, small in tracking.F90")
! poincare circulation start
    do j=1,ntrackp(1)
! ptrackedi(i,j),i=1~3,represent for psi,theta,zeta of particle at present
! ptrackedi0(i,j),i=1~3,represent for psi,theta,zeta of particle at the previous time
! We use the abslute change of zeta in a time step larger than pi 
! as the criterion of passing the zeta=0 surface.
! What's more, there is an requirement for the motion of an single particle,
! that the particle's zeta change in a time step should be smaller than pi
! if not, we can't distinguish the direction of motion,and would make some 
! mistakes when we estimate whether the particle pass the zeta=0 surface.
! Since we can't know the direction of particle motion ahead of time,
! the linear interpolation process could be troublesome. So,we define 
! large and small to save the value of zeta,
! the larger one(e.g.zeta=6.28...) and the smaller one(e.g.zeta=0.0...)  
      if(abs(ptrackedi(3,j)-ptrackedi0(3,j))>pi)then  
         if(ptrackedi(3,j)>ptrackedi0(3,j))then 
           do i=1,3
             large(i,j)=ptrackedi(i,j)
             small(i,j)=ptrackedi0(i,j)
           enddo  
         else
           do i=1,3
             large(i,j)=ptrackedi0(i,j)
             small(i,j)=ptrackedi(i,j)
           enddo
         endif


         ptrackedii=0.0
! write the poincare data under the data name
         open(58,file=poincare,status='old',position='append')
! linear interpolation process in (psi,theta,zeta) configuration space 
! (psii,thetai,0)is the point between(large(1),large(2),large(3))and (small(1),small(2),small(3))
! pay attention the effect of 2pi 
! that is: large(3,j)-->large(3,j)-pi2
         psii=large(1,j)-(large(3,j)-pi2)*(small(1,j)-large(1,j))/(small(3,j)+pi2-large(3,j))

         if(abs(ptrackedi(2,j)-ptrackedi0(2,j))<pi)then
           thetai=large(2,j)-(large(3,j)-pi2)*(small(2,j)-large(2,j))/(small(3,j)+pi2-large(3,j))
         endif
! if there is a situation that when tha particle passing zeta=0 surface,
! it happens to pass the theta=0 line,the interpolation equation should consider the effect of 2pi
! that is: small(2,j)-->small(2,j)+pi2
         if(abs(ptrackedi(2,j)-ptrackedi0(2,j))>pi)then
            if(small(2,j)>large(2,j))then
               thetai=small(2,j)-small(3,j)*(small(2,j)-large(2,j)-pi2)/(small(3,j)+pi2-large(3,j))
            else if(small(2,j)<large(2,j))then
               thetai=small(2,j)+pi2-small(3,j)*(small(2,j)+pi2-large(2,j))/(small(3,j)+pi2-large(3,j))  
            endif
         endif
! maintain the value of theta in between 0~2pi
       thetai=modulo(thetai,pi2)
! calculate q of particle
       q=spq(psii)
! ptrackedii(i),i=1~7,represent the value of 
! 1:psi coordinate
! 2:theta coordinate
! 3:R coordinate
! 4:Z coordinate
! 5:safety factor
! 6:particle id
! 7:time of passing zeta=0 surface

       ptrackedii(1)=psii
       ptrackedii(2)=thetai
       ptrackedii(3)=spx(psii,thetai)
       ptrackedii(4)=spz(psii,thetai)
       ptrackedii(5)=q
       ptrackedii(6)=ptrackedi(8,j)
       ptrackedii(7)=istep+irun*mstep
       write(58,101)ptrackedii(1:7)
       close(58)
     end if
    end do 

100 format(7a15)
101 format(7e15.7)

end subroutine write_poincare_plot
