content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE ExplicitForAll #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE MonoLocalBinds #-}
  
  -- |
  -- Module      : Data.Massiv.Array.Ops.Sort
  -- Copyright   : (c) Alexey Kuleshevich 2018-2022
  -- License     : BSD3
  -- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
  -- Stability   : experimental
  -- Portability : non-portable
  module Data.Massiv.Array.Ops.Sort (
    tally,
    quicksort,
    quicksortBy,
    quicksortByM,
    quicksortAs,
    quicksortAsBy,
    quicksortAsByM,
    quicksortM_,
    quicksortByM_,
    unsafeUnstablePartitionRegionM,
  ) where
  
  import Control.Monad (when)
  import Control.Monad.IO.Unlift
  import Control.Monad.Primitive
  import Control.Scheduler
  import Data.Bits (countLeadingZeros)
  import Data.Massiv.Array.Delayed.Stream
  import Data.Massiv.Array.Mutable
  import Data.Massiv.Array.Ops.Transform
  import Data.Massiv.Core.Common
  import Data.Massiv.Vector (scatMaybes, sunfoldrN)
  import Data.Word (Word64)
  import System.IO.Unsafe
  
  -- | Count number of occurrences of each element in the array. Results will be
  -- sorted in ascending order of the element.
  --
  -- ==== __Example__
  --
  -- >>> import Data.Massiv.Array as A
  -- >>> xs = fromList Seq [2, 4, 3, 2, 4, 5, 2, 1] :: Array P Ix1 Int
  -- >>> xs
  -- Array P Seq (Sz1 8)
  --   [ 2, 4, 3, 2, 4, 5, 2, 1 ]
  -- >>> tally xs
  -- Array DS Seq (Sz1 5)
  --   [ (1,1), (2,3), (3,1), (4,2), (5,1) ]
  --
  -- @since 0.4.4
  tally :: (Manifest r e, Load r ix e, Ord e) => Array r ix e -> Vector DS (e, Int)
  tally arr
    | isEmpty arr = setComp (getComp arr) empty
    | otherwise = scatMaybes $ sunfoldrN (liftSz2 (+) sz oneSz) count (0, 0, sorted ! 0)
    where
      sz@(Sz k) = size sorted
      count (!i, !n, !prev)
        | i < k =
            let !e' = unsafeLinearIndex sorted i
             in if prev == e'
                  then Just (Nothing, (i + 1, n + 1, prev))
                  else Just (Just (prev, n), (i + 1, 1, e'))
        | otherwise = Just (Just (prev, n), (i + 1, n, prev))
      {-# INLINE count #-}
      sorted = quicksort $ flatten arr
  {-# INLINE tally #-}
  
  -- | Partition a segment of a vector. Starting and ending indices are unchecked.
  --
  -- @since 1.0.0
  unsafeUnstablePartitionRegionM
    :: forall r e m
     . (Manifest r e, PrimMonad m)
    => MVector (PrimState m) r e
    -> (e -> m Bool)
    -> Ix1
    -- ^ Start index of the region
    -> Ix1
    -- ^ End index of the region
    -> m Ix1
  unsafeUnstablePartitionRegionM marr f start end = fromLeft start (end + 1)
    where
      fromLeft i j
        | i == j = pure i
        | otherwise = do
            e <- f =<< unsafeLinearRead marr i
            if e
              then fromLeft (i + 1) j
              else fromRight i (j - 1)
      fromRight i j
        | i == j = pure i
        | otherwise = do
            x <- unsafeLinearRead marr j
            e <- f x
            if e
              then do
                unsafeLinearWrite marr j =<< unsafeLinearRead marr i
                unsafeLinearWrite marr i x
                fromLeft (i + 1) j
              else fromRight i (j - 1)
  {-# INLINE unsafeUnstablePartitionRegionM #-}
  
  -- | Same as `quicksort` except it accepts any array that is computable.
  --
  -- @since 1.0.2
  quicksortAs
    :: (Load r Ix1 e, Manifest r' e, Ord e) => r' -> Vector r e -> Vector r' e
  quicksortAs _ arr = unsafePerformIO $ withLoadMArray_ arr quicksortM_
  {-# INLINE quicksortAs #-}
  
  -- | Same as `quicksortBy` except it accepts any array that is computable.
  --
  -- @since 1.0.2
  quicksortAsBy
    :: (Load r Ix1 e, Manifest r' e) => r' -> (e -> e -> Ordering) -> Vector r e -> Vector r' e
  quicksortAsBy _ f arr =
    unsafePerformIO $ withLoadMArray_ arr (quicksortByM_ (\x y -> pure $ f x y))
  {-# INLINE quicksortAsBy #-}
  
  -- | Same as `quicksortByM` except it accepts any array that is computable.
  --
  -- @since 1.0.2
  quicksortAsByM
    :: (Load r Ix1 e, Manifest r' e, MonadUnliftIO m)
    => r'
    -> (e -> e -> m Ordering)
    -> Vector r e
    -> m (Vector r' e)
  quicksortAsByM _ f arr =
    withRunInIO $ \run -> withLoadMArray_ arr (quicksortByM_ (\x y -> run (f x y)))
  {-# INLINE quicksortAsByM #-}
  
  -- | This is an implementation of
  -- [Quicksort](https://en.wikipedia.org/wiki/Quicksort), which is an efficient,
  -- but unstable sort. This implementation uses Median-of-three for pivot
  -- choosing, as such it performs very well not only for random values, but also
  -- for common edge cases like already sorted, reversed sorted and arrays with
  -- many duplicate elements. It will also respect the computation strategy and
  -- will result in a nice speed up for systems with multiple CPUs.
  --
  -- @since 0.3.2
  quicksort
    :: (Manifest r e, Ord e) => Vector r e -> Vector r e
  quicksort arr = unsafePerformIO $ withMArray_ arr quicksortM_
  {-# INLINE quicksort #-}
  
  -- | Same as `quicksortBy`, but instead of `Ord` constraint expects a custom `Ordering`.
  --
  -- @since 0.6.1
  quicksortByM
    :: (Manifest r e, MonadUnliftIO m) => (e -> e -> m Ordering) -> Vector r e -> m (Vector r e)
  quicksortByM f arr = withRunInIO $ \run -> withMArray_ arr (quicksortByM_ (\x y -> run (f x y)))
  {-# INLINE quicksortByM #-}
  
  -- | Same as `quicksortBy`, but instead of `Ord` constraint expects a custom `Ordering`.
  --
  -- @since 0.6.1
  quicksortBy :: Manifest r e => (e -> e -> Ordering) -> Vector r e -> Vector r e
  quicksortBy f arr =
    unsafePerformIO $ withMArray_ arr (quicksortByM_ (\x y -> pure $ f x y))
  {-# INLINE quicksortBy #-}
  
  -- | Manifest version of `quicksort`
  --
  -- @since 0.3.2
  quicksortM_
    :: (Ord e, Manifest r e, MonadPrimBase s m)
    => Scheduler s ()
    -> MVector s r e
    -> m ()
  quicksortM_ = quicksortInternalM_ (\e1 e2 -> pure $ e1 < e2) (\e1 e2 -> pure $ e1 == e2)
  {-# INLINE quicksortM_ #-}
  
  -- | Same as `quicksortM_`, but instead of `Ord` constraint expects a custom `Ordering`.
  --
  -- @since 0.6.1
  quicksortByM_
    :: (Manifest r e, MonadPrimBase s m)
    => (e -> e -> m Ordering)
    -> Scheduler s ()
    -> MVector s r e
    -> m ()
  quicksortByM_ compareM =
    quicksortInternalM_ (\x y -> (LT ==) <$> compareM x y) (\x y -> (EQ ==) <$> compareM x y)
  {-# INLINE quicksortByM_ #-}
  
  quicksortInternalM_
    :: (Manifest r e, MonadPrimBase s m)
    => (e -> e -> m Bool)
    -> (e -> e -> m Bool)
    -> Scheduler s ()
    -> MVector s r e
    -> m ()
  quicksortInternalM_ fLT fEQ scheduler marr
    | numWorkers scheduler < 2 || depthPar <= 0 = qsortSeq 0 (k - 1)
    | otherwise = qsortPar depthPar 0 (k - 1)
    where
      -- How deep into the search tree should we continue scheduling jobs. Constants below
      -- were discovered empirically:
      depthPar = min (logNumWorkers + 4) (logSize - 10)
      k = unSz (sizeOfMArray marr)
      -- We must use log becuase decinding into a tree creates an exponential number of jobs
      logNumWorkers = 63 - countLeadingZeros (fromIntegral (numWorkers scheduler) :: Word64)
      -- Using many cores on small vectors only makes things slower
      logSize = 63 - countLeadingZeros (fromIntegral k :: Word64)
      ltSwap i j = do
        ei <- unsafeLinearRead marr i
        ej <- unsafeLinearRead marr j
        lt <- fLT ei ej
        if lt
          then do
            unsafeLinearWrite marr i ej
            unsafeLinearWrite marr j ei
            pure ei
          else pure ej
      {-# INLINE ltSwap #-}
      getPivot lo hi = do
        let !mid = (hi + lo) `div` 2
        _ <- ltSwap mid lo
        _ <- ltSwap hi lo
        ltSwap mid hi
      {-# INLINE getPivot #-}
      qsortPar !n !lo !hi =
        when (lo < hi) $ do
          p <- getPivot lo hi
          l <- unsafeUnstablePartitionRegionM marr (`fLT` p) lo (hi - 1)
          h <- unsafeUnstablePartitionRegionM marr (`fEQ` p) l hi
          if n > 0
            then do
              let !n' = n - 1
              scheduleWork scheduler $ qsortPar n' lo (l - 1)
              scheduleWork scheduler $ qsortPar n' h hi
            else do
              qsortSeq lo (l - 1)
              qsortSeq h hi
      qsortSeq !lo !hi =
        when (lo < hi) $ do
          p <- getPivot lo hi
          l <- unsafeUnstablePartitionRegionM marr (`fLT` p) lo (hi - 1)
          h <- unsafeUnstablePartitionRegionM marr (`fEQ` p) l hi
          qsortSeq lo (l - 1)
          qsortSeq h hi
  {-# INLINE quicksortInternalM_ #-}
  