{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ForeignFunctionInterface #-}
module Synthesizer.LLVM.Causal.Process where

import qualified Synthesizer.LLVM.Simple.Signal as Sig
import qualified LLVM.Extra.Representation as Rep
import qualified Synthesizer.LLVM.Sample as Sample
import qualified Synthesizer.LLVM.Execution as Exec
import qualified LLVM.Extra.MaybeContinuation as Maybe
-- import qualified LLVM.Extra.Control as U

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
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 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 Data.Tuple.HT (swap, )

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


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


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 =
   first (fromSignal sig) <<^ (\b -> ((),b))

feedSnd :: Sig.T a -> T b (b,a)
feedSnd sig =
   swap ^<< feedFst 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


mix ::
   (IsArithmetic a) =>
   T (Value a, Value a) (Value a)
mix = map (uncurry Sample.mixMono)

mixStereo ::
   (IsArithmetic a) =>
   T (Stereo.T (Value a), Stereo.T (Value a)) (Stereo.T (Value a))
mixStereo = map (uncurry Sample.mixStereo)


envelope ::
   (IsArithmetic a) =>
   T (Value a, Value a) (Value a)
envelope = map (uncurry Sample.amplifyMono)

envelopeStereo ::
   (IsArithmetic a) =>
   T (Value a, Stereo.T (Value a)) (Stereo.T (Value a))
envelopeStereo = map (uncurry Sample.amplifyStereo)

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

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



applyStorable ::
   (Storable a, MakeValueTuple a valueA, Rep.Memory valueA structA,
    Storable b, MakeValueTuple b valueB, Rep.Memory 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 $
         createFunction ExternalLinkage $ \ size alPtr blPtr -> do
            s <- start ioContext
            (pos,_) <- Maybe.arrayLoop2 size alPtr blPtr s $
                  \ aPtri bPtri s0 -> do
               a <- Maybe.lift $ Rep.load aPtri
               (b,s1) <- next ioContext a s0
               Maybe.lift $ Rep.store b bPtri
               return s1
            ret (pos :: Value Word32)
      fmap (fromIntegral :: Word32 -> Int) $
         fill (fromIntegral len)
            (Rep.castStorablePtr aPtr)
            (Rep.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 ::
   (Rep.Memory aValue aStruct,
    Rep.Memory bValue bStruct,
    Rep.Memory 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 (,,)
         (createFunction ExternalLinkage $
          do
             -- FIXME: size computation in LLVM currently does not work for structs!
             pptr <- Rep.malloc
             flip Rep.store pptr =<< start
             ret pptr)
         (createFunction ExternalLinkage $
          \ pptr -> Rep.free pptr >> ret ())
         (createFunction ExternalLinkage $
          \ sptr loopLen aPtr bPtr -> do
             sInit <- Rep.load sptr
             (pos,sExit) <- Maybe.arrayLoop2 loopLen aPtr bPtr sInit $
                   \ aPtri bPtri s0 -> do
                a <- Maybe.lift $ Rep.load aPtri
                (b,s1) <- next a s0
                Maybe.lift $ Rep.store b bPtri
                return s1
             Rep.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, Rep.Memory valueA structA,
    Storable b, MakeValueTuple b valueB, Rep.Memory 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 <- Rep.newForeignPtr (deleteIOContext ioContext) False

   return $ \sig -> SVL.fromChunks $ unsafePerformIO $ do
      statePtr <- Rep.newForeignPtrInit 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)
                            (Rep.castStorablePtr aPtr) .
                         Rep.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, Rep.Memory valueA structA,
    Storable b, MakeValueTuple b valueB, Rep.Memory 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 <- Rep.newForeignPtrInit 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 <- Rep.newForeignPtr (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)
                            (Rep.castStorablePtr aPtr) .
                         Rep.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)