{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE MagicHash            #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE UnboxedTuples        #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

-- |
-- Module      :  Internal.Vector
-- Copyright   :  (c) Alberto Ruiz 2007-15
-- License     :  BSD3
-- Maintainer  :  Alberto Ruiz
-- Stability   :  provisional
--

module Internal.Vector(
    I,Z,R,C,
    fi,ti,
    Vector, fromList, unsafeToForeignPtr, unsafeFromForeignPtr, unsafeWith,
    createVector, avec, inlinePerformIO,
    toList, dim, (@>), at', (|>),
    vjoin, subVector, takesV, idxs,
    buildVector,
    asReal, asComplex,
    toByteString,fromByteString,
    zipVector, unzipVector, zipVectorWith, unzipVectorWith,
    foldVector, foldVectorG, foldVectorWithIndex, foldLoop,
    mapVector, mapVectorM, mapVectorM_,
    mapVectorWithIndex, mapVectorWithIndexM, mapVectorWithIndexM_
) where

import           Data.Complex
import           Data.Int                      (Int64)
import           Data.Vector.Storable          (Vector, fromList, unsafeFromForeignPtr,
                                                unsafeToForeignPtr, unsafeWith)
import qualified Data.Vector.Storable          as Vector
import           Foreign.C.Types               (CInt)
import           Foreign.ForeignPtr
import           Foreign.Marshal.Array
import           Foreign.Ptr
import           Foreign.Storable
import           GHC.Base                      (IO (IO), realWorld#, when)
import           GHC.ForeignPtr                (mallocPlainForeignPtrBytes)
import           System.IO.Unsafe              (unsafePerformIO)

import           Control.Monad                 (replicateM)
import           Data.Binary
import           Data.Binary.Put
import qualified Data.ByteString.Internal      as BS
import           Data.Vector.Storable.Internal (updPtr)

type I = CInt
type Z = Int64
type R = Float
type C = Complex Float


-- | specialized fromIntegral
fi :: Int -> CInt
fi = fromIntegral

-- | specialized fromIntegral
ti :: CInt -> Int
ti = fromIntegral


-- | Number of elements
dim :: (Storable t) => Vector t -> Int
dim = Vector.length
{-# INLINE dim #-}


-- C-Haskell vector adapter
{-# INLINE avec #-}
avec :: Storable a => Vector a -> (f -> IO r) -> ((CInt -> Ptr a -> f) -> IO r)
avec v f g = unsafeWith v $ \ptr -> f (g (fromIntegral (Vector.length v)) ptr)

-- allocates memory for a new vector
createVector :: Storable a => Int -> IO (Vector a)
createVector n = do
    when (n < 0) $ error ("trying to createVector of negative dim: "++show n)
    fp <- doMalloc undefined
    return $ unsafeFromForeignPtr fp 0 n
  where
    --
    -- Use the much cheaper Haskell heap allocated storage
    -- for foreign pointer space we control
    --
    doMalloc :: Storable b => b -> IO (ForeignPtr b)
    doMalloc dummy = do
        mallocPlainForeignPtrBytes (n * sizeOf dummy)

{- | creates a Vector from a list:

@> fromList [2,3,5,7]
4 |> [2.0,3.0,5.0,7.0]@

-}

safeRead :: Storable a => Vector a -> (Ptr a -> IO c) -> c
safeRead v = inlinePerformIO . unsafeWith v
{-# INLINE safeRead #-}

inlinePerformIO :: IO a -> a
inlinePerformIO (IO m) = case m realWorld# of (# _, r #) -> r
{-# INLINE inlinePerformIO #-}

{- extracts the Vector elements to a list

>>> toList (linspace 5 (1,10))
[1.0,3.25,5.5,7.75,10.0]

-}
toList :: Storable a => Vector a -> [a]
toList v = safeRead v $ peekArray (dim v)

{- | Create a vector from a list of elements and explicit dimension. The input
     list is truncated if it is too long, so it may safely
     be used, for instance, with infinite lists.

>>> 5 |> [1..]
[1.0,2.0,3.0,4.0,5.0]
it :: (Enum a, Num a, Foreign.Storable.Storable a) => Vector a

-}
(|>) :: (Storable a) => Int -> [a] -> Vector a
infixl 9 |>
n |> l
    | length l' == n = fromList l'
    | otherwise      = error "list too short for |>"
  where
    l' = take n l


-- | Create a vector of indexes, useful for matrix extraction using '(??)'
idxs :: [Int] -> Vector I
idxs js = fromList (map fromIntegral js) :: Vector I

{- | takes a number of consecutive elements from a Vector

>>> subVector 2 3 (fromList [1..10])
[3.0,4.0,5.0]
it :: (Enum t, Num t, Foreign.Storable.Storable t) => Vector t

-}
subVector :: Storable t => Int       -- ^ index of the starting element
                        -> Int       -- ^ number of elements to extract
                        -> Vector t  -- ^ source
                        -> Vector t  -- ^ result
subVector = Vector.slice
{-# INLINE subVector #-}


{- | Reads a vector position:

>>> fromList [0..9] @> 7
7.0

-}
(@>) :: Storable t => Vector t -> Int -> t
infixl 9 @>
v @> n
    | n >= 0 && n < dim v = at' v n
    | otherwise = error "vector index out of range"
{-# INLINE (@>) #-}

-- | access to Vector elements without range checking
at' :: Storable a => Vector a -> Int -> a
at' v n = safeRead v $ flip peekElemOff n
{-# INLINE at' #-}

{- | concatenate a list of vectors

>>> vjoin [fromList [1..5::Float], konst 1 3]
[1.0,2.0,3.0,4.0,5.0,1.0,1.0,1.0]
it :: Vector Float

-}
vjoin :: Storable t => [Vector t] -> Vector t
vjoin [] = fromList []
vjoin [v] = v
vjoin as = unsafePerformIO $ do
    let tot = sum (map dim as)
    r <- createVector tot
    unsafeWith r $ \ptr ->
        joiner as tot ptr
    return r
  where joiner [] _ _ = return ()
        joiner (v:cs) _ p = do
            let n = dim v
            unsafeWith v $ \pb -> copyArray p pb n
            joiner cs 0 (advancePtr p n)


{- | Extract consecutive subvectors of the given sizes.

>>> takesV [3,4] (linspace 10 (1,10::Float))
[[1.0,2.0,3.0],[4.0,5.0,6.0,7.0]]
it :: [Vector Float]

-}
takesV :: Storable t => [Int] -> Vector t -> [Vector t]
takesV ms w | sum ms > dim w = error $ "takesV " ++ show ms ++ " on dim = " ++ (show $ dim w)
            | otherwise = go ms w
    where go [] _ = []
          go (n:ns) v = subVector 0 n v
                      : go ns (subVector n (dim v - n) v)

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

-- | transforms a complex vector into a real vector with alternating real and imaginary parts
asReal :: (RealFloat a, Storable a) => Vector (Complex a) -> Vector a
asReal v = unsafeFromForeignPtr (castForeignPtr fp) (2*i) (2*n)
    where (fp,i,n) = unsafeToForeignPtr v

-- | transforms a real vector into a complex vector with alternating real and imaginary parts
asComplex :: (RealFloat a, Storable a) => Vector a -> Vector (Complex a)
asComplex v = unsafeFromForeignPtr (castForeignPtr fp) (i `div` 2) (n `div` 2)
    where (fp,i,n) = unsafeToForeignPtr v

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


-- | map on Vectors
mapVector :: (Storable a, Storable b) => (a-> b) -> Vector a -> Vector b
mapVector f v = unsafePerformIO $ do
    w <- createVector (dim v)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) = return ()
                go !k = do x <- peekElemOff p k
                           pokeElemOff      q k (f x)
                           go (k-1)
            go (dim v -1)
    return w
{-# INLINE mapVector #-}

-- | zipWith for Vectors
zipVectorWith :: (Storable a, Storable b, Storable c) => (a-> b -> c) -> Vector a -> Vector b -> Vector c
zipVectorWith f u v = unsafePerformIO $ do
    let n = min (dim u) (dim v)
    w <- createVector n
    unsafeWith u $ \pu ->
        unsafeWith v $ \pv ->
            unsafeWith w $ \pw -> do
                let go (-1) = return ()
                    go !k = do x <- peekElemOff pu k
                               y <- peekElemOff pv k
                               pokeElemOff      pw k (f x y)
                               go (k-1)
                go (n -1)
    return w
{-# INLINE zipVectorWith #-}

-- | unzipWith for Vectors
unzipVectorWith :: (Storable (a,b), Storable c, Storable d)
                   => ((a,b) -> (c,d)) -> Vector (a,b) -> (Vector c,Vector d)
unzipVectorWith f u = unsafePerformIO $ do
      let n = dim u
      v <- createVector n
      w <- createVector n
      unsafeWith u $ \pu ->
          unsafeWith v $ \pv ->
              unsafeWith w $ \pw -> do
                  let go (-1) = return ()
                      go !k   = do z <- peekElemOff pu k
                                   let (x,y) = f z
                                   pokeElemOff      pv k x
                                   pokeElemOff      pw k y
                                   go (k-1)
                  go (n-1)
      return (v,w)
{-# INLINE unzipVectorWith #-}

foldVector :: Storable a => (a -> b -> b) -> b -> Vector a -> b
foldVector f x v = unsafePerformIO $
    unsafeWith v $ \p -> do
        let go (-1) s = return s
            go !k !s = do y <- peekElemOff p k
                          go (k-1::Int) (f y s)
        go (dim v -1) x
{-# INLINE foldVector #-}

-- the zero-indexed index is passed to the folding function
foldVectorWithIndex :: Storable a => (Int -> a -> b -> b) -> b -> Vector a -> b
foldVectorWithIndex f x v = unsafePerformIO $
    unsafeWith v $ \p -> do
        let go (-1) s = return s
            go !k !s = do y <- peekElemOff p k
                          go (k-1::Int) (f k y s)
        go (dim v -1) x
{-# INLINE foldVectorWithIndex #-}

foldLoop :: (Int -> t -> t) -> t -> Int -> t
foldLoop f s0 d = go (d - 1) s0
     where
       go 0 s   = f (0::Int) s
       go !j !s = go (j - 1) (f j s)

foldVectorG :: Storable t1 => (Int -> (Int -> t1) -> t -> t) -> t -> Vector t1 -> t
foldVectorG f s0 v = foldLoop g s0 (dim v)
    where g !k !s = f k (safeRead v . flip peekElemOff) s
          {-# INLINE g #-} -- Thanks to Ryan Ingram (http://permalink.gmane.org/gmane.comp.lang.haskell.cafe/46479)
{-# INLINE foldVectorG #-}

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

-- | monadic map over Vectors
--    the monad @m@ must be strict
mapVectorM :: (Storable a, Storable b, Monad m) => (a -> m b) -> Vector a -> m (Vector b)
mapVectorM f v = do
    w <- return $! unsafePerformIO $! createVector (dim v)
    mapVectorM' w 0 (dim v -1)
    return w
    where mapVectorM' w' !k !t
              | k == t               = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f x
                                       return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
              | otherwise            = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f x
                                       _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
                                       mapVectorM' w' (k+1) t
{-# INLINE mapVectorM #-}

-- | monadic map over Vectors
mapVectorM_ :: (Storable a, Monad m) => (a -> m ()) -> Vector a -> m ()
mapVectorM_ f v = do
    mapVectorM' 0 (dim v -1)
    where mapVectorM' !k !t
              | k == t            = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    f x
              | otherwise         = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    _ <- f x
                                    mapVectorM' (k+1) t
{-# INLINE mapVectorM_ #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
--    the monad @m@ must be strict
mapVectorWithIndexM :: (Storable a, Storable b, Monad m) => (Int -> a -> m b) -> Vector a -> m (Vector b)
mapVectorWithIndexM f v = do
    w <- return $! unsafePerformIO $! createVector (dim v)
    mapVectorM' w 0 (dim v -1)
    return w
    where mapVectorM' w' !k !t
              | k == t               = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f k x
                                       return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
              | otherwise            = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f k x
                                       _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
                                       mapVectorM' w' (k+1) t
{-# INLINE mapVectorWithIndexM #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
mapVectorWithIndexM_ :: (Storable a, Monad m) => (Int -> a -> m ()) -> Vector a -> m ()
mapVectorWithIndexM_ f v = do
    mapVectorM' 0 (dim v -1)
    where mapVectorM' !k !t
              | k == t            = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    f k x
              | otherwise         = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    _ <- f k x
                                    mapVectorM' (k+1) t
{-# INLINE mapVectorWithIndexM_ #-}


mapVectorWithIndex :: (Storable a, Storable b) => (Int -> a -> b) -> Vector a -> Vector b
--mapVectorWithIndex g = head . mapVectorWithIndexM (\a b -> [g a b])
mapVectorWithIndex f v = unsafePerformIO $ do
    w <- createVector (dim v)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) = return ()
                go !k = do x <- peekElemOff p k
                           pokeElemOff      q k (f k x)
                           go (k-1)
            go (dim v -1)
    return w
{-# INLINE mapVectorWithIndex #-}

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


-- a 64K cache, with a Float taking 13 bytes in Bytestring,
-- implies a chunk size of 5041
chunk :: Int
chunk = 5000

chunks :: Int -> [Int]
chunks d = let c = d `div` chunk
               m = d `mod` chunk
           in if m /= 0 then reverse (m:(replicate c chunk)) else (replicate c chunk)

putVector :: (Storable t, Binary t) => Vector t -> Data.Binary.Put.PutM ()
putVector v = mapM_ put $! toList v

getVector :: (Storable a, Binary a) => Int -> Get (Vector a)
getVector d = do
              xs <- replicateM d get
              return $! fromList xs

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

toByteString :: Storable t => Vector t -> BS.ByteString
toByteString v = BS.PS (castForeignPtr fp) (sz*o) (sz * dim v)
  where
    (fp,o,_n) = unsafeToForeignPtr v
    sz = sizeOf (v@>0)


fromByteString :: Storable t => BS.ByteString -> Vector t
fromByteString (BS.PS fp o n) = r
  where
    r = unsafeFromForeignPtr (castForeignPtr (updPtr (`plusPtr` o) fp)) 0 n'
    n' = n `div` sz
    sz = sizeOf (r@>0)

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

instance (Binary a, Storable a) => Binary (Vector a) where

    put v = do
            let d = dim v
            put d
            mapM_ putVector $! takesV (chunks d) v

    -- put = put . v2bs

    get = do
          d <- get
          vs <- mapM getVector $ chunks d
          return $! vjoin vs

    -- get = fmap bs2v get


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

{- | creates a Vector of the specified length using the supplied function to
     to map the index to the value at that index.

@> buildVector 4 fromIntegral
4 |> [0.0,1.0,2.0,3.0]@

-}
buildVector :: Storable a => Int -> (Int -> a) -> Vector a
buildVector len f =
    fromList $ map f [0 .. (len - 1)]


-- | zip for Vectors
zipVector :: (Storable a, Storable b, Storable (a,b)) => Vector a -> Vector b -> Vector (a,b)
zipVector = zipVectorWith (,)

-- | unzip for Vectors
unzipVector :: (Storable a, Storable b, Storable (a,b)) => Vector (a,b) -> (Vector a,Vector b)
unzipVector = unzipVectorWith id

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