module Data.Vector.Split
  ( chunksOf
  , splitPlaces
  , splitPlacesBlanks
  , chop
  , divvy
  , module Data.Vector.Split.Internal
  ) where


import           Data.Vector.Generic (Vector)
import qualified Data.Vector.Generic as V

import           Data.List (unfoldr)

import           Data.Vector.Split.Internal



-- | @'chunksOf' n@ splits a vector into length-n pieces.  The last
--   piece will be shorter if @n@ does not evenly divide the length of
--   the vector.  If @n <= 0@, @'chunksOf' n l@ returns an infinite list
--   of empty vectors.  For example:
--
--   Note that @'chunksOf' n []@ is @[]@, not @[[]]@.  This is
--   intentional, and is consistent with a recursive definition of
--   'chunksOf'; it satisfies the property that
--
--   @chunksOf n xs ++ chunksOf n ys == chunksOf n (xs ++ ys)@
--
--   whenever @n@ evenly divides the length of @xs@.
chunksOf :: Vector v a => Int -> v a -> [v a]
chunksOf :: forall (v :: * -> *) a. Vector v a => Int -> v a -> [v a]
chunksOf Int
i = forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr forall {v :: * -> *} {a}. Vector v a => v a -> Maybe (v a, v a)
go
  where go :: v a -> Maybe (v a, v a)
go v a
v | forall (v :: * -> *) a. Vector v a => v a -> Bool
V.null v a
v = forall a. Maybe a
Nothing
             | Bool
otherwise = forall a. a -> Maybe a
Just (forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
V.splitAt Int
i v a
v)

-- | Split a vector into chunks of the given lengths. For example:
--
-- > splitPlaces [2,3,4] [1..20] == [[1,2],[3,4,5],[6,7,8,9]]
-- > splitPlaces [4,9] [1..10] == [[1,2,3,4],[5,6,7,8,9,10]]
-- > splitPlaces [4,9,3] [1..10] == [[1,2,3,4],[5,6,7,8,9,10]]
--
--   If the input vector is longer than the total of the given lengths,
--   then the remaining elements are dropped. If the vector is shorter
--   than the total of the given lengths, then the result may contain
--   fewer chunks than requested, and the last chunk may be shorter
--   than requested.
splitPlaces :: Vector v a => [Int] -> v a -> [v a]
splitPlaces :: forall (v :: * -> *) a. Vector v a => [Int] -> v a -> [v a]
splitPlaces [Int]
is v a
v = forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr forall {v :: * -> *} {a}.
Vector v a =>
([Int], v a) -> Maybe (v a, ([Int], v a))
go ([Int]
is,v a
v)
  where go :: ([Int], v a) -> Maybe (v a, ([Int], v a))
go ([],v a
_)   = forall a. Maybe a
Nothing
        go (Int
x:[Int]
xs,v a
y) | forall (v :: * -> *) a. Vector v a => v a -> Bool
V.null v a
y = forall a. Maybe a
Nothing
                    | Bool
otherwise = let (v a
l,v a
r) = forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
V.splitAt Int
x v a
y in forall a. a -> Maybe a
Just (v a
l,([Int]
xs,v a
r))


-- | Split a vector into chunks of the given lengths. Unlike
--   'splitPlaces', the output list will always be the same length as
--   the first input argument. If the input vector is longer than the
--   total of the given lengths, then the remaining elements are
--   dropped. If the vector is shorter than the total of the given
--   lengths, then the last several chunks will be shorter than
--   requested or empty. For example:
--
-- > splitPlacesBlanks [2,3,4] [1..20] == [[1,2],[3,4,5],[6,7,8,9]]
-- > splitPlacesBlanks [4,9] [1..10] == [[1,2,3,4],[5,6,7,8,9,10]]
-- > splitPlacesBlanks [4,9,3] [1..10] == [[1,2,3,4],[5,6,7,8,9,10],[]]
--
--   Notice the empty list in the output of the third example, which
--   differs from the behavior of 'splitPlaces'.
splitPlacesBlanks :: Vector v a => [Int] -> v a -> [v a]
splitPlacesBlanks :: forall (v :: * -> *) a. Vector v a => [Int] -> v a -> [v a]
splitPlacesBlanks [Int]
is v a
v = forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr forall {v :: * -> *} {a}.
Vector v a =>
([Int], v a) -> Maybe (v a, ([Int], v a))
go ([Int]
is,v a
v)
  where go :: ([Int], v a) -> Maybe (v a, ([Int], v a))
go ([],v a
_)   = forall a. Maybe a
Nothing
        go (Int
x:[Int]
xs,v a
y) = let (v a
l,v a
r) = forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
V.splitAt Int
x v a
y in forall a. a -> Maybe a
Just (v a
l,([Int]
xs,v a
r))



-- | A useful recursion pattern for processing a list to produce a new
--   list, often used for \"chopping\" up the input list.  Typically
--   chop is called with some function that will consume an initial
--   prefix of the list and produce a value and the rest of the list.
--
--   For example, many common Prelude functions can be implemented in
--   terms of @chop@:
--
-- > group :: (Eq a) => [a] -> [[a]]
-- > group = chop (\ xs@(x:_) -> span (==x) xs)
-- >
-- > words :: String -> [String]
-- > words = filter (not . null) . chop (span (not . isSpace) . dropWhile isSpace)
chop :: Vector v a => (v a -> (b, v a)) -> v a -> [b]
chop :: forall (v :: * -> *) a b.
Vector v a =>
(v a -> (b, v a)) -> v a -> [b]
chop v a -> (b, v a)
f v a
v | forall (v :: * -> *) a. Vector v a => v a -> Bool
V.null v a
v  = []
         | Bool
otherwise = b
b forall a. a -> [a] -> [a]
: forall (v :: * -> *) a b.
Vector v a =>
(v a -> (b, v a)) -> v a -> [b]
chop v a -> (b, v a)
f v a
v'
            where (b
b, v a
v') = v a -> (b, v a)
f v a
v

-- | Divides up an input vector into a set of subvectors, according to 'n' and 'm'
--   input specifications you provide. Each subvector will have 'n' items, and the
--   start of each subvector will be offset by 'm' items from the previous one.
--
-- > divvy 5 5 [1..20] == [[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15],[16,17,18,19,20]]
--
--   In the case where a source vector's trailing elements do no fill an entire
--   subvector, those trailing elements will be dropped.
--
-- > divvy 5 2 [1..10] == [[1,2,3,4,5],[3,4,5,6,7],[5,6,7,8,9]]
--
--   As an example, you can generate a moving average over a vector of prices:
-- 
-- > type Prices = [Float]
-- > type AveragePrices = [Float]
-- > 
-- > average :: [Float] -> Float
-- > average xs = sum xs / (fromIntegral $ length xs)
-- > 
-- > simpleMovingAverage :: Prices -> AveragePrices
-- > simpleMovingAverage priceList =
-- >   map average divvyedPrices
-- >     where divvyedPrices = divvy 20 1 priceList
divvy :: Vector v a => Int -> Int -> v a -> [v a]
divvy :: forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> [v a]
divvy Int
n Int
m v a
v | forall (v :: * -> *) a. Vector v a => v a -> Bool
V.null v a
v = []
            | Bool
otherwise = forall a. (a -> Bool) -> [a] -> [a]
filter (\v a
ws -> Int
n forall a. Eq a => a -> a -> Bool
== forall (v :: * -> *) a. Vector v a => v a -> Int
V.length v a
ws)
                        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a b.
Vector v a =>
(v a -> (b, v a)) -> v a -> [b]
chop (\v a
xs -> (forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
V.take Int
n v a
xs, forall (v :: * -> *) a. Vector v a => Int -> v a -> v a
V.drop Int
m v a
xs))
                        forall a b. (a -> b) -> a -> b
$ v a
v