{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# 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, ValueTuple, )
import LLVM.Util.Loop (Phi, )
import LLVM.Core (Value, valueOf, 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, )

import qualified Prelude as P


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


simple ::
   (Storable startParamTuple,
    Storable nextParamTuple,
    MakeValueTuple startParamTuple, ValueTuple startParamTuple ~ startParamValue,
    MakeValueTuple nextParamTuple, ValueTuple nextParamTuple ~ nextParamValue,
    Memory.C startParamValue,
    Memory.C nextParamValue,
    Memory.C state) =>
   (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, ValueTuple pnh ~ pnl, Memory.C pnl,
    Storable psh, MakeValueTuple psh, ValueTuple psh ~ psl, Memory.C psl,
    Memory.C s) =>
   (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, ValueTuple ph ~ pl, Memory.C pl) =>
   (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


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

instance (A.Field b, A.Real b, A.RationalConstant b) => P.Fractional (T p a b) where
   fromRational x = pure (A.fromRational' 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, ValueTuple ch ~ c,
    Memory.C c) =>
   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, ValueTuple ph ~ pl, Memory.C pl) =>
   (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, ValueTuple a ~ al, Memory.C al) =>
   Param.T p a ->
   T p al al
integrate =
   flip loop (arr snd &&& zipWithSimple A.add)

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