{-# 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.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.DSL.Execution as Exec

import qualified LLVM.Extra.Storable as Storable
import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Tuple as Tuple

import qualified LLVM.Core as LLVM
import LLVM.Core (IsPrimitive, ret)

import qualified Type.Data.Num.Decimal as TypeNum

import qualified Control.Category as Cat

import qualified Data.List.HT as ListHT
import Data.Word (Word)

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.C a, IsPrimitive a, TypeNum.Positive n) =>
   SV.Vector (Serial.Plain n a) -> SV.Vector a
unpackStrict v =
   let getDim ::
          (TypeNum.Positive n) =>
          SV.Vector (Serial.Plain n a) -> TypeNum.Singleton n -> Int
       getDim _ = TypeNum.integralFromSingleton
       d = getDim v TypeNum.singleton
       (fptr,s,l) = SVB.toForeignPtr v
   in  SVB.SV (castForeignPtr fptr) (s*d) (l*d)

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


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

unpackStereo ::
   (Storable.C a, IsPrimitive a, TypeNum.Positive 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.C a, Tuple.ValueOf a ~ va,
    Storable.C v, Tuple.ValueOf v ~ vv) =>
   IO (SV.Vector v -> SV.Vector a)
makeUnpackGenericStrict =
   let vectorSize ::
          (Serial.C vl, n ~ Serial.Size vl, al ~ Serial.Element vl,
           Storable.C v, Tuple.ValueOf v ~ vl) =>
          SV.Vector v -> TypeNum.Singleton n
       vectorSize _ = TypeNum.singleton
   in  fmap (\f v -> f (TypeNum.integralFromSingleton (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.C a, Tuple.ValueOf a ~ va,
    Storable.C v, Tuple.ValueOf v ~ vv) =>
   IO (SVL.Vector v -> SVL.Vector a)
makeUnpackGeneric =
   fmap (\f -> SVL.fromChunks . map f . SVL.chunks) $
   makeUnpackGenericStrict


makeReverser ::
   (Storable.C a, Tuple.ValueOf a ~ value) =>
   (value -> LLVM.CodeGenFunction () value) ->
   IO (Word -> Ptr a -> Ptr a -> IO ())
--   (Memory.C a struct, Serial.C a) =>
--   IO (Word -> Ptr struct -> Ptr struct -> IO ())
makeReverser rev =
   Exec.compile "reverse" $
   Exec.createFunction derefMixPtr "reverse" $ \ size ptrA ptrB -> do
      sizeInt <- LLVM.bitcast size
      ptrAEnd <- Storable.advancePtr sizeInt ptrA
      _ <- Storable.arrayLoop size ptrB ptrAEnd $ \ ptrBi ptrAj0 -> do
         ptrAj1 <- Storable.decrementPtr ptrAj0
         flip Storable.store ptrBi
            =<< rev
            =<< Storable.load ptrAj1
         return ptrAj1
      ret ()

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

makeReversePacked ::
   (Serial.C vv, n ~ Serial.Size vv, va ~ Serial.Element vv,
    Storable.C v, Tuple.ValueOf v ~ 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.Positive n, Storable.C 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,
    Tuple.Value v, Tuple.ValueOf 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 (Word -> Ptr a -> IO ())

{- |
'fillBuffer' is not only more general than filling with zeros,
it also simplifies type inference.
-}
fillBuffer ::
   (Storable.C a, Tuple.ValueOf a ~ value) =>
   value -> IO (Word -> Ptr a -> IO ())
fillBuffer x =
   Exec.compile "constant" $
   Exec.createFunction derefFillPtr "constantfill" $ \ size ptr -> do
      Storable.arrayLoop size ptr () $ \ ptri () -> Storable.store x ptri
      ret ()


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

makeMixer ::
   (Storable.C a, Tuple.ValueOf a ~ value) =>
   (value -> value -> LLVM.CodeGenFunction () value) ->
   IO (Word -> Ptr a -> Ptr a -> IO ())
makeMixer add =
   Exec.compile "mixer" $
   Exec.createFunction derefMixPtr "mix" $ \ size srcPtr dstPtr -> do
      _ <-
         Storable.arrayLoop2 size srcPtr dstPtr () $
            \srcPtri dstPtri () -> do
         y <- Storable.load srcPtri
         Storable.modify (add y) dstPtri
      ret ()


addToBuffer ::
   (Storable a) =>
   (Word -> 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.C a, Tuple.ValueOf a ~ value, A.Additive value) =>
   IO (SVL.ChunkSize ->
       EventList.T NonNeg.Int (SVL.Vector a) ->
       SVL.Vector a)
makeArranger = do
   mixer <- makeMixer A.add
   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.C a, Tuple.ValueOf a ~ value, A.Additive 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