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

import qualified Synthesizer.LLVM.CausalParameterized.ProcessPrivate as CausalP
import qualified Synthesizer.LLVM.Parameterized.Signal as Signal
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.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 @CausalP.T p a b@
data Code p a b =
   forall context state ioContext parameters.
      (Storable parameters,
       MakeValueTuple parameters,
       Memory.C (ValueTuple parameters),
       Memory.C context,
       Memory.C state) =>
   Code
      (forall r c.
       (Phi c) =>
       context -> a -> state ->
       StateT Vault (Maybe.T r c) (b, state))
          -- compute next value
      (forall r.
       ValueTuple parameters ->
       CodeGenFunction r (context, state))
          -- initial state
      (forall r.
       context -> state ->
       CodeGenFunction r ())
          -- cleanup
      (p -> 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



instance Category (Code p) where
   id = arr id
   Code nextB startB stopB createIOContextB deleteIOContextB .
      Code nextA startA stopA createIOContextA deleteIOContextA = Code
         (CausalP.composeNext
             (State.mapStateT . Maybe.onFail)
             stopA stopB nextA nextB)
         (CausalP.composeStart startA startB)
         (CausalP.composeStop stopA stopB)
         (CausalP.composeCreate createIOContextA createIOContextB)
         (CausalP.composeDelete deleteIOContextA deleteIOContextB)


instance Arrow (Code p) where
   arr f = Code
      (\ _p a () -> return (f a, ()))
      (\() -> return ((),()))
      (\() () -> return ())
      (const $ return ((),()))
      (const $ return ())
   first (Code next start stop create delete) = Code
      (CausalP.firstNext next) start stop
      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 (CausalP.mapSimple (const f))

lift1 :: (forall r. a -> CodeGenFunction r out) -> T p inp a -> T p inp out
lift1 f x = CausalP.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 = CausalP.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.Field b, A.Real b, A.RationalConstant b) => Fractional (T p a b) where
   fromRational x = pure (A.fromRational' x)
   (/) = lift2 A.fdiv


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 $&

($&) :: CausalP.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 :: CausalP.T p inp out -> Code p inp out
liftCode (CausalP.Cons next start stop create delete) =
   Code
      (\p a state -> Trans.lift (next p a state))
      start stop create delete

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

fromSignal :: Signal.T p out -> T p inp out
fromSignal = lift . CausalP.fromSignal

tag :: Vault.Key out -> Code p inp out -> T p inp out
tag key (Code next start stop 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 stop 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 -> CausalP.T p inp out
initialize (Code next start stop create delete) =
   CausalP.Cons
      (\p a state -> State.evalStateT (next p a state) Vault.empty)
      start stop create delete

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

compileSignal :: T p () out -> Signal.T p out
compileSignal f = CausalP.toSignal $ compile f


{- |
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) -> CausalP.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) -> CausalP.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 () = f ()
instance MakeArguments () 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)


newtype AnyArg a = AnyArg {getAnyArg :: a}

type instance Arguments f (AnyArg a) = f a
instance MakeArguments (AnyArg a) where
   makeArgs = fmap getAnyArg