{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ForeignFunctionInterface #-}
module Synthesizer.LLVM.Causal.Process where

import qualified Synthesizer.LLVM.Simple.Signal as Sig
import qualified Synthesizer.LLVM.Frame as Frame
import qualified Synthesizer.LLVM.Execution as Exec

import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.MaybeContinuation as Maybe
import qualified LLVM.Extra.ForeignPtr as ForeignPtr
import qualified LLVM.Extra.Memory as Memory
-- import qualified LLVM.Extra.Control as U
import LLVM.Extra.Class (MakeValueTuple, )

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

import qualified Synthesizer.LLVM.Frame.Stereo as Stereo

import LLVM.Core as LLVM
import LLVM.Util.Loop (Phi, )
import LLVM.ExecutionEngine (simpleFunction, )

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

import Data.Word (Word32, )
import Foreign.Storable (Storable, )
import Foreign.ForeignPtr (withForeignPtr, touchForeignPtr, )
import Foreign.Ptr (FunPtr, )
import Control.Exception (bracket, )
import System.IO.Unsafe (unsafePerformIO, unsafeInterleaveIO, )

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, )


data T a b =
   forall state packed size ioContext.
      (Memory.C state packed, IsSized packed size) =>
      Cons (forall r c.
            (Phi c) =>
            ioContext ->
            a -> state -> Maybe.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

simple ::
   (Memory.C state packed, IsSized packed size) =>
   (forall r c.
    (Phi c) =>
    a -> state -> Maybe.T r c (b, state)) ->
   (forall r. CodeGenFunction r state) ->
   T a b
simple next start =
   Cons
      (const next)
      (const start)
      (return ())
      (const $ return ())

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 ::
   (forall r. a -> CodeGenFunction r b) ->
   T a b
map f =
   mapAccum (\a s -> fmap (flip (,) s) $ f a) (return ())

mapAccum ::
   (Memory.C state packed, IsSized packed size) =>
   (forall r.
    a -> state -> CodeGenFunction r (b, state)) ->
   (forall r. CodeGenFunction r state) ->
   T a b
mapAccum next =
   simple (\a s -> Maybe.lift $ next a s)

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


apply ::
   T a b -> Sig.T a -> Sig.T b
apply proc sig =
   toSignal (proc <<< fromSignal sig)

feedFst :: Sig.T a -> T b (a,b)
feedFst sig =
   fromSignal sig &&& Cat.id

feedSnd :: Sig.T a -> T b (b,a)
feedSnd sig =
   Cat.id &&& fromSignal sig


applyFst :: T (a,b) c -> Sig.T a -> T b c
applyFst proc sig =
   proc <<< feedFst sig

applySnd :: T (a,b) c -> Sig.T b -> T a c
applySnd proc sig =
   proc <<< feedSnd sig

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


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


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

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

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

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



applyStorable ::
   (Storable a, MakeValueTuple a valueA, Memory.C valueA structA,
    Storable b, MakeValueTuple b valueB, Memory.C valueB structB) =>
   T valueA valueB -> SV.Vector a -> SV.Vector b
applyStorable (Cons next start createIOContext deleteIOContext) as =
   unsafePerformIO $
   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,_) <- Maybe.arrayLoop2 size alPtr blPtr s $
                  \ aPtri bPtri s0 -> do
               a <- Maybe.lift $ Memory.load aPtri
               (b,s1) <- next ioContext a s0
               Maybe.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" derefStartPtr ::
   Exec.Importer (IO (Ptr stateStruct))

foreign import ccall safe "dynamic" derefStopPtr ::
   Exec.Importer (Ptr stateStruct -> IO ())

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


compileChunky ::
   (Memory.C aValue aStruct,
    Memory.C bValue bStruct,
    Memory.C state stateStruct,
    IsSized stateStruct stateSize) =>
   (forall r.
    aValue -> state ->
    Maybe.T r (Value Bool, (Value (Ptr bStruct), state)) (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) <- Maybe.arrayLoop2 loopLen aPtr bPtr sInit $
                   \ aPtri bPtri s0 -> do
                a <- Maybe.lift $ Memory.load aPtri
                (b,s1) <- next a s0
                Maybe.lift $ Memory.store b bPtri
                return s1
             Memory.store 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 valueA, Memory.C valueA structA,
    Storable b, MakeValueTuple b valueB, Memory.C valueB structB) =>
   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 $ unsafePerformIO $ do
      statePtr <- ForeignPtr.newInit stopFunc startFunc
      let go xt =
             unsafeInterleaveIO $
             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 valueA, Memory.C valueA structA,
    Storable b, MakeValueTuple b valueB, Memory.C valueB structB) =>
   T valueA valueB -> SVL.Vector a -> SVL.Vector b
applyStorableChunky
     (Cons next start createIOContext deleteIOContext) sig =
   SVL.fromChunks $ unsafePerformIO $ 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 =
             unsafeInterleaveIO $
             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)