{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
module Synthesizer.LLVM.CausalParameterized.ProcessPrivate where

import qualified Synthesizer.LLVM.Parameterized.SignalPrivate as Sig
import qualified Synthesizer.LLVM.Parameter as Param

import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.MaybeContinuation as Maybe
import qualified LLVM.Extra.Memory as Memory

import qualified LLVM.Core as LLVM
import LLVM.Extra.Class (MakeValueTuple, )
import LLVM.Util.Loop (Phi, )
import LLVM.Core
          (Value, valueOf,
           IsSized, CodeGenFunction, )

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

import Data.Word (Word32, )
import Foreign.Storable.Tuple ()
import Foreign.Storable (Storable, )

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, iterate, map, zip, zipWith, take, takeWhile, )


data T p a b =
   forall state packed size ioContext
        startParamTuple startParamValue startParamPacked startParamSize
        nextParamTuple  nextParamValue  nextParamPacked  nextParamSize.
      (Storable startParamTuple,
       Storable nextParamTuple,
       MakeValueTuple startParamTuple startParamValue,
       MakeValueTuple nextParamTuple  nextParamValue,
       Memory.C     startParamValue startParamPacked,
       Memory.C     nextParamValue  nextParamPacked,
       IsSized        startParamPacked startParamSize,
       IsSized        nextParamPacked  nextParamSize,
       Memory.C state packed,
       IsSized packed size) =>
   Cons
      (forall r c.
       (Phi c) =>
       nextParamValue ->
       a -> state -> Maybe.T r c (b, state))
          -- compute next value
      (forall r.
       startParamValue ->
       CodeGenFunction r state)
          -- initial state
      (p -> IO (ioContext, (nextParamTuple, startParamTuple)))
          {- initialization from IO monad
          This will be run within unsafePerformIO,
          so no observable In/Out actions please!
          -}
      (ioContext -> IO ())
          -- finalization from IO monad, also run within unsafePerformIO


simple ::
   (Storable startParamTuple,
    Storable nextParamTuple,
    MakeValueTuple startParamTuple startParamValue,
    MakeValueTuple nextParamTuple nextParamValue,
    Memory.C startParamValue startParamPacked,
    Memory.C nextParamValue nextParamPacked,
    IsSized    startParamPacked startParamSize,
    IsSized    nextParamPacked  nextParamSize,
    Memory.C state packed,
    IsSized packed size) =>
   (forall r c.
    (Phi c) =>
    nextParamValue ->
    a -> state -> Maybe.T r c (b, state)) ->
   (forall r.
    startParamValue ->
    CodeGenFunction r state) ->
   Param.T p nextParamTuple ->
   Param.T p startParamTuple -> T p a b
simple f start selectParam initial = Cons
   (f . Param.value selectParam)
   (start . Param.value initial)
   (return . (,) () . Param.get (selectParam &&& initial))
   (const $ return ())


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

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


mapAccum ::
   (Storable pnh, MakeValueTuple pnh pnl, Memory.C pnl pnp, IsSized pnp pns,
    Storable psh, MakeValueTuple psh psl, Memory.C psl psp, IsSized psp pss,
    Memory.C s struct, IsSized struct sa) =>
   (forall r. pnl -> a -> s -> CodeGenFunction r (b,s)) ->
   (forall r. psl -> CodeGenFunction r s) ->
   Param.T p pnh ->
   Param.T p psh ->
   T p a b
mapAccum next start selectParamN selectParamS =
   simple
      (\p a s -> Maybe.lift $ next p a s)
      start
      selectParamN selectParamS


map ::
   (Storable ph, MakeValueTuple ph pl, Memory.C pl pp, IsSized pp ps) =>
   (forall r. pl -> a -> CodeGenFunction r b) ->
   Param.T p ph ->
   T p a b
map f selectParamF =
   mapAccum
      (\p a s -> fmap (flip (,) s) $ f p a)
      (const $ return ())
      selectParamF
      (return ())

mapSimple ::
   (forall r. a -> CodeGenFunction r b) ->
   T p a b
mapSimple f =
   map (const f) (return ())

zipWithSimple ::
   (forall r. a -> b -> CodeGenFunction r c) ->
   T p (a,b) c
zipWithSimple f =
   mapSimple (uncurry f)


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

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

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


{-
Very similar to 'apply',
since 'apply' can be considered being of type
@T p a b -> T p () a -> T p () b@.
-}
compose :: T p a b -> T p b c -> T p a c
compose
      (Cons nextA startA createIOContextA deleteIOContextA)
      (Cons nextB startB createIOContextB deleteIOContextB) =
   Cons
      (\(paramA, paramB) a (sa0,sb0) ->
         do (b,sa1) <- nextA paramA a sa0
            (c,sb1) <- nextB paramB b sb0
            return (c, (sa1,sb1)))
      (\(paramA, paramB) ->
         liftM2 (,)
            (startA paramA)
            (startB paramB))
      (\p -> do
         (ca,(nextParamA,startParamA)) <- createIOContextA p
         (cb,(nextParamB,startParamB)) <- createIOContextB p
         return ((ca,cb),
            ((nextParamA,  nextParamB),
             (startParamA, startParamB))))
      (\(ca,cb) ->
         deleteIOContextA ca >>
         deleteIOContextB cb)


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


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

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


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

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


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

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

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


{- |
Not quite the loop of ArrowLoop
because we need a delay of one time step
and thus an initialization value.

For a real ArrowLoop.loop, that is a zero-delay loop,
we would formally need a MonadFix instance of CodeGenFunction.
But this will not become reality, since LLVM is not able to re-order code
in a way that allows to access a result before creating the input.
-}
loop ::
   (Storable ch,
    MakeValueTuple ch c,
    Memory.C c cp,
    IsSized cp cs) =>
   Param.T p ch -> T p (a,c) (b,c) -> T p a b
loop initial (Cons next start createIOContext deleteIOContext) =
   Cons
      (\p a0 (c0,s0) -> do
         ((b1,c1), s1) <- next p (a0,c0) s0
         return (b1,(c1,s1)))
      (\(i,p) -> fmap ((,) (Param.value initial i)) $ start p)
      (\p -> do
         (ctx,(nextParam,startParam)) <- createIOContext p
         return (ctx,
            (nextParam, (Param.get initial p, startParam))))
      deleteIOContext


takeWhile ::
   (Storable ph, MakeValueTuple ph pl, Memory.C pl pp, IsSized pp ps) =>
   (forall r. pl -> a -> CodeGenFunction r (Value Bool)) ->
   Param.T p ph ->
   T p a a
takeWhile check selectParam = simple
   (\p a () -> do
      Maybe.guard =<< Maybe.lift (check p a)
      return (a, ()))
   return
   selectParam
   (return ())


take ::
   Param.T p Int ->
   T p a a
take len =
   snd ^<<
   takeWhile (const $ A.cmp LLVM.CmpLT (valueOf 0) . fst) (return ()) <<<
   feedFst
      (Sig.iterate (const A.dec) (return ())
         ((fromIntegral :: Int -> Word32) . max 0 ^<< len))


{- |
The first output value is the initial value.
Thus 'integrate' delays by one sample compared with 'integrate0'.
-}
integrate ::
   (Storable a, A.Additive al,
    MakeValueTuple a al, Memory.C al as, IsSized as size) =>
   Param.T p a ->
   T p al al
integrate =
   flip loop (arr snd &&& zipWithSimple A.add)

integrate0 ::
   (Storable a, A.Additive al,
    MakeValueTuple a al, Memory.C al as, IsSized as size) =>
   Param.T p a ->
   T p al al
integrate0 =
   flip loop ((\a -> (a,a)) ^<< zipWithSimple A.add)