{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ForeignFunctionInterface #-}
module Synthesizer.LLVM.Simple.SignalPrivate where

import qualified LLVM.Extra.Memory as Memory
import qualified LLVM.Extra.MaybeContinuation as MaybeCont
import qualified LLVM.Extra.Arithmetic as A
import LLVM.Extra.Class (MakeValueTuple, ValueTuple, )

import LLVM.Util.Loop (Phi, )
import LLVM.Core (CodeGenFunction, )

import Control.Monad (liftM2, )
import Control.Applicative (Applicative, pure, liftA2, (<*>), )

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

import qualified Prelude as P


{-
We need the forall quantification for 'CodeGenFunction's @r@ parameter.
This type parameter will be unified with the result type of the final function.
Since one piece of code can be used in multiple functions
we cannot yet fix the type @r@ here.


We might avoid code duplication with Causal.Process by defining

> newtype T a = Cons (Causal.T () a)


In earlier versions the createIOContext method created only an ioContext
that was directly used to construct code for 'start' and 'next'.
This had the advantage that we did not need to pass
something via the Memory.C interface to the function.
However, creating both an ioContext and a low-level parameter has those advantages:
We can design Causal.Process such that a process
can be applied to multiple signals without recompilation.
We can lift simple signals and processes to their parameterized counterparts.
-}
data T a =
   forall state ioContext parameters.
      (Storable parameters,
       MakeValueTuple parameters,
       Memory.C (ValueTuple parameters),
       Memory.C state) =>
      Cons (forall r c.
            (Phi c) =>
            ValueTuple parameters ->
            state -> MaybeCont.T r c (a, state))
               -- compute next value
           (forall r.
            ValueTuple parameters ->
            CodeGenFunction r state)
               -- initial state
           (IO (ioContext, parameters))
               {- 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


data Core context initState exitState a =
   forall state.
      (Memory.C state) =>
      Core (forall r c.
            (Phi c) =>
            context ->
            state -> MaybeCont.T r c (a, state))
               -- compute next value
           (forall r.
            initState ->
            CodeGenFunction r state)
               -- initial state
           (state -> exitState)
               -- extract final state for cleanup


class Applicative signal => C signal where
   simple ::
      (Memory.C state) =>
      (forall r c. state -> MaybeCont.T r c (a, state)) ->
      (forall r. CodeGenFunction r state) ->
      signal a

   alter ::
      (forall context initState exitState.
          Core context initState exitState a0 ->
          Core context initState exitState a1) ->
      signal a0 -> signal a1

instance C T where
   simple next start =
      Cons
         (const next)
         (const start)
         (return ((),()))
         (const $ return ())

   alter f (Cons next0 start0 create delete) =
      case f (Core next0 start0 id) of
         Core next1 start1 _ ->
            Cons next1 start1 create delete


map ::
   (C signal) =>
   (forall r. a -> CodeGenFunction r b) -> signal a -> signal b
map f = alter (\(Core next start stop) ->
   Core
      (\ioContext sa0 -> do
         (a,sa1) <- next ioContext sa0
         b <- MaybeCont.lift $ f a
         return (b, sa1))
      start
      stop)

zipWith ::
   (C signal) =>
   (forall r. a -> b -> CodeGenFunction r c) ->
   signal a -> signal b -> signal c
zipWith f a b  =  map (uncurry f) $ liftA2 (,) a b


zip :: T a -> T b -> T (a,b)
zip (Cons nextA startA createIOContextA deleteIOContextA)
    (Cons nextB startB createIOContextB deleteIOContextB) =
   Cons
      (\(paramA, paramB) (sa0,sb0) -> do
         (a,sa1) <- nextA paramA sa0
         (b,sb1) <- nextB paramB sb0
         return ((a,b), (sa1,sb1)))
      (combineStart startA startB)
      (combineCreate createIOContextA createIOContextB)
      (combineDelete deleteIOContextA deleteIOContextB)

combineStart ::
   Monad m =>
   (paramA -> m stateA) ->
   (paramB -> m stateB) ->
   (paramA, paramB) -> m (stateA, stateB)
combineStart startA startB (paramA, paramB) =
   liftM2 (,)
      (startA paramA)
      (startB paramB)

combineCreate ::
   Monad m =>
   m (ioContextA, contextA) ->
   m (ioContextB, contextB) ->
   m ((ioContextA, ioContextB), (contextA, contextB))
combineCreate createIOContextA createIOContextB = do
   (ca,paramA) <- createIOContextA
   (cb,paramB) <- createIOContextB
   return ((ca,cb), (paramA,paramB))

combineDelete :: (Monad m) => (ca -> m ()) -> (cb -> m ()) -> (ca, cb) -> m ()
combineDelete deleteIOContextA deleteIOContextB (ca,cb) =
   deleteIOContextA ca >>
   deleteIOContextB cb


instance Functor T where
   fmap f = map (return . f)

{- |
ZipList semantics
-}
instance Applicative T where
   pure x = simple (\() -> return (x, ())) (return ())
   f <*> a = fmap (uncurry ($)) $ zip f a

instance (A.Additive a) => Additive.C (T a) where
   zero = pure A.zero
   negate = map A.neg
   (+) = zipWith A.add
   (-) = zipWith A.sub

instance (A.PseudoRing a, A.IntegerConstant a) => Ring.C (T a) where
   one = pure A.one
   fromInteger n = pure (A.fromInteger' n)
   (*) = zipWith A.mul

instance (A.Field a, A.RationalConstant a) => Field.C (T a) where
   fromRational' x = pure (A.fromRational' $ Ratio.toRational98 x)
   (/) = zipWith A.fdiv


instance (A.PseudoRing a, A.Real a, A.IntegerConstant a) => P.Num (T a) where
   fromInteger n = pure (A.fromInteger' n)
   negate = map A.neg
   (+) = zipWith A.add
   (-) = zipWith A.sub
   (*) = zipWith A.mul
   abs = map A.abs
   signum = map A.signum

instance (A.Field a, A.Real a, A.RationalConstant a) => P.Fractional (T a) where
   fromRational x = pure (A.fromRational' x)
   (/) = zipWith A.fdiv