{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ForeignFunctionInterface #-}
module Synthesizer.LLVM.CausalParameterized.ProcessPacked where

import Synthesizer.LLVM.CausalParameterized.Process (T(Cons), )
import qualified Synthesizer.LLVM.CausalParameterized.Process as Causal
import qualified Synthesizer.LLVM.Parameter as Param
import qualified Synthesizer.LLVM.Frame as Frame
import qualified Synthesizer.LLVM.Frame.SerialVector as Serial
import qualified Synthesizer.LLVM.Frame.Stereo as Stereo

import qualified LLVM.Extra.ScalarOrVector as SoV
import qualified LLVM.Extra.Vector as Vector
import qualified LLVM.Extra.MaybeContinuation as Maybe
import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.Class as Class
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Control as C
import LLVM.Extra.Class (MakeValueTuple, ValueTuple, Undefined, undefTuple, )

import LLVM.Core as LLVM

import qualified Types.Data.Bool as TypeBool
import qualified Types.Data.Num as TypeNum
import Types.Data.Num.Ops ((:*:), )
import Types.Data.Ord ((:<:), )

import qualified Control.Monad.Trans.Class as MT
import qualified Control.Monad.Trans.State as MS
import qualified Control.Category as Cat
import qualified Control.Arrow    as Arr
import Control.Arrow ((^<<), (<<<), )

import qualified Algebra.Additive as Additive

import Data.Tuple.HT (swap, )

import Data.Word (Word32, )
import Foreign.Storable (Storable, )

import NumericPrelude.Numeric
import NumericPrelude.Base hiding (and, iterate, map, zip, zipWith, )



{- |
Run a scalar process on packed data.
If the signal length is not divisible by the chunk size,
then the last chunk is dropped.
-}
pack ::
   (Serial.Read va, n ~ Serial.Size va, a ~ Serial.Element va,
    Serial.C    vb, n ~ Serial.Size vb, b ~ Serial.Element vb) =>
   T p a b -> T p va vb
pack (Cons next start createIOContext deleteIOContext) = Cons
   (\param a s -> do
      r <- Maybe.lift $ Serial.readStart a
      ((_,w2),(_,s2)) <-
         Maybe.fromBool $
         C.whileLoop
            (valueOf True,
             let w = undefTuple
             in  ((r,w), (valueOf $ (fromIntegral $ Serial.sizeOfIterator w :: Word32), s)))
            (\(cont,(_rw0,(i0,_s0))) ->
               A.and cont =<<
                  A.cmp CmpGT i0 (value LLVM.zero))
            (\(_,((r0,w0),(i0,s0))) -> Maybe.toBool $ do
               (ai,r1) <- Maybe.lift $ Serial.readNext r0
               (bi,s1) <- next param ai s0
               Maybe.lift $ do
                  w1 <- Serial.writeNext bi w0
                  i1 <- A.dec i0
                  return ((r1,w1),(i1,s1)))
      b <- Maybe.lift $ Serial.writeStop w2
      return (b, s2))
   start
   createIOContext
   deleteIOContext

{- |
Like 'pack' but duplicates the code for the scalar process.
That is, for vectors of size n,
the code for the scalar causal process will be written n times.
This is efficient only for simple input processes.
-}
packSmall ::
   (Serial.Read va, n ~ Serial.Size va, a ~ Serial.Element va,
    Serial.C    vb, n ~ Serial.Size vb, b ~ Serial.Element vb) =>
   T p a b -> T p va vb
packSmall (Cons next start createIOContext deleteIOContext) = Cons
   (\param a ->
      MS.runStateT $
         (MT.lift . Maybe.lift . Serial.assemble)
         =<<
         mapM (MS.StateT . next param)
         =<<
         (MT.lift $ Maybe.lift $ Serial.extractAll a))
   start
   createIOContext
   deleteIOContext


{- |
Run a packed process on scalar data.
If the signal length is not divisible by the chunk size,
then the last chunk is dropped.
In order to stay causal, we have to delay the output by @n@ samples.
-}
unpack ::
   (Serial.Zero va, n ~ Serial.Size va, a ~ Serial.Element va,
    Serial.Read vb, n ~ Serial.Size vb, b ~ Serial.Element vb,
    Memory.C (Serial.WriteIt va), Memory.C (Serial.ReadIt vb),
    Memory.C va,
    Memory.C vb) =>
   T p va vb -> T p a b
unpack (Cons next start createIOContext deleteIOContext) = Cons
   (\param ai ((w0,r0),(i0,s0)) -> do
      endOfVector <- Maybe.lift $ A.cmp CmpEQ i0 (valueOf 0)
      ((w2,r2),(i2,s2)) <-
         Maybe.fromBool $
         C.ifThen endOfVector (valueOf True, ((w0,r0),(i0,s0))) $ do
            a0 <- Serial.writeStop w0
            (cont1, (b1,s1)) <- Maybe.toBool $ next param a0 s0
            r1 <- Serial.readStart b1
            w1 <- Serial.writeStart
            return (cont1,
                      ((w1, r1),
                       (valueOf $ fromIntegral $ Serial.size a0, s1)))
      Maybe.lift $ do
         w3 <- Serial.writeNext ai w2
         (bi,r3) <- Serial.readNext r2
         i3 <- A.dec i2
         return (bi, ((w3,r3),(i3,s2))))
   (\p -> do
      s <- start p
      w <- Serial.writeZero
      return ((w, Class.undefTuple), (valueOf (0::Word32), s)))
   createIOContext
   deleteIOContext


raise ::
   (Storable a, IsArithmetic a, IsConst a,
    MakeValueTuple a, ValueTuple a ~ (Value a),
    Memory.FirstClass a, Memory.Stored a ~ am,
    IsPrimitive a,  IsSized a, SizeOf a ~ as,
    IsPrimitive am, IsSized am, SizeOf am ~ amsize,
    TypeNum.PositiveT (n :*: as),
    TypeNum.PositiveT (n :*: amsize),
    TypeNum.PositiveT n) =>
   Param.T p a ->
   T p (Serial.Value n a) (Serial.Value n a)
raise x =
   Causal.map Frame.mix (Serial.replicate ^<< x)


amplify ::
   (Storable a, IsArithmetic a, IsConst a,
    MakeValueTuple a, ValueTuple a ~ (Value a),
    Memory.FirstClass a, Memory.Stored a ~ am,
    IsPrimitive a,  IsSized a, SizeOf a ~ as,
    IsPrimitive am, IsSized am, SizeOf am ~ amsize,
    TypeNum.PositiveT (n :*: as),
    TypeNum.PositiveT (n :*: amsize),
    TypeNum.PositiveT n) =>
   Param.T p a ->
   T p (Serial.Value n a) (Serial.Value n a)
amplify p =
   Causal.map Frame.amplifyMono (Serial.replicate ^<< p)

amplifyStereo ::
   (Storable a, IsArithmetic a, IsConst a,
    MakeValueTuple a, ValueTuple a ~ (Value a),
    Memory.FirstClass a, Memory.Stored a ~ am,
    IsPrimitive a,  IsSized a, SizeOf a ~ as,
    IsPrimitive am, IsSized am, SizeOf am ~ amsize,
    TypeNum.PositiveT (n :*: as),
    TypeNum.PositiveT (n :*: amsize),
    TypeNum.PositiveT n) =>
   Param.T p a ->
   T p (Stereo.T (Serial.Value n a)) (Stereo.T (Serial.Value n a))
amplifyStereo p =
   Causal.map Frame.amplifyStereo (Serial.replicate ^<< p)


osciCore ::
   (Memory.FirstClass t, Memory.Stored t ~ tm, IsSized t, IsSized tm,
    SoV.Fraction t, IsConst t,
    Vector.Real t, IsPrimitive t,
    TypeNum.PositiveT n,
    Additive.C t) =>
   T p (Serial.Value n t, Serial.Value n t) (Serial.Value n t)
osciCore =
   Causal.zipWithSimple A.addToPhase <<<
   Arr.second
      (Causal.mapAccumSimple
         (\a phase0 -> do
            (phase1,b1) <- Serial.cumulate phase0 a
            phase2 <- A.signedFraction phase1
            return (b1,phase2))
         (return (valueOf Additive.zero)))

osciSimple ::
   (Memory.FirstClass t, Memory.Stored t ~ tm, IsSized t, IsSized tm,
    SoV.Fraction t, IsConst t,
    Vector.Real t, IsPrimitive t,
    TypeNum.PositiveT n,
    Additive.C t) =>
   (forall r. Serial.Value n t -> CodeGenFunction r y) ->
   T p (Serial.Value n t, Serial.Value n t) y
osciSimple wave =
   Causal.mapSimple wave <<< osciCore

shapeModOsci ::
   (Memory.FirstClass t, Memory.Stored t ~ tm, IsSized t, IsSized tm,
    SoV.Fraction t, IsConst t,
    Vector.Real t, IsPrimitive t,
    TypeNum.PositiveT n,
    Additive.C t) =>
   (forall r. c -> Serial.Value n t -> CodeGenFunction r y) ->
   T p (c, (Serial.Value n t, Serial.Value n t)) y
shapeModOsci wave =
   Causal.zipWithSimple wave <<< Arr.second osciCore


delay1 ::
   (Serial.C va, n ~ Serial.Size va, al ~ Serial.Element va,
    Storable a,
    MakeValueTuple a, ValueTuple a ~ al,
    Memory.C al) =>
   Param.T p a -> T p va va
delay1 initial =
   Causal.loop initial $
   Causal.mapSimple (fmap swap . uncurry Serial.shiftUp . swap)

differentiate ::
   (Serial.C va, n ~ Serial.Size va, al ~ Serial.Element va,
    A.Additive va,
    Storable a,
    MakeValueTuple a, ValueTuple a ~ al,
    Memory.C al) =>
   Param.T p a -> T p va va
differentiate initial =
   Cat.id - delay1 initial


arrayElement ::
   (IsFirstClass a, LLVM.Value a ~ Serial.Element v, Serial.C v,
    LLVM.GetValue (LLVM.Array dim a) index,
    LLVM.ValueType (LLVM.Array dim a) index ~ a,
    TypeNum.NaturalT index, TypeNum.NaturalT dim,
    (index :<: dim) ~ TypeBool.True) =>
   index -> T p (Value (LLVM.Array dim a)) v
arrayElement i =
   Causal.mapSimple Serial.upsample <<< Causal.arrayElement i