{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleContexts #-}
module Synthesizer.LLVM.CausalParameterized.FunctionalPlug (
   T,
   ($&), (&|&),
   run, runPlugOut,
   fromSignal, plug, askParameter, Input,
   withArgs, withArgsPlugOut,
   MakeArguments, Arguments, makeArgs,
   ) where

import qualified Synthesizer.LLVM.Plug.Input as PIn
import qualified Synthesizer.LLVM.Plug.Output as POut

import qualified Synthesizer.LLVM.CausalParameterized.Process as CausalP
import qualified Synthesizer.LLVM.Parameterized.Signal as Sig

import qualified Synthesizer.CausalIO.Process as PIO
import qualified Synthesizer.Generic.Cut as CutG
import qualified Synthesizer.Zip as Zip

import qualified Data.EventList.Relative.BodyTime as EventListBT
import qualified Data.StorableVector as SV

import qualified LLVM.Extra.Arithmetic as A
import LLVM.Core (CodeGenFunction, )

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.Reader as MR
import qualified Control.Monad.Trans.State as MS

import qualified Data.Set as Set
import qualified Data.Vault.Lazy as Vault
import Data.Vault.Lazy (Vault, )
import Data.Unique (Unique, newUnique, )
import Data.Maybe (fromMaybe, )

import Control.Arrow ((^<<), (<<^), arr, first, second, )
import Control.Category (id, (.), )
import Control.Applicative (Applicative, (<*>), pure, liftA2, liftA3, )

import qualified System.Unsafe as Unsafe

import Prelude hiding (id, (.), )


{- |
This data type detects sharing.
-}
{-
There are two levels of the use of keys.
At the top level, in T's State monad,
we store an object id in order to check,
whether we have already seen a certain object.
If we encounter a known object
then we use the Simple constructor
and fetch the stored CausalP output
within the causal process enclosed in Simple.
This and the causal process in the Plugged constructor
are the second level.
These arrows handle a Vault like a state monad
and insert all values they produce into the Vault.
-}
newtype T pp pl inp out =
   Cons (MS.State (Set.Set Unique) (Core pp pl inp out))

{-
We need to hide the x and y types
since these types grow when combining Cores,
and then we could not define numeric instances.
-}
data Core pp pl inp out =
   forall x y. CutG.Read x =>
   Plugged
      (pp -> inp -> x)
      (PIn.T x y)
      (CausalP.T pl (y, Vault) (out, Vault))
   |
   {-
   The Simple constructor is needed for reusing shared CausalP processes
   and for input without external representation. (a Plug.Input)
   -}
   Simple (CausalP.T pl Vault (out, Vault))


applyCore ::
   CausalP.T pl (a, Vault) (b, Vault) ->
   Core pp pl inp a ->
   Core pp pl inp b
applyCore f core =
   case core of
      Plugged prep plg process -> Plugged prep plg (f . process)
      Simple process -> Simple (f . process)

combineCore ::
   Core pp pl inp a ->
   Core pp pl inp b ->
   Core pp pl inp (a,b)
combineCore (Plugged prepA plugA processA) (Plugged prepB plugB processB) =
   Plugged
      (\p -> Zip.arrowFanout (prepA p) (prepB p))
      (PIn.split plugA plugB)
      ((\(a,(b,v)) -> ((a,b), v)) ^<< second processB
       . arr (\((a,v),b) -> (a,(b,v))) .
       first processA <<^ (\((a,b),v) -> ((a,v),b)))
combineCore (Simple processA) (Plugged prepB plugB processB) =
   Plugged prepB plugB
      ((\(b,(a,v)) -> ((a,b), v)) ^<< second processA . processB)
combineCore (Plugged prepA plugA processA) (Simple processB) =
   Plugged prepA plugA
      ((\(a,(b,v)) -> ((a,b), v)) ^<< second processB . processA)
combineCore (Simple processA) (Simple processB) =
   Simple ((\(a,(b,v)) -> ((a,b), v)) ^<< second processB . processA)


reuseCore :: Vault.Key out -> Core pp pl inp out
reuseCore key =
   Simple $ arr $ \vault ->
      (fromMaybe (error "key must have been lost") $ Vault.lookup key vault,
       vault)


tag ::
   Unique -> Vault.Key out ->
   MS.State (Set.Set Unique) (Core pp pl inp out) ->
   T pp pl inp out
tag unique key stateCore = Cons $ do
   alreadySeen <- MS.gets (Set.member unique)
   if alreadySeen
      then return $ reuseCore key
      else do
         MS.modify (Set.insert unique)
         fmap (applyCore (arr $ \(a,v) -> (a, Vault.insert key a v))) stateCore

tagUnique ::
   MS.State (Set.Set Unique) (Core pp pl inp out) ->
   T pp pl inp out
tagUnique core =
   Unsafe.performIO $
   liftA3 tag newUnique Vault.newKey (pure core)


infixr 0 $&

($&) ::
   CausalP.T pl a b ->
   T pp pl inp a ->
   T pp pl inp b
f  $&  Cons core =
   tagUnique $ fmap (applyCore $ first f) core


infixr 3 &|&

(&|&) ::
   T pp pl inp a ->
   T pp pl inp b ->
   T pp pl inp (a,b)
Cons coreA  &|&  Cons coreB =
   tagUnique $ liftA2 combineCore coreA coreB


instance Functor (Core pp pl inp) where
   fmap f (Simple process) = Simple (fmap (first f) process)
   fmap f (Plugged prep plg process) = Plugged prep plg (fmap (first f) process)

instance Applicative (Core pp pl inp) where
   pure a = lift0Core $ pure a
   f <*> x = fmap (uncurry ($))  $  combineCore f x

lift0Core :: (forall r. CodeGenFunction r out) -> Core pp pl inp out
lift0Core f = Simple (CausalP.mapSimple (\v -> fmap (flip (,) v) f))


instance Functor (T pp pl inp) where
   fmap f (Cons x) = tagUnique $ fmap (fmap f) x

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


lift0 :: (forall r. CodeGenFunction r out) -> T pp pl inp out
lift0 f = tagUnique $ pure $ lift0Core f

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

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


instance (A.PseudoRing b, A.Real b, A.IntegerConstant b) => Num (T pp pl 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 pp pl a b) where
   fromRational x = pure (A.fromRational' x)
   (/) = lift2 A.fdiv


instance (A.Additive b) => Additive.C (T pp pl 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 pp pl 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 pp pl a b) where
   fromRational' x = pure (A.fromRational' $ Ratio.toRational98 x)
   (/) = lift2 A.fdiv

instance (A.Transcendental b, A.RationalConstant b) => Algebraic.C (T pp pl 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 pp pl 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"



fromSignal ::
   Sig.T pl a ->
   T pp pl inp a
fromSignal sig =
   tagUnique $ pure $ Simple (CausalP.feedFst sig)



type Input pp a = MR.Reader (pp, a)

plug ::
   (CutG.Read b, PIn.Default b) =>
   Input pp a b ->
   T pp pl a (PIn.Element b)
plug accessor =
   tagUnique $ pure $
   Plugged
      (curry $ MR.runReader accessor)
      PIn.deflt
      id

askParameter :: Input pp a pp
askParameter = MR.asks fst


runPlugOut ::
   T pp pl a x -> POut.T x b ->
   IO (pp -> pl -> PIO.T a b)
runPlugOut (Cons core) pout =
   case MS.evalState core Set.empty of
      Plugged prep pin process ->
         fmap (\f pp pl -> f pl <<^ prep pp) $
         CausalP.processIOCore
            pin
            (fst ^<< process <<^ flip (,) Vault.empty)
            pout
      Simple _ ->
         error "FunctionalPlug.runPlugOut: no substantial input available"
         -- Simple process ->
         --    CausalP.processIOCore pin process pout

run ::
   (POut.Default b) =>
   T pp pl a (POut.Element b) ->
   IO (pp -> pl -> PIO.T a b)
run f = runPlugOut f POut.deflt


{- |
Cf. 'F.withArgs'.
-}
withArgs ::
   (MakeArguments a, POut.Default b) =>
   (Arguments (Input pp a) a -> T pp pl a (POut.Element b)) ->
   IO (pp -> pl -> PIO.T a b)
withArgs f = withArgsPlugOut f POut.deflt

withArgsPlugOut ::
   (MakeArguments a) =>
   (Arguments (Input pp a) a -> T pp pl a x) ->
   POut.T x b ->
   IO (pp -> pl -> PIO.T a b)
withArgsPlugOut = withArgsPlugOutStart (MR.asks snd)

withArgsPlugOutStart ::
   (MakeArguments a) =>
   Input pp a a ->
   (Arguments (Input pp a) a -> T pp pl a x) ->
   POut.T x b ->
   IO (pp -> pl -> PIO.T a b)
withArgsPlugOutStart fid f = runPlugOut (f (makeArgs fid))



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

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


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

type instance Arguments f (SV.Vector a) = f (SV.Vector a)
instance MakeArguments (SV.Vector a) where
   makeArgs = id

type instance Arguments f (Zip.T a b) = (Arguments f a, Arguments f b)
instance (MakeArguments a, MakeArguments b) =>
      MakeArguments (Zip.T a b) where
   makeArgs f = (makeArgs $ fmap Zip.first f, makeArgs $ fmap Zip.second f)