{- Copyright 2015-2019 NGLess Authors
 - License: MIT
 -}

module Utils.Vector
 ( unsafeIncrement
 , unsafeIncrement'

 , binarySearch
 , binarySearchBy
 , binarySearchByRange
 , binarySearchByExact

 , binaryFindBy

 , sortParallel

 , withVector
 ) where

import           Control.Monad (guard)
import           Control.Monad.Primitive (PrimMonad(..))
import           Control.Monad.IO.Class (MonadIO(..), liftIO)
import qualified Control.Concurrent.Async as A
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import qualified Data.Vector.Generic.Mutable as VGM
import           Data.Vector.Algorithms.Intro (sortByBounds)

-- | increment by 1
unsafeIncrement :: (Num a, PrimMonad m, VUM.Unbox a) => VUM.MVector (PrimState m) a -> Int -> m ()
unsafeIncrement v i = unsafeIncrement' v i 1
{-# INLINE unsafeIncrement #-}

-- | increment by given value
unsafeIncrement' :: (Num a, PrimMonad m, VUM.Unbox a) => VUM.MVector (PrimState m) a -> Int -> a -> m ()
unsafeIncrement' v i inc = VUM.unsafeModify v (+ inc) i
{-# INLINE unsafeIncrement' #-}

binarySearch :: (Ord a) => V.Vector a -> a -> Int
binarySearch v = binarySearchByRange 0 (V.length v) compare v

binarySearchByRange :: Int -> Int -> (a -> b -> Ordering) -> V.Vector a -> b -> Int
binarySearchByRange start end comp v target = loop start end
    where
        loop :: Int -> Int -> Int
        loop s e
            | s == e = s
            | otherwise = let mid = (s + e) `div` 2
                            in if comp (V.unsafeIndex v mid) target == LT
                                    then loop (mid + 1) e
                                    else loop s mid
binarySearchBy :: (a -> b -> Ordering) -> V.Vector a -> b -> Int
binarySearchBy comp v = binarySearchByRange 0 (V.length v) comp v

binarySearchByExact :: (a -> b -> Ordering) -> V.Vector a -> b -> Maybe Int
binarySearchByExact comp v target = do
    let ix = binarySearchBy comp v target
    guard (ix < V.length v)
    guard (comp (V.unsafeIndex v ix) target == EQ)
    return ix

binaryFindBy :: (a -> b -> Ordering) -> V.Vector a -> b -> Maybe a
binaryFindBy comp v target = do
    ix <- binarySearchByExact comp v target
    return $! V.unsafeIndex v ix


-- sort a vector in parallel threads
sortParallel :: (Ord e) =>
                        Int -- ^ nr of threads
                        -> VM.IOVector e -- ^ vector
                        -> IO ()
sortParallel n v = sortPByBounds n v 0 (VM.length v)

sortPByBounds :: (Ord e) => Int -> VM.IOVector e -> Int -> Int -> IO ()
sortPByBounds 1 v start end = sortByBounds compare v start end
sortPByBounds threads v start end
    | end - start < 1024 = sortByBounds compare v start end
    | otherwise = do
        mid <- pivot v start end
        k <- unstablePartition (< mid) v start end
        let t1 :: Int
            t1
                | k - start < 1024 = 1
                | end - k < 1024 = threads - 1
                | otherwise = threads `div` 2
            t2 = threads - t1
        A.concurrently_
            (sortPByBounds t1 v start k)
            (sortPByBounds t2 v k end)

pivot v start end = do
        a <- VM.read v start
        b <- VM.read v (end - 1)
        c <- VM.read v ((start + end) `div` 2)
        return $! median a b c
    where
        median a b c = if a <= b
            then min b c
            else max a c

unstablePartition f v start end = (+ start) <$> VGM.unstablePartition f (VM.unsafeSlice start (end - start) v)


withVector :: (MonadIO m, VUM.Unbox a) => VUM.IOVector a -> (VU.Vector a -> b) -> m b
withVector v f = liftIO $ do
    v' <- VU.unsafeFreeze v
    return $! f v'