{- |
Functions for 'StorableVector' that allow control of the size of individual chunks.

This is import for an application like the following:
You want to mix audio signals that are relatively shifted.
The structure of chunks of three streams may be illustrated as:

> [____] [____] [____] [____] ...
>   [____] [____] [____] [____] ...
>     [____] [____] [____] [____] ...

When we mix the streams (@zipWith3 (\x y z -> x+y+z)@)
with respect to the chunk structure of the first signal,
computing the first chunk requires full evaluation of all leading chunks of the stream.
However the last value of the third leading chunk
is much later in time than the last value of the first leading chunk.
We like to reduce these dependencies using a different chunk structure,
say

> [____] [____] [____] [____] ...
>   [__] [____] [____] [____] ...
>     [] [____] [____] [____] ...

-}
module Data.StorableVector.LazyVarying (
   Vector,
   ChunkSize,
   chunkSize,
   defaultChunkSize,

   empty,
   singleton,
   pack,
   unpack,
   packWith,
   unpackWith,
   unfoldrN,
   iterateN,
   cycle,
   replicate,
   null,
   length,
   cons,
   append,
   concat,
   map,
   reverse,
   foldl,
   foldl',
   any,
   all,
   maximum,
   minimum,
   viewL,
   viewR,
   switchL,
   switchR,
   scanl,
   mapAccumL,
   mapAccumR,
   crochetL,
   take,
   drop,
   splitAt,
   dropMarginRem,
   dropMargin,
   dropWhile,
   takeWhile,
   span,
   filter,
   zipWith,
   zipWith3,
   zipWith4,
   zipWithSize,
   zipWithSize3,
   zipWithSize4,
{-
   pad,
   cancelNullVector,
   fromChunk,
   hGetContentsAsync,
   hPut,
   readFileAsync,
   writeFile,
   appendFile,
-}
   ) where

import qualified Data.StorableVector.LazySize as LS
import qualified Data.StorableVector.Lazy as LSV
import qualified Data.StorableVector as V

import Data.StorableVector.Lazy (Vector(SV), ChunkSize(ChunkSize))

import Data.StorableVector.Lazy (
   chunkSize, defaultChunkSize,
   empty, singleton, unpack, unpackWith, cycle,
   null, cons, append, concat, map, reverse,
   foldl, foldl', any, all, maximum, minimum,
   viewL, viewR, switchL, switchR,
   scanl, mapAccumL, mapAccumR, crochetL,
   dropMarginRem, dropMargin,
   dropWhile, takeWhile, span, filter, 
   zipWith, zipWith3, zipWith4, 
   )

import qualified Data.List as List

import Data.Maybe (Maybe(Just, Nothing), )
import Data.StorableVector.Utility (viewListL, mapPair, mapFst, )

import Control.Monad (liftM2, liftM3, liftM4, guard, )

import Foreign.Storable (Storable)

{-
import Prelude hiding
   (length, (++), iterate, foldl, map, repeat, replicate, null,
    zip, zipWith, zipWith3, drop, take, splitAt, takeWhile, dropWhile, reverse)
-}
import Prelude ((.), ($), fst, flip, curry, return, fmap, not, )


type LazySize = LS.T

-- * Introducing and eliminating 'Vector's

pack :: (Storable a) => LazySize -> [a] -> Vector a
pack size =
   fst . unfoldrN size viewListL


{-# INLINE packWith #-}
packWith :: (Storable b) => LazySize -> (a -> b) -> [a] -> Vector b
packWith size f =
   fst . unfoldrN size (fmap (\(a,b) -> (f a, b)) . viewListL)


{-
{-# INLINE unfoldrNAlt #-}
unfoldrNAlt :: (Storable b) =>
      LazySize
   -> (a -> Maybe (b,a))
   -> a
   -> (Vector b, Maybe a)
unfoldrNAlt (LS.Cons size) f x =
   let go sz y =
          case sz of
             [] -> ([], y)
             (ChunkSize s : ss) ->
                maybe
                   ([], Nothing)
                   ((\(c,a1) -> mapFst (c:) $ go ss a1) .
                    V.unfoldrN s (fmap (mapSnd f)))
                   (f y)
   in  mapFst SV $ go size (Just x)
-}

{-# INLINE unfoldrN #-}
unfoldrN :: (Storable b) =>
      LazySize
   -> (a -> Maybe (b,a))
   -> a
   -> (Vector b, Maybe a)
unfoldrN (LS.Cons size) f x =
   let go sz y =
          case sz of
             [] -> ([], y)
             (ChunkSize s : ss) ->
                let m =
                       do a0 <- y
                          let p = V.unfoldrN s f a0
                          guard (not (V.null (fst p)))
                          return p
                in  case m of
                       Nothing -> ([], Nothing)
                       Just (c,a1) -> mapFst (c:) $ go ss a1
   in  mapFst SV $ go size (Just x)


{-# INLINE iterateN #-}
iterateN :: Storable a => LazySize -> (a -> a) -> a -> Vector a
iterateN size f =
   fst . unfoldrN size (\x -> Just (x, f x))

replicate :: Storable a => LazySize -> a -> Vector a
replicate size x =
   SV $ List.map (\(ChunkSize m) -> V.replicate m x) (LS.decons size)



-- * Basic interface

length :: Vector a -> LazySize
length = LS.Cons . List.map chunkLength . LSV.chunks

chunkLength :: V.Vector a -> ChunkSize
chunkLength = ChunkSize . V.length


-- * sub-vectors

{-# INLINE take #-}
take :: (Storable a) => LazySize -> Vector a -> Vector a
take _ (SV []) = empty
take (LS.Cons []) _ = empty
take n (SV (x:xs)) =
   let remain = LS.decrementLimit (chunkLength x) n
   in  SV $
       if LS.isNull remain
         then [V.take (LS.toInt n) x]
         else
           let (SV ys) = take remain $ SV xs
           in  x:ys

{-# INLINE drop #-}
drop :: (Storable a) => LazySize -> Vector a -> Vector a
drop (LS.Cons size) xs =
   List.foldl (flip (LSV.drop . LS.intFromChunkSize)) xs size

{-# INLINE splitAt #-}
splitAt :: (Storable a) => LazySize -> Vector a -> (Vector a, Vector a)
splitAt (LS.Cons []) = (,) empty
splitAt n0 =
   let recurse _ [] = ([], [])
       recurse n (x:xs) =
          let remain = LS.decrementLimit (chunkLength x) n
          in  if LS.isNull remain
                then mapPair ((:[]), (:xs)) $ V.splitAt (LS.toInt n) x
                else mapFst (x:) $ recurse remain xs
   in  mapPair (SV, SV) . recurse n0 . LSV.chunks


{-# INLINE [0] zipWithSize #-}
zipWithSize :: (Storable a, Storable b, Storable c) =>
      LazySize
   -> (a -> b -> c)
   -> Vector a
   -> Vector b
   -> Vector c
zipWithSize size f =
   curry (fst . unfoldrN size (\(xt,yt) ->
      liftM2
         (\(x,xs) (y,ys) -> (f x y, (xs,ys)))
         (viewL xt)
         (viewL yt)))

{-# INLINE zipWithSize3 #-}
zipWithSize3 ::
   (Storable a, Storable b, Storable c, Storable d) =>
   LazySize -> (a -> b -> c -> d) ->
   (Vector a -> Vector b -> Vector c -> Vector d)
zipWithSize3 size f s0 s1 s2 =
   fst $ unfoldrN size (\(xt,yt,zt) ->
      liftM3
         (\(x,xs) (y,ys) (z,zs) ->
             (f x y z, (xs,ys,zs)))
         (viewL xt)
         (viewL yt)
         (viewL zt))
      (s0,s1,s2)

{-# INLINE zipWithSize4 #-}
zipWithSize4 ::
   (Storable a, Storable b, Storable c, Storable d, Storable e) =>
   LazySize -> (a -> b -> c -> d -> e) ->
   (Vector a -> Vector b -> Vector c -> Vector d -> Vector e)
zipWithSize4 size f s0 s1 s2 s3 =
   fst $ unfoldrN size (\(xt,yt,zt,wt) ->
      liftM4
         (\(x,xs) (y,ys) (z,zs) (w,ws) ->
             (f x y z w, (xs,ys,zs,ws)))
         (viewL xt)
         (viewL yt)
         (viewL zt)
         (viewL wt))
      (s0,s1,s2,s3)

{-
{- |
Ensure a minimal length of the list by appending pad values.
-}
{-# ONLINE pad #-}
pad :: (Storable a) => ChunkSize -> a -> Int -> Vector a -> Vector a
pad size y n0 =
   let recurse n xt =
          if n<=0
            then xt
            else
              case xt of
                 [] -> chunks $ replicate size n y
                 x:xs -> x : recurse (n - V.length x) xs
   in  SV . recurse n0 . chunks

padAlt :: (Storable a) => ChunkSize -> a -> Int -> Vector a -> Vector a
padAlt size x n xs =
   append xs
      (let m = length xs
       in  if n>m
             then replicate size (n-m) x
             else empty)





-- * Helper functions for StorableVector


{-# INLINE cancelNullVector #-}
cancelNullVector :: (V.Vector a, b) -> Maybe (V.Vector a, b)
cancelNullVector y =
   toMaybe (not (V.null (fst y))) y

-- if the chunk has length zero, an empty sequence is generated
{-# INLINE fromChunk #-}
fromChunk :: (Storable a) => V.Vector a -> Vector a
fromChunk x =
   if V.null x
     then empty
     else SV [x]



{- * IO -}

{- |
Read the rest of a file lazily and
provide the reason of termination as IOError.
If IOError is EOF (check with @System.Error.isEOFError err@),
then the file was read successfully.
Only access the final IOError after you have consumed the file contents,
since finding out the terminating reason forces to read the entire file.
Make also sure you read the file completely,
because it is only closed when the file end is reached
(or an exception is encountered).

TODO:
In ByteString.Lazy the chunk size is reduced
if data is not immediately available.
Maybe we should adapt that behaviour
but when working with realtime streams
that may mean that the chunks are very small.
-}
hGetContentsAsync :: Storable a =>
   ChunkSize -> Handle -> IO (IOError, Vector a)
hGetContentsAsync (ChunkSize size) h =
   let go =
          unsafeInterleaveIO $
          flip catch (\err -> return (err,[])) $
          do v <- V.hGet h size
             if V.null v
               then hClose h >>
                    return (Exc.mkIOError Exc.eofErrorType
                      "StorableVector.Lazy.hGetContentsAsync" (Just h) Nothing, [])
               else liftM (\ ~(err,rest) -> (err, v:rest)) go
{-
          unsafeInterleaveIO $
          flip catch (\err -> return (err,[])) $
          liftM2 (\ chunk ~(err,rest) -> (err,chunk:rest))
             (V.hGet h size) go
-}
   in  fmap (mapSnd SV) go

{-
hGetContentsSync :: Storable a =>
   ChunkSize -> Handle -> IO (IOError, Vector a)
hGetContentsSync (ChunkSize size) h =
   let go =
          flip catch (\err -> return (err,[])) $
          do v <- V.hGet h size
             if V.null v
               then return (Exc.mkIOError Exc.eofErrorType
                      "StorableVector.Lazy.hGetContentsAsync" (Just h) Nothing, [])
               else liftM (\ ~(err,rest) -> (err, v:rest)) go
   in  fmap (mapSnd SV) go
-}

hPut :: Storable a => Handle -> Vector a -> IO ()
hPut h = mapM_ (V.hPut h) . chunks

{-
*Data.StorableVector.Lazy> print . mapSnd (length :: Vector Data.Int.Int16 -> Int) =<< readFileAsync (ChunkSize 1000) "dist/build/libHSstorablevector-0.1.3.a"
(dist/build/libHSstorablevector-0.1.3.a: hGetBuf: illegal operation (handle is closed),0)
-}
{- |
The file can only closed after all values are consumed.
That is you must always assert that you consume all elements of the stream,
and that no values are missed due to lazy evaluation.
This requirement makes this function useless in many applications.
-}
readFileAsync :: Storable a => ChunkSize -> FilePath -> IO (IOError, Vector a)
readFileAsync size path =
   openBinaryFile path ReadMode >>= hGetContentsAsync size

writeFile :: Storable a => FilePath -> Vector a -> IO ()
writeFile path =
   bracket (openBinaryFile path WriteMode) hClose . flip hPut

appendFile :: Storable a => FilePath -> Vector a -> IO ()
appendFile path =
   bracket (openBinaryFile path AppendMode) hClose . flip hPut
-}