{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# 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.Sample as Sample
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.Representation as Rep
import qualified LLVM.Extra.Class as Class
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Control as C
import LLVM.Extra.Control (whileLoop, ifThen, )

import LLVM.Core as LLVM

import qualified Data.TypeLevel.Num as TypeNum
import qualified Data.TypeLevel.Num.Sets as TypeSet

import qualified Control.Arrow    as Arr
import Control.Arrow ((^<<), (<<<), )

-- import qualified Algebra.Transcendental as Trans
-- import qualified Algebra.Field as Field
import qualified Algebra.Ring as Ring
import qualified Algebra.Additive as Additive

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 ::
   (Vector.Access n a va, Vector.Access n b vb) =>
   T p a b -> T p va vb
pack (Cons next start createIOContext deleteIOContext) = Cons
   (\param a s -> do
      ((_,b2),(_,s2)) <-
         Maybe.fromBool $
         C.whileLoop
            (valueOf True,
             let b = undefTuple
             in  ((a,b), (valueOf $ (fromIntegral $ Vector.sizeInTuple b :: Word32), s)))
            (\(cont,(_ab0,(i0,_s0))) ->
               A.and cont =<<
                  A.icmp IntUGT i0 (value LLVM.zero))
            (\(_,((a0,b0),(i0,s0))) -> Maybe.toBool $ do
               ai <- Maybe.lift $ Vector.extract (valueOf 0) a0
               (bi,s1) <- next param ai s0
               Maybe.lift $ do
                  a1 <- Vector.rotateDown a0
                  b1 <- fmap snd $ Vector.shiftDown bi b0
                  i1 <- A.dec i0
                  return ((a1,b1),(i1,s1)))
      return (b2, 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 ::
   (Vector.Access n a va, Vector.Access n b vb) =>
   T p a b -> T p va vb
packSmall (Cons next start createIOContext deleteIOContext) = Cons
   (\param a s ->
      let vundef = LLVM.undefTuple
      in  foldr
             (\i rest (v0,s0) -> do
                ai <- Maybe.lift $ Vector.extract (valueOf i) a
                (bi,s1) <- next param ai s0
                v1 <- Maybe.lift $ Vector.insert (valueOf i) bi v0
                rest (v1,s1))
             return
             (take (Vector.sizeInTuple vundef) [0..])
             (vundef, s))
   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 ::
   (Vector.Access n a va, Vector.Access n b vb,
    Class.Zero va, LLVM.Undefined b,
    Rep.Memory va vap, IsSized vap vas,
    Rep.Memory vb vbp, IsSized vbp vbs) =>
   T p va vb -> T p a b
unpack (Cons next start createIOContext deleteIOContext) = Cons
   (\param ai ((a0,b0),(i0,s0)) -> do
      endOfVector <- Maybe.lift $ A.icmp IntEQ i0 (valueOf 0)
      ((a2,b2),(i2,s2)) <-
         Maybe.fromBool $
         C.ifThen endOfVector (valueOf True, ((a0,b0),(i0,s0))) $ do
            (cont1, (b1,s1)) <- Maybe.toBool $ next param a0 s0
            return (cont1,
                      ((LLVM.undefTuple, b1),
                       (valueOf $ fromIntegral $ Vector.sizeInTuple a0, s1)))
      Maybe.lift $ do
         a3 <- fmap snd $ Vector.shiftDown ai a2
         (bi,b3) <- Vector.shiftDown (LLVM.undefTuple) b2
         i3 <- A.dec i2
         return (bi, ((a3,b3),(i3,s2))))
   (\p -> do
      s <- start p
      return ((Class.zeroTuple, LLVM.undefTuple), (valueOf (0::Word32), s)))
   createIOContext
   deleteIOContext


raise ::
   (Storable a, IsArithmetic a, IsPrimitive a, IsConst a,
    MakeValueTuple a (Value a), IsFirstClass a, IsSized a size,
    IsPowerOf2 n, TypeNum.Mul n size ps, TypeSet.Pos ps) =>
   Param.T p a ->
   T p (Value (Vector n a)) (Value (Vector n a))
raise x =
   Causal.map Sample.mixMono (LLVM.vector . (:[]) ^<< x)


amplify ::
   (Storable a, IsArithmetic a, IsPrimitive a, IsConst a,
    MakeValueTuple a (Value a), IsFirstClass a, IsSized a size,
    IsPowerOf2 n, TypeNum.Mul n size ps, TypeSet.Pos ps) =>
   Param.T p a ->
   T p (Value (Vector n a)) (Value (Vector n a))
amplify p =
   Causal.map Sample.amplifyMono (LLVM.vector . (:[]) ^<< p)

amplifyStereo ::
   (Storable a, IsArithmetic a, IsPrimitive a, IsConst a,
    MakeValueTuple a (Value a), IsFirstClass a, IsSized a size,
    IsPowerOf2 n, TypeNum.Mul n size ps, TypeSet.Pos ps) =>
   Param.T p a ->
   T p (Stereo.T (Value (Vector n a))) (Stereo.T (Value (Vector n a)))
amplifyStereo p =
   Causal.map Sample.amplifyStereo (LLVM.vector . (:[]) ^<< p)


osciCore ::
   (IsFirstClass t, IsSized t size,
    SoV.Fraction t, IsConst t,
    Vector.Real t, IsPrimitive t,
    IsPowerOf2 n,
    Additive.C t) =>
   T p (Value (Vector n t), Value (Vector n t)) (Value (Vector n t))
osciCore =
   Causal.mapSimple (uncurry SoV.addToPhase) <<<
   Arr.second
      (Causal.mapAccumSimple
         (\a phase0 -> do
            (phase1,b1) <- Vector.cumulate phase0 a
            phase2 <- SoV.signedFraction phase1
            return (b1,phase2))
         (return (valueOf Additive.zero)))

osciSimple ::
   (IsFirstClass t, IsSized t size,
    SoV.Fraction t, IsConst t,
    Vector.Real t, IsPrimitive t,
    IsPowerOf2 n,
    Additive.C t) =>
   (forall r. Value (Vector n t) -> CodeGenFunction r y) ->
   T p (Value (Vector n t), Value (Vector n t)) y
osciSimple wave =
   Causal.mapSimple wave <<< osciCore

shapeModOsci ::
   (IsFirstClass t, IsSized t size,
    SoV.Fraction t, IsConst t,
    Vector.Real t, IsPrimitive t,
    IsPowerOf2 n,
    Additive.C t) =>
   (forall r. c -> Value (Vector n t) -> CodeGenFunction r y) ->
   T p (c, (Value (Vector n t), Value (Vector n t))) y
shapeModOsci wave =
   Causal.mapSimple (uncurry wave) <<< Arr.second osciCore