{-# LANGUAGE CPP           #-}
{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE TypeOperators #-}
-----------------------------------------------------------------------------
-- |
-- Module      : Data.Array.Vector.Strict.Basics
-- Copyright   : (c) [2001..2002] Manuel M T Chakravarty & Gabriele Keller
--               (c) 2006         Manuel M T Chakravarty & Roman Leshchinskiy
-- License     : see libraries/ndp/LICENSE
-- 
-- Maintainer  : Roman Leshchinskiy <rl@cse.unsw.edu.au>
-- Stability   : internal
-- Portability : portable
--
-- Description ---------------------------------------------------------------
--
--  Basic operations on flat unlifted arrays.
--
-- Todo ----------------------------------------------------------------------
--


#include "fusion-phases.h"

module Data.Array.Vector.Strict.Basics where

import Data.Array.Vector.Stream
import Data.Array.Vector.UArr hiding (lengthU, indexU)
import qualified Data.Array.Vector.UArr as Prim (lengthU, indexU)

import Data.Array.Vector.Strict.Stream

import Data.Array.Vector.Prim.Debug
import Data.Array.Vector.Prim.Hyperstrict
import GHC.ST

import Debug.Trace

------------------------------------------------------------------------

instance (Eq e, UA e) => Eq (UArr e) where (==) = eqU
   -- not really fusible

{-# INLINE_U eqU #-}
eqU :: (Eq e, UA e) => UArr e -> UArr e -> Bool
eqU a1 a2 = lengthU a1 == lengthU a2 && foldlU cmp True (zipU a1 a2)
   where
       cmp r (e1 :*: e2) = e1 == e2 && r

------------------------------------------------------------------------

-- XXX general rule:
-- If we have direct implementations with better complexity than streams:
--
--   length, index, take, drop
--
-- Then use the direct version


-- |/O(1)/. 'lengthU' returns the length of a 'UArr' as an 'Int'.
lengthU :: UA e => UArr e -> Int
lengthU = foldlU (const . (+1)) 0
-- lengthU = Prim.lengthU
{-# INLINE_U lengthU #-}

{-

Unfused version:

$ time ./length
100000000
./length  1.10s user 1.13s system 91% cpu 2.433 total

Fusible version:

100000000
./length  0.31s user 0.00s system 97% cpu 0.318 total

-}

-- lengthU_stream :: UA e => UArr e -> Int
-- {-# INLINE lengthU_stream #-}

-- lengthU is reexported from UArr

-- |/O(1)/. 'nullU' tests whether the given array is empty.
--
nullU :: UA e => UArr e -> Bool
nullU  = nullS . streamU  -- better code if we short circuit
{-# INLINE_U nullU #-}
-- nullU  = (== 0) . lengthU

-- |/O(1)/. 'emptyU' yields an empty array.
--
emptyU :: UA e => UArr e
emptyU = unstreamU emptyS
{-# INLINE_U emptyU #-}

-- emptyU = newU 0 (const $ return ())

-- |/O(1)/. 'singletonU' yields a singleton array containing the given element.
--
singletonU :: UA e => e -> UArr e
{-# INLINE_U singletonU #-}
singletonU = unstreamU . singletonS

-- |/O(n)/. 'consU' prepends the given element to an array.
--
consU :: UA e => e -> UArr e -> UArr e
{-# INLINE_U consU #-}
consU x = unstreamU . consS x . streamU

-- |/O(n)/. 'snocU' appends the given element to an array.
--
snocU :: UA e => UArr e -> e -> UArr e
{-# INLINE_U snocU #-}
snocU s x = unstreamU (snocS (streamU s) x)

-- unitsU is reexported from Loop

-- |/O(n)/. @'replicateU' n e@ yields an array containing @n@ repetitions of @e@.
--
replicateU :: UA e => Int -> e -> UArr e
{-# INLINE_U replicateU #-}
replicateU n e = unstreamU (replicateS n e)

-- |/O(n)/. @'replicateEachU' n r e@ yields an array such that each element in
-- @e@ is repeated as many times as the value contained at the corresponding
-- index in @r@. For example:
--
-- @replicateEachU 10 (toU [1..3]) (toU [3..5])@ yields @toU [3.0,4.0,4.0,5.0,5.0,5.0]@
--
-- /N.B/: the @n@ parameter specifies how many elements are /allocated/ for the
-- output array, but the function will happily overrun the allocated buffer for
-- all sorts of interesting effects! The caller is expected to ensure that 
-- @n <= sumU r@.
--
replicateEachU :: UA e => Int -> UArr Int -> UArr e -> UArr e
{-# INLINE_U replicateEachU #-}
replicateEachU n ns es = unstreamU
                       . replicateEachS n
                       $ zipS (streamU ns) (streamU es)

-- |/O(n)/. 'indexU' extracts an element out of an immutable unboxed array.
--
indexU :: UA e => UArr e -> Int -> e
indexU arr n = indexS (streamU arr) n
{-# INLINE_U indexU #-}

-- |/O(1)/. 'headU' yields the first element of an array.
--
headU :: UA e => UArr e -> e
headU = headS . streamU
{-# INLINE_U headU #-}

-- |/O(n)/. 'lastU' yields the last element of an array.
--
lastU :: UA e => UArr e -> e
lastU = foldl1U (flip const)
{-# INLINE lastU #-}

-- |/O(n)/. 'appendU' concatenates two arrays.
--
appendU :: UA e => UArr e -> UArr e -> UArr e
{-# INLINE_U appendU #-}
a1 `appendU` a2 = unstreamU (streamU a1 +++ streamU a2)

-- |/O(n)/. Concatenate a list of arrays.
concatU :: UA e => [UArr e] -> UArr e
concatU as = unstreamU (foldr (+++) emptyS (map streamU as))
{-# INLINE concatU #-}

-- |/O(n)/. 'initU' yields the input array without its last element.
initU :: UA e => UArr e -> UArr e -- not unboxing
initU = unstreamU . initS . streamU
{-# INLINE initU #-}

-- |/O(n)/. @'repeatU' n u@ repeats an array @u@ @n@ times.
--
repeatU :: UA e => Int -> UArr e -> UArr e
repeatU n = unstreamU . repS n
{-# INLINE_U repeatU #-}

-- No work duplicated
repS :: UA e => Int -> UArr e -> Stream e
{-# INLINE_STREAM repS #-}
repS k xs = Stream next (0 :*: k) (k*n)
  where
    n = lengthU xs

    {-# INLINE next #-}
    next (i :*: 0) = Done
    next (i :*: k) | i == n    = Skip (0 :*: k-1)
                   | otherwise = Yield (xs `Prim.indexU` i) (i+1 :*: k)

-- *Indexing
-- ---------

-- |/O(n)/. 'indexedU' associates each element of the array with its index.
--
indexedU :: UA e => UArr e -> UArr (Int :*: e)
{-# INLINE_U indexedU #-}
indexedU = unstreamU . indexedS . streamU

-- *Conversion
-- -----------

-- |/O(n)/. 'toU' turns a list into a 'UArr'.
--
toU :: UA e => [e] -> UArr e
{-# INLINE_U toU #-}
toU = unstreamU . toStream

-- |/O(n)/. 'fromU' collects the elements of a 'UArr' into a list.
--
fromU :: UA e => UArr e -> [e]
{-# INLINE_U fromU #-}
fromU a = [a `Prim.indexU` i | i <- [0..lengthU a - 1]]

------------------------------------------------------------------------


here s = "Data.Array.Vector.Strict.Combinators." ++ s

-- |/O(n)/. @'iterateU' n f a@ constructs an array of size @n@ by iteratively 
-- applying @f@ to @a@.
--
iterateU :: (UA a) => Int -> (a -> a) -> a -> UArr a
iterateU n f = unfoldU n (\x -> JustS $ x :*: f x)

-- |/O(n)/. 'mapU' maps a function over an array.
--
mapU :: (UA e, UA e') => (e -> e') -> UArr e -> UArr e'
{-# INLINE_U mapU #-}
mapU f = unstreamU . mapS f . streamU

-- |/O(n)/. 'filterU' extracts all elements from an array that satisfy 
-- the given predicate.
--
filterU :: UA e => (e -> Bool) -> UArr e -> UArr e 
{-# INLINE_U filterU #-}
filterU p = unstreamU . filterS p . streamU

-- |/O(n)/. 'packU' extracts all elements from an array according to the 
-- provided flag array. For example:
--
-- @packU (toU [1..5]) (toU [True,False,False,False,True])@
--
-- yields @toU [1.0,5.0]@.
-- 
packU:: UA e => UArr e -> UArr Bool -> UArr e
{-# INLINE_U packU #-}
packU xs = fstU . filterU sndS . zipU xs


-- |/O(n)/. 'foldlU' reduces an array proceeding from the left.
--
foldlU :: UA a => (b -> a -> b) -> b -> UArr a -> b
{-# INLINE_U foldlU #-}
foldlU f z = foldS f z . streamU

-- |/O(n)/. 'foldl1U' is a variant of 'foldlU' that assumes a non-empty input
-- array, but requires no starting element. Throws an exception if the input
-- array is empty.
--

-- FIXME: Rewrite for 'Stream's.

-- foldl1U :: UA a => (a -> a -> a) -> UArr a -> a
-- {-# INLINE_U foldl1U #-}
-- foldl1U f arr = checkNotEmpty (here "foldl1U") (lengthU arr) $
--                 foldlU f (arr `Prim.indexU` 0) (sliceU arr 1 (lengthU arr - 1))

foldl1U :: UA a => (a -> a -> a) -> UArr a -> a
foldl1U f = foldl1S f . streamU
{-# INLINE foldl1U #-}

-- | /O(n)/ 'foldrU', applied to a binary operator, a starting value
-- (typically the right-identity of the operator), and a 'UArr a',
-- reduces the 'UArr a' using the binary operator, from right to left.
foldrU :: UA a => (a -> b -> b) -> b -> UArr a -> b
foldrU f z = foldrS f z . streamU
{-# INLINE foldrU #-}

-- | /O(n)/ A variant of 'foldr' that has no starting value argument,
-- and thus must be applied to a non-empty 'UArr a'.
foldr1U :: UA a => (a -> a -> a) -> UArr a -> a
foldr1U f = foldr1S f . streamU
{-# INLINE foldr1U #-}

-- |/O(n)/. 'foldl1MaybeU' behaves like 'foldl1U' but returns 'NothingS' if the
-- input array is empty.
foldl1MaybeU :: UA a => (a -> a -> a) -> UArr a -> MaybeS a
{-# INLINE_U foldl1MaybeU #-}
foldl1MaybeU f = fold1MaybeS f . streamU

-- |/O(n)/. 'foldU' reduces an array using an associative combination function
-- and its unit.
--
foldU :: UA a => (a -> a -> a) -> a -> UArr a -> a
{-# INLINE_U foldU #-}
foldU = foldlU

-- |/O(n)/. 'fold1MaybeU' behaves like 'fold1U' but returns 'NothingS' if the
-- input array is empty.
fold1MaybeU :: UA a => (a -> a -> a) -> UArr a -> MaybeS a
{-# INLINE_U fold1MaybeU #-}
fold1MaybeU = foldl1MaybeU

-- |/O(n)/. 'fold1U' is a variant of 'foldU' that requires a non-empty input
-- array. Throws an exception if its input array is empty.
--
fold1U :: UA a => (a -> a -> a) -> UArr a -> a
{-# INLINE_U fold1U #-}
fold1U = foldl1U

-- |/O(n)/. 'scanlU' is equivalent to 'foldlU' on all prefixes (except the
-- array itself) of the input array.
--
-- /N.B/: the behavior of this function differs from that of Data.List. Compare:
--
-- @scanl (+) 0.0 [1..5]@ gives @[0.0,1.0,3.0,6.0,10.0,15.0]@
--
-- @scanlU (+) 0.0 $ toU [1..5]@ gives @toU [0.0,1.0,3.0,6.0,10.0]@
--
-- To get behavior closer to the List counterpart, see 'scanResU'.
--
scanlU :: (UA a, UA b) => (b -> a -> b) -> b -> UArr a -> UArr b
{-# INLINE_U scanlU #-}
scanlU f z = unstreamU . scanS f z . streamU

-- |/O(n)/. 'scanl1U' is like 'scanlU', but requires no starting value.
--
scanl1U :: UA a => (a -> a -> a) -> UArr a -> UArr a
{-# INLINE_U scanl1U #-}
scanl1U f arr = checkNotEmpty (here "scanl1U") (lengthU arr) $
                unstreamU (scan1S f (streamU arr))

-- |/O(n)/. 'scanU' is equivalent to 'foldU' on all prefixes (except the array
-- itself) of the input array.
--
scanU :: UA a => (a -> a -> a) -> a -> UArr a -> UArr a
{-# INLINE_U scanU #-}
scanU = scanlU

-- |/O(n)/. 'scan1U' is like 'scanU', but requires no starting value.
--
scan1U :: UA a => (a -> a -> a) -> UArr a -> UArr a
{-# INLINE_U scan1U #-}
scan1U = scanl1U

-- |/O(n)/. 'scanResU' behaves like 'scanU', but yields a strict pair with the
-- 'scanU' result as its @fstS@ and the "missing" element ('foldU' on the same
-- arguments) as its @sndS@. Compare:
--
-- @scanl (+) 0.0 [1..5]@ gives @[0.0,1.0,3.0,6.0,10.0,15.0]@
--
-- @scanlU (+) 0.0 $ toU [1..5]@ gives @toU [0.0,1.0,3.0,6.0,10.0]@
--
-- @scanResU (+) 0.0 $ toU [1..5]@ gives @toU [0.0,1.0,3.0,6.0,10.0] :*: 15.0@.
--
scanResU :: UA a => (a -> a -> a) -> a -> UArr a -> UArr a :*: a
{-# INLINE_U scanResU #-}
scanResU f z = unstreamScan f z . streamU

unstreamScan :: UA a => (a -> a -> a) -> a -> Stream a -> UArr a :*: a
{-# INLINE_STREAM unstreamScan #-}
unstreamScan f z st@(Stream _ _ n)
  = newDynResU n (\marr -> unstreamScanM marr f z st)

unstreamScanM :: UA a => MUArr a s -> (a -> a -> a) -> a -> Stream a
                      -> ST s (Int :*: a)
{-# INLINE_U unstreamScanM #-}
unstreamScanM marr f z (Stream next s n) = fill s z 0
  where
    fill s !z !i = case next s of
                     Done       -> return (i :*: z)
                     Skip    s' -> s' `seq` fill s' z i
                     Yield x s' -> s' `seq`
                                   do
                                     writeMU marr i z
                                     fill s' (f z x) (i+1)

-- |/O(n)/. 'mapAccumLU' is an accumulating map from left to right. Unlike its
-- List counterpart, it does not return the accumulator.
--

-- FIXME: Naming inconsistent with lists.
--
mapAccumLU :: (UA a, UA b) => (c -> a -> c :*: b) -> c -> UArr a -> UArr b
{-# INLINE_U mapAccumLU #-}
mapAccumLU f z = unstreamU . mapAccumS f z . streamU

-- zipU is re-exported from UArr

-- |/O(1)/. 'zip3U' takes three arrays and returns an array of triples.
--
zip3U :: (UA e1, UA e2, UA e3) 
      => UArr e1 -> UArr e2 -> UArr e3 -> UArr (e1 :*: e2 :*: e3)
{-# INLINE_U zip3U #-}
zip3U a1 a2 a3 = (a1 `zipU` a2) `zipU` a3

-- |/O(n)/. 'zipWithU' applies a function to corresponding elements of two 
-- arrays, yielding an array containing the results.
--
zipWithU :: (UA a, UA b, UA c) 
         => (a -> b -> c) -> UArr a -> UArr b -> UArr c
{-# INLINE_U zipWithU #-}
zipWithU f a1 a2 = unstreamU (zipWithS f (streamU a1) (streamU a2))

-- |/O(n)/. 'zipWith3U' applies a function to corresponding elements of three
-- arrays, yielding an array with the results.
--
zipWith3U :: (UA a, UA b, UA c, UA d) 
          => (a -> b -> c -> d) -> UArr a -> UArr b -> UArr c -> UArr d
{-# INLINE_U zipWith3U #-}
zipWith3U f a1 a2 a3 = unstreamU (zipWith3S f (streamU a1)
                                              (streamU a2)
                                              (streamU a3))

-- unzipU is re-exported from UArr

-- |/O(1)/. 'unzip3U' unpairs an array of strict triples into three arrays.
unzip3U :: (UA e1, UA e2, UA e3) 
        => UArr (e1 :*: e2 :*: e3) -> (UArr e1 :*: UArr e2 :*: UArr e3)
{-# INLINE_U unzip3U #-}
unzip3U a = let (a12 :*: a3) = unzipU a
                (a1  :*: a2) = unzipU a12
            in
            (a1 :*: a2 :*: a3)

-- fstU and sndU reexported from UArr

-- |/O(n)/. @'combineU' f a1 a2@ yields an array by picking elements from @a1@
-- if @f@ is @True@ at the given position, and picking elements from @a2@ 
-- otherwise. For example:
--
-- @combineU (toU [True,True,False,True,False,False]) (toU [1..3]) (toU [4..6])@
--
-- yields @toU [1.0,2.0,4.0,3.0,5.0,6.0]@.
combineU :: UA a
         => UArr Bool -> UArr a -> UArr a -> UArr a
{-# INLINE_U combineU #-}
combineU f a1 a2 = checkEq (here "combineU") 
     ("flag length not equal to sum of arg length")
     (lengthU f) (lengthU a1 + lengthU a2) $ 
--  trace ("combineU:\n\t"  ++ show (lengthU f)  ++ "\n\t" ++ show (lengthU a1) ++ "\n\t" ++ show (lengthU a2) ++ "\n")
  unstreamU (combineS (streamU f) (streamU a1) (streamU a2))

------------------------------------------------------------------------

-- sliceU reexported from UArr

-- {-# INLINE_U extractU #-}
-- extractU :: UA a => UArr a -> Int -> Int -> UArr a
-- extractU arr i n = newU n $ \marr -> copyMU marr 0 (sliceU arr i n)

-- |/O(n)/. 'tailU' yields the given array without its initial element.
tailU :: UA e => UArr e -> UArr e
tailU = unstreamU . tailS . streamU
{-# INLINE_U tailU #-}

-- |/O(n)/. 'dropU' yields the suffix obtained by dropping the given number
-- of elements from an array.
dropU :: UA e=> Int -> UArr e -> UArr e
dropU n = unstreamU . dropS n . streamU
{-# INLINE dropU #-}

-- |/O(n)/. 'takeU' yields the prefix of the given length of an array.
takeU :: UA e=> Int -> UArr e -> UArr e
takeU n = unstreamU . takeS n . streamU
{-# INLINE takeU #-}

-- |/O(n)/. 'splitAtU' splits an array into two subarrays at the given index.
--
splitAtU :: UA e => Int -> UArr e -> (UArr e, UArr e)
splitAtU n a = (takeU n a, dropU n a)
{-# INLINE splitAtU #-}

{-

-- |Extract a prefix of an array
--
takeU :: UA e=> Int -> UArr e -> UArr e
{-# INLINE_U takeU #-}
takeU n a = extractU a 0 n

-- |Extract a suffix of an array
--
dropU :: UA e => Int -> UArr e -> UArr e
{-# INLINE_U dropU #-}
dropU n a = let len = lengthU a
            in
            extractU a n (len - n)

-- |Split an array into two halves at the given index
--
splitAtU :: UA e => Int -> UArr e -> (UArr e, UArr e)
{-# INLINE_U splitAtU #-}
splitAtU n a = (takeU n a, dropU n a)

-}

------------------------------------------------------------------------

-- |/O(n)/. 'takeWhileU', applied to a predicate @p@ and a UArr @xs@,
-- returns the longest prefix (possibly empty) of @xs@ of elements that satisfy @p@.
--
takeWhileU :: UA e => (e -> Bool) -> UArr e -> UArr e
takeWhileU f = unstreamU . takeWhileS f . streamU
{-# INLINE_U takeWhileU #-}

-- |/O(n)/. 'dropWhileU' @p xs@ returns the suffix remaining after 'takeWhileU' @p xs@.
dropWhileU :: UA e => (e -> Bool) -> UArr e -> UArr e
dropWhileU f = unstreamU . dropWhileS f . streamU
{-# INLINE_U dropWhileU #-}

------------------------------------------------------------------------
-- ** Searching with a predicate

-- |/O(n)/, /fusion/. The 'findU' function takes a predicate and an array
-- and returns the first element in the list matching the predicate, or
-- 'Nothing' if there is no such element.
findU :: UA a => (a -> Bool) -> UArr a -> Maybe a
{-# INLINE_U findU #-}
findU p = findS p . streamU

-- |/O(n)/, /fusion/. The 'findIndexU' function takes a predicate and an array and returns
-- the index of the first element in the array satisfying the predicate,
-- or 'Nothing' if there is no such element.
--
findIndexU :: UA a => (a -> Bool) -> UArr a -> Maybe Int
{-# INLINE_U findIndexU #-}
findIndexU p = findIndexS p . streamU

-- |/O(n)/, /fusion/. 'lookupU' @key assocs@ looks up a key in an array
-- of pairs treated as an association table.
--
lookupU :: (Eq a, UA a, UA b) => a -> UArr (a :*: b) -> Maybe b
{-# INLINE_U lookupU #-}
lookupU p = lookupS p . streamU

------------------------------------------------------------------------

-- |/O(n)/. @'unfoldU' n f z@ builds an array of size @n@ from a seed value
-- @z@ by iteratively applying @f@, stopping when @n@ elements have been 
-- generated or @f@ returns 'NothingS'.
--
unfoldU :: UA a => Int -> (b -> MaybeS (a :*: b)) -> b -> UArr a
{-# INLINE_U unfoldU #-}
unfoldU n f z = unstreamU (unfoldS n f z)