{-# OPTIONS_HADDOCK showe-extensions #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -funbox-strict-fields #-}

-- |
-- Module      :  CaseBi.Unboxed
-- Copyright   :  (c) OleksandrZhabenko 2019-2021
-- License     :  MIT
-- Stability   :  Experimental
-- Maintainer  :  olexandr543@yahoo.com
--

module CaseBi.Unboxed (
-- * Function that can be used instead of @case ... of@ construction
--
-- > case var of 
-- >   a1 -> b1
-- >   a2 -> b2
-- >   a3 -> b3
-- >   ...
-- >   an -> bn
-- >   _  -> def
-- 
-- for efficiency or other data representation
  getBFst', getBFst, getBFstV, 
-- * Additional functions that are used to sort a list of pairs (which can be obtained e. g. by 'Prelude.zip')
  sortFst, sortFstV,
-- ** Function that can be used for changing the @Vector (a, b)@ during its creation 
  filterP
) where

import qualified Data.Vector.Unboxed as V (Vector,unsafeIndex,unsafeSlice,length,fromList,map)
import qualified Data.List as L (groupBy,nubBy)
import Data.Vector.Unboxed.Base
import GHC.Exts

-- | A variant of the 'CaseBi.getBFst'' function that operates on the unboxed 'V.Vector'.
-- 
getBFst' 
  :: (Ord a, Unbox a, Unbox b) 
  => (b, V.Vector (a, b)) -- ^ @b@ is a default value that can be substituted if there is no correspendence in the set of @(a, b)@ tuples (the 'otherwise' or irrefutable pattern analogue). Vector of the @(a, b)@ tuples that must be sorted in ascending order for the first argument. If there are several pairs @(a, b)@ with the same @a@, the function gives a resulting @b@ as if there is only the first one
  -> a -- ^ an element for which the corresponding resulting b must be found
  -> b -- ^ the result
getBFst' (def, vec) = getBFst'' (# 0, k #) (# j, m #) vec def
  where !j = V.length vec - 1
        !k = V.unsafeIndex vec 0
        !m = V.unsafeIndex vec j
{-# INLINE getBFst' #-}
                   
-- | A variant of the 'CaseBi.getBFstV' that operates on the unboxed 'V.Vector'.
-- 
getBFstV :: (Ord a, Unbox a, Unbox b) => V.Vector (a, b) -- ^ Vector of the @(a, b)@ tuples that are sorted in ascending order for the first argument
  -> b -- ^ a default value that can be substituted if there is no correspendence in the set of @(a, b)@ tuples (the 'otherwise' or irrefutable pattern analogue)
  -> V.Vector a -- ^ a Vector needed to be transformed accordingly to the correct @(a, b)@ tuple pairs
  -> V.Vector b -- ^ the resulting Vector
getBFstV c y = V.map (getBFst' (y, c))
{-# INLINE getBFstV #-}

-- | A variant of the 'CaseBi.getBFst' that operates on the unboxed 'V.Vector'.
-- 
getBFst :: (Ord a, Unbox a, Unbox b) => V.Vector (a, b) -- ^ Vector of the @(a, b)@ tuples that must be sorted in ascending order for the first argument
  -> b -- ^ a default value that can be substituted if there is no correspendence in the set of @(a, b)@ tuples (the 'otherwise' or irrefutable pattern analogue)
  -> [a] -- ^ a list of values needed to be transformed accordingly to the correct @(a, b)@ tuple pairs
  -> [b] -- ^ the resulting list
getBFst c y = map (getBFst' (y, c))
{-# INLINE getBFst #-}

-- | A variant of the 'CaseBi.sortFst' that operates on the unboxed 'V.Vector'. It is inspired by the work: https://wiki.haskell.org/Introduction
sortFst :: (Ord a, Unbox a, Unbox b) => [(a, b)] -> [(a, b)]
sortFst xs
 | null xs = []
 | otherwise = sortFst (filter (\(x, _) -> x < fst (head xs)) xs) ++ filter (\(x, _) -> x == (fst (head xs))) xs ++ 
     sortFst (filter (\(x, _) -> x > fst (head xs)) xs)
{-# INLINABLE sortFst #-}           

-- | A variant of the 'CaseBi.sortFstV' that operates on the unboxed 'V.Vector'.
--
sortFstV 
  :: (Ord a, Unbox a, Unbox b) => [(a, b)] -- ^ The list of conditions that is then converted to the corresponding Vector
   -> V.Vector (a, b) -- ^ the resulting sorted Vector that can be used further in getBFst' and its successors.
sortFstV = V.fromList . L.nubBy (\(x, _) (y, _) -> x == y) . sortFst
{-# INLINE sortFstV #-}

-- | A variant of the 'CaseBi.filterP' that operates on the unboxed 'V.Vector'.
filterP 
  :: (Ord a, Unbox a, Unbox b) => ((a, b) -> Bool) -- ^ The predicate @p@ used to select the only one value of @b@ in the pairs @(a, b)@ with the same @a@. 
  -- ^ If there are several pairs @(a, b)@ for the same @a@ that satisfies a predicate then the first one is used. For large @[(a, b)]@ 
  -- ^ it can be rather complex.
  -> [(a, b)] -- ^ The list of @(a, b)@ sorted in the ascending order by the first element a (e. g. by the 'sortFst' function)
  -> V.Vector (a, b) -- ^ The resulting filtered @Vector (a, b)@ that can be used for getFstB' and its successor functions.
--
filterP p xs = V.fromList . concatMap (\x -> take 1 . dropWhile (not . p) $ x) . L.groupBy (\(x1,_) (x2,_) -> x1 == x2) $ xs
{-# INLINE filterP #-}


{- | The function that can be used instead of the 'case ... of' function
@
case var of { a1 -> b1 ; a2 -> b2 ; a3 -> b3 ; ... ; an -> bn ; ~z -> defaultValue }
@
If we follow a lot of teaching materials that explain the workflow of the construction
we think that the complexity of it is about /O(n)/ for the transformation of @a@ to @b@ here.
David Feuer (david.feuer (at) gmail.com) said that 'case ... of' is already optimized in GHC.
Some benchmarks show that its behaviour  tends to be about of /O(log n)/ complexity, the same as
the proposed function 'getBFst''. Nevertheless, the last one shows better performance in some situations,
is rather general and can be used for another data representation. Therefore, it can be preferred in
some situations. 'getBFst'' uses binary search algorithm and an 'V.Vector' of @(a, b)@ as somewhat like a complicated filter or like a special sieve. The 'V.Vector' must be sorted in ascending order here for the algorithm to be used correctly. For this you can use the function 'sortFstV' or the similar ones. 
-}
getBFst''
  :: (Ord a, Unbox a, Unbox b) => (# Int, (a, b) #) -- ^ The first unboxed tuple of the index and the element of the array.
  -> (# Int, (a, b) #) -- ^ The second unboxed tuple of the index and the element of the array.
  -> V.Vector (a, b) -- ^ The array of the pairs of the compared value and the result that is used in case the last argument is equal to the compared value.
  -> b -- ^ The default value that is used if no first element in the array tuples equals to the compared value.
  -> a -- ^ The compared value, well, the @main@ function argument, to which it is applied.
  -> b -- ^ The resulting branch value.
getBFst'' (# (I# i#), k #) (# (I# j#), m #) vec def x
 | if x < fst k then True else x > fst m = def
 | otherwise = gBF3 (# i# , k #) (# j# , m #) vec def x
{-# INLINE getBFst'' #-}

-- | The meaning of the arguments is the same as for 'getBFst'''. Is used internally in it.
gBF3 :: (Ord a, Unbox a, Unbox b) => (# Int#, (a, b) #) -> (# Int#, (a, b) #) -> V.Vector (a, b) -> b -> a -> b
gBF3 (# !i#, k #) (# !j#, m #) vec def x
 | isTrue# ((j# -# i#) ># 1# ) = 
    case compare x (fst p) of
     GT -> gBF3 (# n#, p #) (# j#, m #) vec def x
     LT  -> gBF3 (# i#, k #) (# n#, p #) vec def x
     _ -> snd p
 | x == fst m = snd m
 | x == fst k = snd k
 | otherwise = def
     where !n# = (i# +# j#) `quotInt#` 2#
           !p = V.unsafeIndex vec (I# n#)
{-# INLINABLE gBF3 #-}