{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# 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, ValueTuple, )

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

import qualified Prelude as P


data T a b =
   forall state ioContext.
      (Memory.C state) =>
      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) =>
   (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) =>
   (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


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 ::
   (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, 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,_) <- 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, Memory.Struct aValue ~ aStruct,
    Memory.C bValue, Memory.Struct bValue ~ bStruct,
    Memory.C state, Memory.Struct state ~ stateStruct) =>
   (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, 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)