!------------------------------------------------------------------------------
! EPFL/Swiss Plasma Center
!------------------------------------------------------------------------------
!
! MODULE: sort
!
!> @author
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> Module containing particle sorting algorithms.
!------------------------------------------------------------------------------
MODULE sort
IMPLICIT NONE

CONTAINS
!---------------------------------------------------------------------------  
!> @author 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief Sorts the particles according to their Z position using quick sort algorithm
!
!> @param[inout] p the particles array that need to be sorted
!> @param[in] leftlimit the left limit index for the sub array considered. 
!> @param[in] rightlimit the right limit index for the sub array considered.      
!---------------------------------------------------------------------------  
  RECURSIVE SUBROUTINE quicksortparts(p, leftlimit, rightlimit)
  Use beam, ONLY: particles, exchange_parts
  USE constants
  TYPE(particles), INTENT(INOUT):: p
  INTEGER,INTENT(IN):: leftlimit, rightlimit
  REAL(kind=db):: pivot
  INTEGER::i, cnt, mid
    IF(leftlimit .ge. rightlimit) RETURN ! Impossible indices, return
    
    mid=(leftlimit+rightlimit)/2 ! Compute middle index

    IF(p%Z(mid).lt.p%Z(leftlimit)) CALL exchange_parts(p,leftlimit,mid)
    IF(p%Z(rightlimit).lt.p%Z(leftlimit)) CALL exchange_parts(p,leftlimit,rightlimit)
    IF(p%Z(mid).lt.p%Z(rightlimit)) CALL exchange_parts(p,rightlimit,mid)

    ! Store the pivot point for comparison
    pivot=p%Z(rightlimit)

    cnt=leftlimit
    
    ! Move all parts with Z smaller than pivot to the left of pivot
    DO i=leftlimit, rightlimit
      IF(p%Z(i) .le. pivot) THEN
        CALL exchange_parts(p, i,cnt)
        cnt=cnt+1
      END IF
    END DO

    ! Quicksort the sub-arrays
    CALL quicksortparts(p, leftlimit,cnt-2)
    CALL quicksortparts(p, cnt,rightlimit)

  END SUBROUTINE quicksortparts

!---------------------------------------------------------------------------  
!> @author 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief Sorts the particles according to their linear index in the poisson solver grid.
!
!> @param[inout] p the particles array that need to be sorted
!> @param[in] leftlimit the left limit index for the sub array considered. 
!> @param[in] rightlimit the right limit index for the sub array considered.     
!---------------------------------------------------------------------------  
  RECURSIVE SUBROUTINE gridsort(p, leftlimit, rightlimit)
    Use beam, ONLY: particles, exchange_parts
  USE constants
  TYPE(particles), INTENT(INOUT):: p
  INTEGER,INTENT(IN):: leftlimit, rightlimit
  REAL(kind=db):: pivot
  INTEGER::i, cnt, mid
    IF(leftlimit .ge. rightlimit) RETURN ! Impossible indices, return
    
    mid=(leftlimit+rightlimit)/2 ! Compute middle index

    IF(linindex(p,mid).lt.linindex(p,leftlimit)) CALL exchange_parts(p,leftlimit,mid)
    IF(linindex(p,rightlimit).lt.linindex(p,leftlimit)) CALL exchange_parts(p,leftlimit,rightlimit)
    IF(linindex(p,mid).lt.linindex(p,rightlimit)) CALL exchange_parts(p,rightlimit,mid)

    ! Store the pivot point for comparison
    pivot=linindex(p,rightlimit)

    cnt=leftlimit
    
    ! Move all parts with Z smaller than pivot to the left of pivot
    DO i=leftlimit, rightlimit
      IF(linindex(p,i) .le. pivot) THEN
        CALL exchange_parts(p, i,cnt)
        cnt=cnt+1
      END IF
    END DO

    ! Quicksort the sub-arrays
    CALL gridsort(p, leftlimit,cnt-2)
    CALL gridsort(p, cnt,rightlimit)

  END SUBROUTINE gridsort

!---------------------------------------------------------------------------  
!> @author 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief Computes the linear index of a given particle in the Poisson solver 2D grid
!
!> @param[in] p the particles array where the particle is stored
!> @param[in] pid the index in the particle array of the particle of interest 
!> @param[out] linindex the computed linear index.     
!---------------------------------------------------------------------------  

  FUNCTION linindex(p,pid)
    USE beam, ONLY: particles
    USE basic, ONLY: nz 
    INTEGER :: linindex, pid
    TYPE(particles):: p
  
    linindex=p%Zindex(pid)+p%rindex(pid)*nz
  
  END FUNCTION linindex

END MODULE sort