{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
module Synthesizer.LLVM.CausalParameterized.Functional (
   T,
   lift,
   ($&), (&|&),
   compile,
   withArgs, MakeArguments, Arguments, makeArgs,
   ) where

import qualified Synthesizer.LLVM.CausalParameterized.Process as Causal
import qualified Synthesizer.LLVM.Frame.Stereo as Stereo

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

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

import qualified Number.Ratio as Ratio
import qualified Algebra.Transcendental as Trans
import qualified Algebra.Algebraic as Algebraic
import qualified Algebra.Field as Field
import qualified Algebra.Ring as Ring
import qualified Algebra.Additive as Additive

import qualified Control.Monad.Trans.State as State
import qualified Control.Monad.Trans.Class as Trans
import Control.Monad.Trans.State (StateT, )

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

import Foreign.Storable (Storable, )

import Data.Tuple.HT (fst3, snd3, thd3, )

import qualified System.Unsafe as Unsafe

import Prelude hiding ((.), )


newtype T p inp out = Cons (Code p inp out)


-- | similar to @Causal.T p a b@
data Code 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) =>
   Code
      (forall r c.
       (Phi c) =>
       ValueTuple nextParamTuple ->
       a -> state ->
       StateT Vault (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



instance Category (Code p) where
   id = arr id
   Code nextB startB createIOContextB deleteIOContextB .
      Code nextA startA createIOContextA deleteIOContextA = Code
         (\(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)

instance Arrow (Code p) where
   arr f = Code
      (\ _p a state -> return (f a, state))
      (const $ return ())
      (const $ return ((),((),())))
      (const $ return ())
   first (Code next start create delete) = Code
      (\ioContext (b,d) sa0 ->
         do (c,sa1) <- next ioContext b sa0
            return ((c,d), sa1))
      start create delete


{-
We must not define Category and Arrow instances
because in osci***osci the result of osci would be shared,
although it depends on the particular input.

instance Category (T p) where
   id = tagUnique Cat.id
   Cons a . Cons b = tagUnique (a . b)

instance Arrow (T p) where
   arr f = tagUnique $ arr f
   first (Cons a) = tagUnique $ first a
-}

instance Functor (T p inp) where
   fmap f (Cons x) =
      tagUnique $ x >>^ f

instance Applicative (T p inp) where
   pure a = tagUnique $ arr (const a)
   f <*> x = fmap (uncurry ($))  $  f &|& x


lift0 :: (forall r. CodeGenFunction r out) -> T p inp out
lift0 f = lift (Causal.mapSimple (const f))

lift1 :: (forall r. a -> CodeGenFunction r out) -> T p inp a -> T p inp out
lift1 f x = Causal.mapSimple f $& x

lift2 :: (forall r. a -> b -> CodeGenFunction r out) -> T p inp a -> T p inp b -> T p inp out
lift2 f x y = Causal.zipWithSimple f $& x&|&y


instance (A.PseudoRing b, A.Real b, A.IntegerConstant b) => Num (T p a b) where
   fromInteger n = pure (A.fromInteger' n)
   (+) = lift2 A.add
   (-) = lift2 A.sub
   (*) = lift2 A.mul
   abs = lift1 A.abs
   signum = lift1 A.signum


instance (A.Additive b) => Additive.C (T p a b) where
   zero = pure A.zero
   (+) = lift2 A.add
   (-) = lift2 A.sub
   negate = lift1 A.neg

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

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

instance (A.Transcendental b, A.RationalConstant b) => Algebraic.C (T p a b) where
   sqrt = lift1 A.sqrt
   root n x = lift2 A.pow x (Field.recip $ Ring.fromInteger n)
   x^/r = lift2 A.pow x (Field.fromRational' r)

instance (A.Transcendental b, A.RationalConstant b) => Trans.C (T p a b) where
   pi = lift0 A.pi
   sin = lift1 A.sin
   cos = lift1 A.cos
   (**) = lift2 A.pow
   exp = lift1 A.exp
   log = lift1 A.log

   asin _ = error "LLVM missing intrinsic: asin"
   acos _ = error "LLVM missing intrinsic: acos"
   atan _ = error "LLVM missing intrinsic: atan"


infixr 0 $&

($&) :: Causal.T p b c -> T p a b -> T p a c
f $& (Cons b) =
   tagUnique $  liftCode f . b


infixr 3 &|&

(&|&) :: T p a b -> T p a c -> T p a (b,c)
Cons b &|& Cons c =
   tagUnique $  b &&& c


liftCode :: Causal.T p inp out -> Code p inp out
liftCode (Causal.Cons next start create delete) =
   Code
      (\p a state -> Trans.lift (next p a state))
      start create delete

lift :: Causal.T p inp out -> T p inp out
lift = tagUnique . liftCode

tag :: Vault.Key out -> Code p inp out -> T p inp out
tag key (Code next start create delete) =
   Cons $
   Code
      (\p a s0 -> do
         mb <- State.gets (Vault.lookup key)
         case mb of
            Just b -> return (b,s0)
            Nothing -> do
               bs@(b,_) <- next p a s0
               State.modify (Vault.insert key b)
               return bs)
      start create delete

-- dummy for debugging
_tag :: Vault.Key out -> Code p inp out -> T p inp out
_tag _ = Cons

tagUnique :: Code p inp out -> T p inp out
tagUnique code =
   Unsafe.performIO $
   fmap (flip tag code) Vault.newKey

initialize :: Code p inp out -> Causal.T p inp out
initialize (Code next start create delete) =
   Causal.Cons
      (\p a state -> State.evalStateT (next p a state) Vault.empty)
      start create delete

compile :: T p inp out -> Causal.T p inp out
compile (Cons code) = initialize code


{- |
Using 'withArgs' you can simplify

> let x = F.lift (arr fst)
>     y = F.lift (arr (fst.snd))
>     z = F.lift (arr (snd.snd))
> in  F.compile (f x y z)

to

> withArgs $ \(x,(y,z)) -> f x y z
-}
withArgs ::
   (MakeArguments inp) =>
   (Arguments (T p inp) inp -> T p inp out) -> Causal.T p inp out
withArgs = withArgsStart (lift Cat.id)

withArgsStart ::
   (MakeArguments inp) =>
   T p inp inp -> (Arguments (T p inp) inp -> T p inp out) -> Causal.T p inp out
withArgsStart fid f = compile (f (makeArgs fid))


type family Arguments (f :: * -> *) (arg :: *)

class MakeArguments arg where
   makeArgs :: Functor f => f arg -> Arguments f arg


{-
I have thought about an Arg type, that marks where to stop descending.
This way we can throw away all of these FlexibleContext instances
and the user can freely choose the granularity of arguments.
However this does not work so easily,
because we would need a functional depedency from, say,
@(Arg a, Arg b)@ to @(a,b)@.
This is the opposite direction to the dependency we use currently.
-}
type instance Arguments f (LLVM.Value a) = f (LLVM.Value a)
instance MakeArguments (LLVM.Value a) where
   makeArgs = id

type instance Arguments f (Stereo.T a) = f (Stereo.T a)
instance MakeArguments (Stereo.T a) where
   makeArgs = id

type instance Arguments f (a,b) = (Arguments f a, Arguments f b)
instance (MakeArguments a, MakeArguments b) =>
      MakeArguments (a,b) where
   makeArgs f = (makeArgs $ fmap fst f, makeArgs $ fmap snd f)

type instance Arguments f (a,b,c) = (Arguments f a, Arguments f b, Arguments f c)
instance (MakeArguments a, MakeArguments b, MakeArguments c) =>
      MakeArguments (a,b,c) where
   makeArgs f = (makeArgs $ fmap fst3 f, makeArgs $ fmap snd3 f, makeArgs $ fmap thd3 f)