{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{- |
Functions on lazy storable vectors that are implemented using LLVM.
-}
module Synthesizer.LLVM.Storable.Signal (
   unpackStrict, unpack,
   unpackStereoStrict, unpackStereo,
   makeUnpackGenericStrict, makeUnpackGeneric,
   makeReversePackedStrict, makeReversePacked,
   continue, continuePacked, continuePackedGeneric,
   -- should be moved to a private module
   fillBuffer, makeMixer, addToBuffer,
   makeArranger, arrange,
   ) where

import qualified Synthesizer.LLVM.Parameterized.Signal as SigP
import qualified Synthesizer.LLVM.Parameterized.SignalPacked as SigPS

import qualified Synthesizer.LLVM.Execution as Exec

import qualified Synthesizer.LLVM.Frame.SerialVector as Serial
import qualified Synthesizer.LLVM.Frame.Stereo as Stereo
import qualified Synthesizer.LLVM.Frame.StereoInterleaved as StereoVector

import qualified Data.StorableVector.Lazy as SVL
import qualified Data.StorableVector as SV
import qualified Data.StorableVector.Base as SVB

import qualified Data.EventList.Relative.TimeBody  as EventList
import qualified Data.EventList.Relative.TimeMixed as EventListTM
import qualified Data.EventList.Absolute.TimeBody  as AbsEventList
import qualified Number.NonNegative as NonNeg

import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Arithmetic as A
import LLVM.Extra.Arithmetic (advanceArrayElementPtr, )
import LLVM.Extra.Control (arrayLoop, )
import LLVM.Extra.Class (MakeValueTuple, ValueTuple, )
import LLVM.Core
   (Linkage(ExternalLinkage), createNamedFunction, ret,
    IsPrimitive, getElementPtr, )
import qualified Types.Data.Num as TypeNum

import qualified Control.Category as Cat

import qualified Data.List.HT as ListHT
import Data.Word (Word32, )
import Data.Int (Int32, )
import Foreign.Ptr (Ptr, )
import Foreign.ForeignPtr (castForeignPtr, )
import Foreign.Storable (Storable, )
import Foreign.Marshal.Array (advancePtr, )

import qualified System.Unsafe as Unsafe

import NumericPrelude.Numeric
import NumericPrelude.Base


{- |
This function needs only constant time
in contrast to 'Synthesizer.LLVM.Parameterized.SignalPacked.unpack'.

We cannot provide a 'pack' function
since the array size may not line up.
It would also need copying since the source data may not be aligned properly.
-}
unpackStrict ::
   (Storable a, IsPrimitive a, TypeNum.PositiveT n) =>
   SV.Vector (Serial.Plain n a) -> SV.Vector a
unpackStrict v =
   let getDim :: (TypeNum.PositiveT n) => SV.Vector (Serial.Plain n a) -> n -> Int
       getDim _ = TypeNum.fromIntegerT
       d = getDim v undefined
       (fptr,s,l) = SVB.toForeignPtr v
   in  SVB.SV (castForeignPtr fptr) (s*d) (l*d)

unpack ::
   (Storable a, IsPrimitive a, TypeNum.PositiveT n) =>
   SVL.Vector (Serial.Plain n a) -> SVL.Vector a
unpack =
   SVL.fromChunks . map unpackStrict . SVL.chunks


unpackStereoStrict ::
   (Storable a, IsPrimitive a, TypeNum.PositiveT n) =>
   SV.Vector (StereoVector.T n a) -> SV.Vector (Stereo.T a)
unpackStereoStrict v =
   let getDim :: (TypeNum.PositiveT n) => SV.Vector (StereoVector.T n a) -> n -> Int
       getDim _ = TypeNum.fromIntegerT
       d = getDim v undefined
       (fptr,s,l) = SVB.toForeignPtr v
   in  SVB.SV (castForeignPtr fptr) (s*d) (l*d)

unpackStereo ::
   (Storable a, IsPrimitive a, TypeNum.PositiveT n) =>
   SVL.Vector (StereoVector.T n a) -> SVL.Vector (Stereo.T a)
unpackStereo =
   SVL.fromChunks . map unpackStereoStrict . SVL.chunks

{- |
This is similar to 'unpackStrict' but performs rearrangement of data.
This is for instance necessary for stereo signals
where the data layout of packed and unpacked data is different,
thus simple casting of the data is not possible.
However, for vectorized Stereo data the StereoInterleaved type
still uses vector operations for interleaving and thus is more efficient.
-}
makeUnpackGenericStrict ::
   (Serial.C vv, n ~ Serial.Size vv, va ~ Serial.Element vv,
    Memory.C (Serial.ReadIt vv),
    Storable a, MakeValueTuple a, ValueTuple a ~ va, Memory.C va,
    Storable v, MakeValueTuple v, ValueTuple v ~ vv, Memory.C vv) =>
   IO (SV.Vector v -> SV.Vector a)
makeUnpackGenericStrict =
   let vectorSize ::
          (Serial.C vl, n ~ Serial.Size vl, al ~ Serial.Element vl, Storable v, MakeValueTuple v, ValueTuple v ~ vl) =>
          SV.Vector v -> n
       vectorSize _ = undefined
   in  fmap (\f v -> f (TypeNum.fromIntegerT (vectorSize v) * SV.length v) v) $
       SigP.run (SigPS.unpack $ SigP.fromStorableVector Cat.id)

makeUnpackGeneric ::
   (Serial.C vv, n ~ Serial.Size vv, va ~ Serial.Element vv,
    Memory.C (Serial.ReadIt vv),
    Storable a, MakeValueTuple a, ValueTuple a ~ va, Memory.C va,
    Storable v, MakeValueTuple v, ValueTuple v ~ vv, Memory.C vv) =>
   IO (SVL.Vector v -> SVL.Vector a)
makeUnpackGeneric =
   fmap (\f -> SVL.fromChunks . map f . SVL.chunks) $
   makeUnpackGenericStrict


makeReverser ::
   (Storable a, Serial.C value,
    MakeValueTuple a, ValueTuple a ~ value, Memory.C value) =>
   value -> IO (Word32 -> Ptr a -> Ptr a -> IO ())
--   (Memory.C a struct, Serial.C a) =>
--   IO (Word32 -> Ptr struct -> Ptr struct -> IO ())
makeReverser dummy =
   fmap (\f len srcPtr dstPtr ->
      f len (Memory.castStorablePtr srcPtr) (Memory.castStorablePtr dstPtr)) $
   fmap derefMixPtr $
   Exec.compileModule $
   createNamedFunction ExternalLinkage "reverse" $ \ size ptrA ptrB -> do
      ptrAEnd <- getElementPtr ptrA (size, ())
      _ <- arrayLoop size ptrB ptrAEnd $ \ ptrBi ptrAj0 -> do
         ptrAj1 <- getElementPtr ptrAj0 (-1 :: Int32, ())
         flip Memory.store ptrBi
            =<< Serial.reverse
            . flip asTypeOf dummy
            =<< Memory.load ptrAj1
         return ptrAj1
      ret ()

makeReversePackedStrict ::
   (Storable v, Serial.C vv, n ~ Serial.Size vv, va ~ Serial.Element vv,
    MakeValueTuple v, ValueTuple v ~ vv, Memory.C vv) =>
   IO (SV.Vector v -> SV.Vector v)
makeReversePackedStrict = do
   rev <- makeReverser undefined
   return $ \v ->
      Unsafe.performIO $
      SVB.withStartPtr v $ \ptrA len ->
      SVB.create len $ \ptrB ->
      rev (fromIntegral len) ptrA ptrB

makeReversePacked ::
   (Storable v, Serial.C vv, n ~ Serial.Size vv, va ~ Serial.Element vv,
    MakeValueTuple v, ValueTuple v ~ vv, Memory.C vv) =>
   IO (SVL.Vector v -> SVL.Vector v)
makeReversePacked =
   fmap (\f -> SVL.fromChunks . reverse . map f . SVL.chunks) $
   makeReversePackedStrict


{- |
Append two signals where the second signal
gets the last value of the first signal as parameter.
If the first signal is empty
then there is no parameter for the second signal
and thus we simply return an empty signal in that case.
-}
continue ::
   (Storable a) =>
   SVL.Vector a -> (a -> SVL.Vector a) -> SVL.Vector a
continue x y =
   SVL.fromChunks $
   withLast SV.empty
      (SVL.chunks x)
      (SV.switchR [] $ \_ -> SVL.chunks . y)

_continueNeglectLast ::
   (Storable a) =>
   SVL.Vector a -> (a -> SVL.Vector a) -> SVL.Vector a
_continueNeglectLast x y =
   SVL.switchR SVL.empty
      (\body l -> SVL.append body (y l)) x

continuePacked ::
   (TypeNum.PositiveT n, Storable a, IsPrimitive a) =>
   SVL.Vector (Serial.Plain n a) ->
   (a -> SVL.Vector (Serial.Plain n a)) ->
   SVL.Vector (Serial.Plain n a)
continuePacked x y =
   SVL.fromChunks $
   withLast SV.empty
      (SVL.chunks x)
      (SV.switchR [] (\_ -> SVL.chunks . y) .
       unpackStrict)

{-
This function reduces the last chunk to size one, repacks that
and takes the last value.
It would be certainly more efficient to use
a single @Memory.load@, @extractelement@ and @store@
instead of a loop of count 1.
However, this implementation is the simplest one, so far.
-}
{- |
Use this like

> do unpackGeneric <- makeUnpackGenericStrict
>    return (continuePackedGeneric unpackGeneric x y)
-}
continuePackedGeneric ::
{-
   (Storable v, Serial.C v, n ~ Serial.Size v, a ~ Serial.Element v,
    MakeValueTuple v, ValueTuple v ~ vv, Memory.C vv) =>
-}
   (Storable v, Storable a) =>
   (SV.Vector v -> SV.Vector a) ->
   SVL.Vector v -> (a -> SVL.Vector v) -> SVL.Vector v
continuePackedGeneric unpackGeneric x y =
   SVL.fromChunks $
   withLast SV.empty
      (SVL.chunks x)
      (\lastChunk ->
         SV.switchR [] (\_ -> SVL.chunks . y) $ unpackGeneric $
         SV.drop (SV.length lastChunk - 1) $ lastChunk)


-- candidate for utility-ht
withLast :: a -> [a] -> (a -> [a]) -> [a]
withLast deflt x y =
   foldr
      (\a cont _ -> a : cont a)
      y x deflt

{-
This version is too strict, since it looks one element ahead.
-}
_withLast :: [a] -> (a -> [a]) -> [a]
_withLast x y =
   ListHT.switchR []
      (\body end -> body ++ end : y end)
      x



foreign import ccall safe "dynamic" derefFillPtr ::
   Exec.Importer (Word32 -> Ptr a -> IO ())

{- |
'fillBuffer' is not only more general than filling with zeros,
it also simplifies type inference.
-}
fillBuffer ::
   (MakeValueTuple a, ValueTuple a ~ value, Memory.C value) =>
   value -> IO (Word32 -> Ptr a -> IO ())
fillBuffer x =
   fmap (\f len ptr -> f len (Memory.castStorablePtr ptr)) $
   fmap derefFillPtr $
   Exec.compileModule $
   createNamedFunction ExternalLinkage "constantfill" $ \ size ptr -> do
      arrayLoop size ptr () $ \ ptri () -> do
         Memory.store x ptri
         return ()
      ret ()


foreign import ccall safe "dynamic" derefMixPtr ::
   Exec.Importer (Word32 -> Ptr a -> Ptr a -> IO ())

makeMixer ::
   (Storable a, A.Additive value,
    MakeValueTuple a, ValueTuple a ~ value, Memory.C value) =>
   value -> IO (Word32 -> Ptr a -> Ptr a -> IO ())
makeMixer dummy =
   fmap (\f len srcPtr dstPtr ->
      f len (Memory.castStorablePtr srcPtr) (Memory.castStorablePtr dstPtr)) $
   fmap derefMixPtr $
   Exec.compileModule $
   createNamedFunction ExternalLinkage "mix" $ \ size srcPtr dstPtr -> do
      _ <- arrayLoop size srcPtr dstPtr $ \ srcPtri dstPtri -> do
         y <- Memory.load srcPtri
         Memory.modify (A.add (y `asTypeOf` dummy)) dstPtri
         advanceArrayElementPtr dstPtri
      ret ()


addToBuffer ::
   (Storable a) =>
   (Word32 -> Ptr a -> Ptr a -> IO ()) ->
   Int -> Ptr a -> Int -> SVL.Vector a -> IO (Int, SVL.Vector a)
addToBuffer addChunkToBuffer len v start xs =
   let (now,future) = SVL.splitAt (len - start) xs
       go i [] = return i
       go i (c:cs) =
          SVB.withStartPtr c (\ptr l ->
             addChunkToBuffer (fromIntegral l) ptr (advancePtr v i)) >>
          go (i + SV.length c) cs
   in  fmap (flip (,) future) . go start . SVL.chunks $ now


{-
Same algorithm as in Synthesizer.Storable.Cut.arrangeEquidist
-}
makeArranger ::
   (Storable a, A.Additive value,
    MakeValueTuple a, ValueTuple a ~ value, Memory.C value) =>
   IO (SVL.ChunkSize ->
       EventList.T NonNeg.Int (SVL.Vector a) ->
       SVL.Vector a)
makeArranger = do
   mixer <- makeMixer undefined
   fill <- fillBuffer A.zero
   return $ \ (SVL.ChunkSize sz) ->
      let sznn = NonNeg.fromNumberMsg "arrange" sz
          go acc evs =
             let (now,future) = EventListTM.splitAtTime sznn evs
                 xs =
                    AbsEventList.toPairList $
                    EventList.toAbsoluteEventList 0 $
                    EventListTM.switchTimeR const now
                 (chunk,newAcc) =
                    Unsafe.performIO $
                    SVB.createAndTrim' sz $ \ptr -> do
                       fill (fromIntegral sz) ptr
                       newAcc0 <- flip mapM acc $ addToBuffer mixer sz ptr 0
                       newAcc1 <- flip mapM xs $ \(i,s) ->
                          addToBuffer mixer sz ptr (NonNeg.toNumber i) s
                       let (ends, suffixes) = unzip $ newAcc0++newAcc1
                           {- if there are more events to come,
                              we must pad with zeros -}
                           len =
                              if EventList.null future
                                then foldl max 0 ends
                                else sz
                       return (0, len,
                               filter (not . SVL.null) suffixes)
             in  if SV.null chunk
                   then []
                   else chunk : go newAcc future
      in  SVL.fromChunks . go []

{- |
This is unsafe since it relies on the prior initialization of the LLVM JIT.
Better use 'makeArranger'.
-}
{-# DEPRECATED arrange "better use makeArranger" #-}
arrange ::
   (Storable a, A.Additive value,
    MakeValueTuple a, ValueTuple a ~ value, Memory.C value) =>
      SVL.ChunkSize
   -> EventList.T NonNeg.Int (SVL.Vector a)
         {-^ A list of pairs: (relative start time, signal part),
             The start time is relative to the start time
             of the previous event. -}
   -> SVL.Vector a
         {-^ The mixed signal. -}
arrange =
   Unsafe.performIO makeArranger