{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ForeignFunctionInterface #-}
module Synthesizer.LLVM.Causal.Process (
   C(simple, loopConst, replicateControlled),
   T,
   amplify,
   amplifyStereo,
   apply,
   applyFst,
   applySnd,
   compose,
   envelope,
   envelopeStereo,
   feedFst,
   feedSnd,
   fromModifier,
   fromSignal,
   toSignal,
   loopZero,
   feedbackControlledZero,
   map,
   mapAccum,
   zipWith,
   mix,
   pipeline,
   stereoFromVector,
   vectorize,
   replaceChannel,
   arrayElement,
   element,
   applyStorable,
   applyStorableChunky,
   runStorableChunky,
   ) where

import qualified Synthesizer.LLVM.Simple.Signal as Sig
import qualified Synthesizer.LLVM.Simple.Value as Value
import qualified Synthesizer.LLVM.Causal.ProcessPrivate as Causal
import qualified Synthesizer.LLVM.Frame.Stereo as Stereo
import qualified Synthesizer.LLVM.Frame as Frame
import qualified Synthesizer.LLVM.Execution as Exec

import qualified Synthesizer.Plain.Modifier as Modifier
import qualified Synthesizer.Causal.Class as CausalClass

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

import qualified LLVM.Extra.Class as Class
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Vector as Vector
import qualified LLVM.Extra.MaybeContinuation as MaybeCont
import qualified LLVM.Extra.Maybe as Maybe
import qualified LLVM.Extra.ForeignPtr as ForeignPtr
import qualified LLVM.Extra.Memory as Memory
import LLVM.Extra.Class (Undefined, MakeValueTuple, ValueTuple, )

import qualified LLVM.Core as LLVM
import LLVM.ExecutionEngine (simpleFunction, )
import LLVM.Util.Loop (Phi, )
import LLVM.Core
          (CodeGenFunction, ret, Value, valueOf,
           IsConst, IsFirstClass, IsArithmetic, IsPrimitive,
           Linkage(ExternalLinkage), createNamedFunction)

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

import qualified Control.Arrow    as Arr
import qualified Control.Category as Cat
import Control.Monad.Trans.State (runState, )
import Control.Arrow ((<<<), (^<<), (>>>), (&&&), )
import Control.Monad (liftM2, liftM3, )
import Control.Applicative (Applicative, pure, (<*>), )

import qualified Data.List as List
import Data.Word (Word32, )
import Foreign.Storable (Storable, )
import Foreign.ForeignPtr (withForeignPtr, touchForeignPtr, )
import Foreign.Ptr (FunPtr, Ptr, )
import Control.Exception (bracket, )
import qualified System.Unsafe as Unsafe

import qualified Number.Ratio as Ratio
import qualified Algebra.Field as Field
import qualified Algebra.Ring as Ring
import qualified Algebra.Additive as Additive

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

import qualified Prelude as P


data T a b =
   forall state ioContext.
      (Memory.C state) =>
      Cons (forall r c.
            (Phi c) =>
            ioContext ->
            a -> state -> MaybeCont.T r c (b, state))
               -- compute next value
           (forall r.
            ioContext ->
            CodeGenFunction r state)
               -- initial state
           (IO ioContext)
               -- initialization from IO monad
           (ioContext -> IO ())
               -- finalization from IO monad


class CausalClass.C process => C process where
   simple ::
      (Memory.C state) =>
      (forall r c.
       (Phi c) =>
       a -> state -> MaybeCont.T r c (b, state)) ->
      (forall r. CodeGenFunction r state) ->
      process a b
   {- |
   Like 'Synthesizer.LLVM.CausalParameterized.loop'
   but uses zero as initial value
   and it does not need a zero as Haskell value.
   -}
   loopConst ::
      (Memory.C c) =>
      c -> process (a,c) (b,c) -> process a b
   replicateControlled ::
      (Undefined x, Phi x) =>
      Int -> process (c,x) x -> process (c,x) x

instance CausalClass.C T where
   type SignalOf T = Sig.T
   type ProcessOf Sig.T = T
   toSignal = toSignal
   fromSignal = fromSignal

instance C T where
   simple next start =
      Cons
         (const next)
         (const start)
         (return ())
         (const $ return ())
   loopConst init (Cons next start create delete) =
      Cons
         (Causal.loopNext next)
         (fmap ((,) init) . start)
         create
         delete
   {-
   Could be implemented with a machine code loop like in CausalParameterized.
   But to this end we would need a 'stop' function.
   -}
   replicateControlled = CausalClass.replicateControlled

toSignal :: T () a -> Sig.T a
toSignal (Cons next start createIOContext deleteIOContext) = Sig.Cons
   (\ioContext -> next ioContext ())
   start
   createIOContext deleteIOContext

fromSignal :: Sig.T b -> T a b
fromSignal (Sig.Cons next start createIOContext deleteIOContext) = Cons
   (\ioContext _ -> next ioContext)
   start
   createIOContext deleteIOContext


map ::
   (C process) =>
   (forall r. a -> CodeGenFunction r b) ->
   process a b
map f =
   mapAccum (\a s -> fmap (flip (,) s) $ f a) (return ())

mapAccum ::
   (C process, Memory.C state) =>
   (forall r.
    a -> state -> CodeGenFunction r (b, state)) ->
   (forall r. CodeGenFunction r state) ->
   process a b
mapAccum next =
   simple (\a s -> MaybeCont.lift $ next a s)

zipWith ::
   (C process) =>
   (forall r. a -> b -> CodeGenFunction r c) ->
   process (a,b) c
zipWith f = map (uncurry f)


fromModifier ::
   (C process) =>
   (Value.Flatten ah, Value.Registers ah ~ al,
    Value.Flatten bh, Value.Registers bh ~ bl,
    Value.Flatten ch, Value.Registers ch ~ cl,
    Value.Flatten sh, Value.Registers sh ~ sl,
    Memory.C sl) =>
   Modifier.Simple sh ch ah bh -> process (cl,al) bl
fromModifier (Modifier.Simple initial step) =
   mapAccum
      (\(c,a) s ->
         Value.flatten $
         runState
            (step (Value.unfold c) (Value.unfold a))
            (Value.unfold s))
      (Value.flatten initial)


apply :: T a b -> Sig.T a -> Sig.T b
apply = CausalClass.apply

feedFst :: Sig.T a -> T b (a,b)
feedFst = CausalClass.feedFst

feedSnd :: Sig.T a -> T b (b,a)
feedSnd = CausalClass.feedSnd


applyFst :: T (a,b) c -> Sig.T a -> T b c
applyFst = CausalClass.applyFst

applySnd :: T (a,b) c -> Sig.T b -> T a c
applySnd = CausalClass.applySnd

compose :: T a b -> T b c -> T a c
compose
      (Cons nextA startA createIOContextA deleteIOContextA)
      (Cons nextB startB createIOContextB deleteIOContextB) = Cons
   (\(ioContextA, ioContextB) a (sa0,sb0) -> do
      (b,sa1) <- nextA ioContextA a sa0
      (c,sb1) <- nextB ioContextB b sb0
      return (c, (sa1,sb1)))
   (\(ioContextA, ioContextB) ->
      liftM2 (,)
         (startA ioContextA)
         (startB ioContextB))
   (liftM2 (,)
      createIOContextA
      createIOContextB)
   (\(ca,cb) ->
      deleteIOContextA ca >>
      deleteIOContextB cb)


first :: T b c -> T (b, d) (c, d)
first (Cons next start createIOContext deleteIOContext) = Cons
   (\ioContext (b,d) sa0 ->
      fmap
         (\(c,sa1) -> ((c,d), sa1))
         (next ioContext b sa0))
   start
   createIOContext deleteIOContext


instance Cat.Category T where
   id = map return
   (.) = flip compose

instance Arr.Arrow T where
   arr f = map (return . f)
   first = first



instance Functor (T a) where
   fmap = (^<<)

instance Applicative (T a) where
   pure x = Arr.arr (const x)
   f <*> x = uncurry ($) ^<< f&&&x


instance (A.Additive b) => Additive.C (T a b) where
   zero = pure A.zero
   negate x = map A.neg <<< x
   x + y = zipWith A.add <<< x&&&y
   x - y = zipWith A.sub <<< x&&&y

instance (A.PseudoRing b, A.IntegerConstant b) => Ring.C (T a b) where
   one = pure A.one
   fromInteger n = pure (A.fromInteger' n)
   x * y = zipWith A.mul <<< x&&&y

instance (A.Field b, A.RationalConstant b) => Field.C (T a b) where
   fromRational' x = pure (A.fromRational' $ Ratio.toRational98 x)
   x / y = zipWith A.fdiv <<< x&&&y


instance (A.PseudoRing b, A.Real b, A.IntegerConstant b) => P.Num (T a b) where
   fromInteger n = pure (A.fromInteger' n)
   negate x = map A.neg <<< x
   x + y = zipWith A.add <<< x&&&y
   x - y = zipWith A.sub <<< x&&&y
   x * y = zipWith A.mul <<< x&&&y
   abs x = map A.abs <<< x
   signum x = map A.signum <<< x

instance (A.Field b, A.Real b, A.RationalConstant b) => P.Fractional (T a b) where
   fromRational x = pure (A.fromRational' x)
   x / y = zipWith A.fdiv <<< x&&&y


mix ::
   (C process, A.Additive a) =>
   process (a, a) a
mix = zipWith Frame.mix


envelope ::
   (C process, A.PseudoRing a) =>
   process (a, a) a
envelope = zipWith Frame.amplifyMono

envelopeStereo ::
   (C process, A.PseudoRing a) =>
   process (a, Stereo.T a) (Stereo.T a)
envelopeStereo = zipWith Frame.amplifyStereo

amplify ::
   (C process, IsArithmetic a, IsConst a) =>
   a -> process (Value a) (Value a)
amplify x =
   map (Frame.amplifyMono (valueOf x))

amplifyStereo ::
   (C process, IsArithmetic a, IsConst a) =>
   a -> process (Stereo.T (Value a)) (Stereo.T (Value a))
amplifyStereo x =
   map (Frame.amplifyStereo (valueOf x))



loopZero ::
   (C process, A.Additive c, Memory.C c) =>
   process (a,c) (b,c) -> process a b
loopZero = loopConst A.zero

pipeline ::
   (C process,
    Vector.C v, a ~ Vector.Element v,
    Class.Zero v, Memory.C v) =>
   process v v -> process a a
pipeline vectorProcess =
   loopConst Class.zeroTuple $
      map (uncurry Vector.shiftUp)
      >>>
      Arr.second vectorProcess


feedbackControlledZero ::
   (C process, A.Additive c, Memory.C c) =>
   process ((ctrl,a),c) b -> process (ctrl,b) c -> process (ctrl,a) b
feedbackControlledZero forth back =
   loopZero (Causal.feedbackControlledAux forth back)


{-
In order to let this work we have to give the disable-mmx option somewhere,
but where?
-}
stereoFromVector ::
   (C process, IsPrimitive a, IsPrimitive b) =>
   process (Value (LLVM.Vector D2 a)) (Value (LLVM.Vector D2 b)) ->
   process (Stereo.T (Value a)) (Stereo.T (Value b))
stereoFromVector proc =
   map Frame.stereoFromVector <<<
   proc <<<
   map Frame.vectorFromStereo


{-
insert and extract instructions will be in opposite order,
no matter whether we use foldr or foldl
and independent from the order of proc and channel in replaceChannel.
However, LLVM neglects the order anyway.
-}
vectorize ::
   (C process,
    Vector.C va, n ~ Vector.Size va, a ~ Vector.Element va,
    Vector.C vb, n ~ Vector.Size vb, b ~ Vector.Element vb) =>
   process a b -> process va vb
vectorize proc =
   withSize $ \n ->
      foldl
         (\acc i -> replaceChannel i proc acc)
         (Arr.arr (const $ Class.undefTuple)) $
      List.take (TypeNum.fromIntegerT n) [0 ..]

withSize ::
   (Vector.Size bv -> f bv) ->
   f bv
withSize f = f undefined

{- |
Given a vector process, replace the i-th output by output
that is generated by a scalar process from the i-th input.
-}
replaceChannel ::
   (C process,
    Vector.C va, n ~ Vector.Size va, a ~ Vector.Element va,
    Vector.C vb, n ~ Vector.Size vb, b ~ Vector.Element vb) =>
   Int -> process a b -> process va vb -> process va vb
replaceChannel i channel proc =
   let li = valueOf $ fromIntegral i
   in  zipWith (Vector.insert li) <<<
          (channel <<< map (Vector.extract li)) &&&
          proc

{- |
Read the i-th element from each array.
-}
arrayElement ::
   (C process,
    LLVM.Array dim a ~ array,
    LLVM.GetValue array index, IsFirstClass a,
    LLVM.ValueType array index ~ a,
    TypeNum.NaturalT index, TypeNum.NaturalT dim,
    (index :<: dim) ~ TypeBool.True) =>
   index -> process (Value (LLVM.Array dim a)) (Value a)
arrayElement i =
   map (\array -> LLVM.extractvalue array i)

{- |
Read the i-th element from an aggregate type.
-}
element ::
   (C process, IsFirstClass a, LLVM.GetValue agg index,
    LLVM.ValueType agg index ~ a) =>
   index -> process (Value agg) (Value a)
element i =
   map (\array -> LLVM.extractvalue array i)


applyStorable ::
   (Storable a, MakeValueTuple a, ValueTuple a ~ valueA, Memory.C valueA,
    Storable b, MakeValueTuple b, ValueTuple b ~ valueB, Memory.C valueB) =>
   T valueA valueB -> SV.Vector a -> SV.Vector b
applyStorable (Cons next start createIOContext deleteIOContext) as =
   Unsafe.performIO $
   bracket createIOContext deleteIOContext $ \ ioContext ->
   SVB.withStartPtr as $ \ aPtr len ->
   SVB.createAndTrim len $ \ bPtr -> do
      fill <-
         simpleFunction $
         createNamedFunction ExternalLinkage "fillprocessblock" $ \ size alPtr blPtr -> do
            s <- start ioContext
            (pos,_) <- MaybeCont.arrayLoop2 size alPtr blPtr s $
                  \ aPtri bPtri s0 -> do
               a <- MaybeCont.lift $ Memory.load aPtri
               (b,s1) <- next ioContext a s0
               MaybeCont.lift $ Memory.store b bPtri
               return s1
            ret (pos :: Value Word32)
      fmap (fromIntegral :: Word32 -> Int) $
         fill (fromIntegral len)
            (Memory.castStorablePtr aPtr)
            (Memory.castStorablePtr bPtr)


foreign import ccall safe "dynamic" derefChunkPtr ::
   Exec.Importer (Ptr stateStruct -> Word32 ->
             Ptr aStruct -> Ptr bStruct -> IO Word32)


compileChunky ::
   (Memory.C aValue, Memory.Struct aValue ~ aStruct,
    Memory.C bValue, Memory.Struct bValue ~ bStruct,
    Memory.C state, Memory.Struct state ~ stateStruct) =>
   (forall r z.
    (Phi z) =>
    aValue -> state ->
    MaybeCont.T r z (bValue, state)) ->
   (forall r.
    CodeGenFunction r state) ->
   IO (FunPtr (IO (Ptr stateStruct)),
       FunPtr (Ptr stateStruct -> IO ()),
       FunPtr (Ptr stateStruct -> Word32 -> Ptr aStruct -> Ptr bStruct -> IO Word32))
compileChunky next start =
   Exec.compileModule $
      liftM3 (,,)
         (createNamedFunction ExternalLinkage "startprocess" $
          do
             pptr <- LLVM.malloc
             flip Memory.store pptr =<< start
             ret pptr)
         (createNamedFunction ExternalLinkage "stopprocess" $
          \ pptr -> LLVM.free pptr >> ret ())
         (createNamedFunction ExternalLinkage "fillprocess" $
          \ sptr loopLen aPtr bPtr -> do
             sInit <- Memory.load sptr
             (pos,sExit) <- MaybeCont.arrayLoop2 loopLen aPtr bPtr sInit $
                   \ aPtri bPtri s0 -> do
                a <- MaybeCont.lift $ Memory.load aPtri
                (b,s1) <- next a s0
                MaybeCont.lift $ Memory.store b bPtri
                return s1
             Memory.store (Maybe.fromJust sExit) sptr
             ret (pos :: Value Word32))


{-# DEPRECATED runStorableChunky "this function will not work when the process itself depends on a lazy storable vector" #-}
{- |
This function will not work as expected,
since feeding a lazy storable vector to the causal process
means that createIOContext creates a StablePtr to an IORef refering to a chunk list.
The IORef will be created once for all uses of the generated function
of type @(SVL.Vector a -> SVL.Vector b)@.
This means that the pointer into the chunks list will conflict.
An alternative would be to create the StablePtr in a foreign function
that calls back to Haskell.
But this way is disallowed for foreign finalizers.
-}
runStorableChunky ::
   (Storable a, MakeValueTuple a, ValueTuple a ~ valueA, Memory.C valueA,
    Storable b, MakeValueTuple b, ValueTuple b ~ valueB, Memory.C valueB) =>
   T valueA valueB -> IO (SVL.Vector a -> SVL.Vector b)
runStorableChunky
      (Cons next start createIOContext deleteIOContext) = do
   ioContext <- createIOContext
   (startFunc, stopFunc, fill) <-
      compileChunky (next ioContext) (start ioContext)

   {-
   This is a dummy pointer, that we need for correct finalization.
   Concerning the live time the FunPtr 'fill' also has the live time
   that we are after,
   but it is unsafe to treat a FunPtr as a Ptr or ForeignPtr.
   -}
   ioContextPtr <- ForeignPtr.new (deleteIOContext ioContext) False

   return $ \sig -> SVL.fromChunks $ Unsafe.performIO $ do
      statePtr <- ForeignPtr.newInit stopFunc startFunc
      let go xt =
             Unsafe.interleaveIO $
             case xt of
                [] -> return []
                x:xs -> SVB.withStartPtr x $ \aPtr size -> do
                   v <-
                      withForeignPtr statePtr $ \sptr ->
                      SVB.createAndTrim size $
                         fmap (fromIntegral :: Word32 -> Int) .
                         derefChunkPtr fill sptr (fromIntegral size)
                            (Memory.castStorablePtr aPtr) .
                         Memory.castStorablePtr
                   touchForeignPtr ioContextPtr
                   (if SV.length v > 0
                      then fmap (v:)
                      else id) $
                      (if SV.length v < size
                         then return []
                         else go xs)
      go (SVL.chunks sig)


applyStorableChunky ::
   (Storable a, MakeValueTuple a, ValueTuple a ~ valueA, Memory.C valueA,
    Storable b, MakeValueTuple b, ValueTuple b ~ valueB, Memory.C valueB) =>
   T valueA valueB -> SVL.Vector a -> SVL.Vector b
applyStorableChunky
     (Cons next start createIOContext deleteIOContext) sig =
   SVL.fromChunks $ Unsafe.performIO $ do
      ioContext <- createIOContext
      (startFunc, stopFunc, fill) <-
         compileChunky (next ioContext) (start ioContext)

      statePtr <- ForeignPtr.newInit stopFunc startFunc
      {-
      This is a dummy pointer, that we need for correct finalization.
      Concerning the live time the FunPtr 'fill' also has the live time
      that we are after,
      but it is unsafe to treat a FunPtr as a Ptr or ForeignPtr.
      -}
      ioContextPtr <- ForeignPtr.new (deleteIOContext ioContext) False

      let go xt =
             Unsafe.interleaveIO $
             case xt of
                [] -> return []
                x:xs -> SVB.withStartPtr x $ \aPtr size -> do
                   v <-
                      withForeignPtr statePtr $ \sptr ->
                      SVB.createAndTrim size $
                         fmap (fromIntegral :: Word32 -> Int) .
                         derefChunkPtr fill sptr (fromIntegral size)
                            (Memory.castStorablePtr aPtr) .
                         Memory.castStorablePtr
                   touchForeignPtr ioContextPtr
                   (if SV.length v > 0
                      then fmap (v:)
                      else id) $
                      (if SV.length v < size
                         then return []
                         else go xs)
      go (SVL.chunks sig)