{-# LANGUAGE TypeOperators #-}
{-# OPTIONS -fplugin=Rattus.Plugin #-}
{-# LANGUAGE Arrows #-}
{-# LANGUAGE RebindableSyntax #-}

-- | A simply Yampa-like library for signal functions.

module Rattus.Yampa (
  SF,
  identity,
  compose,
  switch,
  rSwitch,
  constant,
  loopPre,
  stepSF,
  initially,
  integral,
  (-->),
  (-:>),
  (>--),
  (-=>),
  (>=-),
  (^>>),
  (>>^),
  (<<^),
  (^<<),
  module Rattus.Arrow, (>>>)) where

import Rattus.Primitives
import Rattus.Plugin
import Rattus.Strict

import GHC.Float
import Control.Category
import Rattus.Arrow
import Prelude hiding (id)

import Data.VectorSpace

-- all functions in this module are in Rattus 
{-# ANN module Rattus #-}

type DTime = Double 

-- | Signal functions from inputs of type @a@ to outputs of type @b@.
data SF a b = SF{
  -- | Run a signal function for one step.
  forall a b. SF a b -> DTime -> a -> O (SF a b) :* b
stepSF :: !(DTime -> a -> (O(SF a b) :* b))}

-- | The identity signal function that does nothing.
identity :: SF a a
identity :: forall a. SF a a
identity = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF (\ DTime
_ a
x -> (forall a. a -> O a
delay forall a. SF a a
identity forall a b. a -> b -> a :* b
:* a
x))

-- | Compose two signal functions.
compose :: SF b c -> SF a b -> SF a c
compose :: forall b c a. SF b c -> SF a b -> SF a c
compose (SF DTime -> b -> O (SF b c) :* c
sf2) (SF DTime -> a -> O (SF a b) :* b
sf1) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF DTime -> a -> O (SF a c) :* c
sf
  where sf :: DTime -> a -> O (SF a c) :* c
sf DTime
d a
a = let (O (SF a b)
r1 :* b
b) = DTime -> a -> O (SF a b) :* b
sf1 DTime
d a
a
                     (O (SF b c)
r2 :* c
c) = DTime -> b -> O (SF b c) :* c
sf2 DTime
d b
b
                 in (forall a. a -> O a
delay (forall b c a. SF b c -> SF a b -> SF a c
compose (forall a. O a -> a
adv O (SF b c)
r2) (forall a. O a -> a
adv O (SF a b)
r1)) forall a b. a -> b -> a :* b
:* c
c)

-- | Compute the integral of a signal. The first argument is the
-- offset.
integral :: (Stable a, VectorSpace a s, Fractional s) => a -> SF a a
integral :: forall a s.
(Stable a, VectorSpace a s, Fractional s) =>
a -> SF a a
integral a
acc = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF forall {s} {a} {p}.
(VectorSpace a s, Fractional s, VectorSpace a a, Real p,
 Fractional a) =>
p -> a -> O (SF a a) :* a
sf'
  where sf' :: p -> a -> O (SF a a) :* a
sf' p
t a
a = let acc' :: a
acc' = a
acc forall v a. VectorSpace v a => v -> v -> v
^+^ (forall a b. (Real a, Fractional b) => a -> b
realToFrac p
t forall v a. VectorSpace v a => a -> v -> v
*^ a
a)
                  in (forall a. a -> O a
delay (forall a s.
(Stable a, VectorSpace a s, Fractional s) =>
a -> SF a a
integral a
acc') forall a b. a -> b -> a :* b
:* a
acc')

-- | @switch s f@ behaves like @s@ composed with @arr fst@ until @s@
-- produces a value of the form @Just' c@ in the second
-- component. From then on it behaves like $f c@.
switch :: SF a (b :* Maybe' c) -> Box (c -> SF a b) -> SF a b
switch :: forall a b c. SF a (b :* Maybe' c) -> Box (c -> SF a b) -> SF a b
switch (SF DTime -> a -> O (SF a (b :* Maybe' c)) :* (b :* Maybe' c)
sf) Box (c -> SF a b)
f = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF DTime -> a -> O (SF a b) :* b
sf'
  where sf' :: DTime -> a -> O (SF a b) :* b
sf' DTime
t a
a = let (O (SF a (b :* Maybe' c))
nxt :* (b
b :* Maybe' c
c')) =  DTime -> a -> O (SF a (b :* Maybe' c)) :* (b :* Maybe' c)
sf DTime
t a
a
                  in case Maybe' c
c' of
                       Just' c
c -> forall a b. SF a b -> DTime -> a -> O (SF a b) :* b
stepSF (forall a. Box a -> a
unbox Box (c -> SF a b)
f c
c) DTime
t a
a
                       Maybe' c
Nothing' -> (forall a. a -> O a
delay (forall a b c. SF a (b :* Maybe' c) -> Box (c -> SF a b) -> SF a b
switch (forall a. O a -> a
adv O (SF a (b :* Maybe' c))
nxt) Box (c -> SF a b)
f)forall a b. a -> b -> a :* b
:* b
b)

-- | @rSwitch s@ behaves like @s@, but every time the second input is
-- of the form @Just' s'@ it will change behaviour to @s'@.
rSwitch :: SF a b -> SF (a :* Maybe' (SF a b)) b
rSwitch :: forall a b. SF a b -> SF (a :* Maybe' (SF a b)) b
rSwitch (SF DTime -> a -> O (SF a b) :* b
sf) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF DTime
-> (a :* Maybe' (SF a b)) -> O (SF (a :* Maybe' (SF a b)) b) :* b
sf'
  where sf' :: DTime
-> (a :* Maybe' (SF a b)) -> O (SF (a :* Maybe' (SF a b)) b) :* b
sf' DTime
t (a
a :* Maybe' (SF a b)
m) = case Maybe' (SF a b)
m of
                        Just' (SF DTime -> a -> O (SF a b) :* b
newSf) ->
                           let (O (SF a b)
nxt :* b
b) = DTime -> a -> O (SF a b) :* b
newSf DTime
t a
a
                           in (forall a. a -> O a
delay (forall a b. SF a b -> SF (a :* Maybe' (SF a b)) b
rSwitch (forall a. O a -> a
adv O (SF a b)
nxt)) forall a b. a -> b -> a :* b
:* b
b)
                        Maybe' (SF a b)
Nothing' -> let (O (SF a b)
nxt :* b
b) = DTime -> a -> O (SF a b) :* b
sf DTime
t a
a
                                   in (forall a. a -> O a
delay (forall a b. SF a b -> SF (a :* Maybe' (SF a b)) b
rSwitch (forall a. O a -> a
adv O (SF a b)
nxt)) forall a b. a -> b -> a :* b
:* b
b)

-- | Constant signal function.
constant :: Stable b => b -> SF a b
constant :: forall b a. Stable b => b -> SF a b
constant b
x = forall {a}. SF a b
run
  where run :: SF a b
run = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF (\ DTime
_ a
_ -> (forall a. a -> O a
delay SF a b
run forall a b. a -> b -> a :* b
:* b
x))

-- | The output at time zero is the first argument, and from that
-- point on it behaves like the signal function passed as second
-- argument.
(-->) :: b -> SF a b -> SF a b
b
b --> :: forall b a. b -> SF a b -> SF a b
--> (SF DTime -> a -> O (SF a b) :* b
sf) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF DTime -> a -> O (SF a b) :* b
sf'
  where sf' :: DTime -> a -> O (SF a b) :* b
sf' DTime
d a
x = (forall a b. (a :* b) -> a
fst' (DTime -> a -> O (SF a b) :* b
sf DTime
d a
x) forall a b. a -> b -> a :* b
:* b
b)

-- | Insert a sample in the output, and from that point on, behave
-- like the given signal function.
--
-- Note: The type of -:> is different from Yampa's: The second
-- argument must be delayed (or boxed).
(-:>) :: b -> O (SF a b) -> SF a b
b
b -:> :: forall b a. b -> O (SF a b) -> SF a b
-:> O (SF a b)
sf = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF forall {p} {p}. p -> p -> O (SF a b) :* b
sf'
  where sf' :: p -> p -> O (SF a b) :* b
sf' p
_d p
_x = (O (SF a b)
sf forall a b. a -> b -> a :* b
:* b
b)

-- | The input at time zero is the first argument, and from that point
-- on it behaves like the signal function passed as second argument.
(>--) :: a -> SF a b -> SF a b
a
a >-- :: forall a b. a -> SF a b -> SF a b
>-- (SF DTime -> a -> O (SF a b) :* b
sf) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF forall {p}. DTime -> p -> O (SF a b) :* b
sf'
  where sf' :: DTime -> p -> O (SF a b) :* b
sf' DTime
d p
_a = DTime -> a -> O (SF a b) :* b
sf DTime
d a
a


-- | Apply a function to the first output value at time
-- zero.
(-=>) :: (b -> b) -> SF a b -> SF a b
b -> b
f -=> :: forall b a. (b -> b) -> SF a b -> SF a b
-=> (SF DTime -> a -> O (SF a b) :* b
sf) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF DTime -> a -> O (SF a b) :* b
sf'
  where sf' :: DTime -> a -> O (SF a b) :* b
sf' DTime
d a
a = let (O (SF a b)
r:*b
b) = DTime -> a -> O (SF a b) :* b
sf DTime
d a
a
                  in (O (SF a b)
rforall a b. a -> b -> a :* b
:*b -> b
f b
b)
                     
-- | Apply a function to the first input value at time
-- zero.
(>=-) :: (a -> a) -> SF a b -> SF a b
a -> a
f >=- :: forall a b. (a -> a) -> SF a b -> SF a b
>=- (SF DTime -> a -> O (SF a b) :* b
sf) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF DTime -> a -> O (SF a b) :* b
sf'
  where sf' :: DTime -> a -> O (SF a b) :* b
sf' DTime
d a
a = DTime -> a -> O (SF a b) :* b
sf DTime
d (a -> a
f a
a)

-- | Override initial value of input signal.

initially :: a -> SF a a
initially :: forall a. a -> SF a a
initially = (forall b a. b -> SF a b -> SF a b
--> forall a. SF a a
identity)

-- | Lift a function to a signal function (applied pointwise).
-- 
-- Note: The type of is different from Yampa's: The function argument
-- must be boxed.
arrPrim :: Box (a -> b) -> SF a b
arrPrim :: forall a b. Box (a -> b) -> SF a b
arrPrim Box (a -> b)
f = SF a b
run where
  run :: SF a b
run = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF (\ DTime
_d a
a -> (forall a. a -> O a
delay SF a b
runforall a b. a -> b -> a :* b
:* forall a. Box a -> a
unbox Box (a -> b)
f a
a ))


{-# ANN firstPrim AllowLazyData #-}
-- | Apply a signal function to the first component.
firstPrim :: SF a b -> SF (a,c) (b,c)
firstPrim :: forall a b c. SF a b -> SF (a, c) (b, c)
firstPrim (SF DTime -> a -> O (SF a b) :* b
sf) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF forall {b} {c}. DTime -> (a, b) -> O (SF (a, c) (b, c)) :* (b, b)
sf'
  where sf' :: DTime -> (a, b) -> O (SF (a, c) (b, c)) :* (b, b)
sf' DTime
d (a
a,b
c) = let (O (SF a b)
r:* b
b) = DTime -> a -> O (SF a b) :* b
sf DTime
d a
a
                       in (forall a. a -> O a
delay (forall a b c. SF a b -> SF (a, c) (b, c)
firstPrim (forall a. O a -> a
adv O (SF a b)
r))forall a b. a -> b -> a :* b
:* (b
b,b
c))

{-# ANN secondPrim AllowLazyData #-}
-- | Apply a signal function to the second component.
secondPrim :: SF a b -> SF (c,a) (c,b)
secondPrim :: forall a b c. SF a b -> SF (c, a) (c, b)
secondPrim (SF DTime -> a -> O (SF a b) :* b
sf) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF forall {a} {c}. DTime -> (a, a) -> O (SF (c, a) (c, b)) :* (a, b)
sf'
  where sf' :: DTime -> (a, a) -> O (SF (c, a) (c, b)) :* (a, b)
sf' DTime
d (a
c,a
a) = let (O (SF a b)
r:* b
b) = DTime -> a -> O (SF a b) :* b
sf DTime
d a
a
                       in (forall a. a -> O a
delay (forall a b c. SF a b -> SF (c, a) (c, b)
secondPrim (forall a. O a -> a
adv O (SF a b)
r))forall a b. a -> b -> a :* b
:* (a
c,b
b))

{-# ANN parSplitPrim AllowLazyData #-}
-- | Apply two signal functions in parallel.
parSplitPrim :: SF a b -> SF c d  -> SF (a,c) (b,d)
parSplitPrim :: forall a b c d. SF a b -> SF c d -> SF (a, c) (b, d)
parSplitPrim (SF DTime -> a -> O (SF a b) :* b
sf1) (SF DTime -> c -> O (SF c d) :* d
sf2) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF DTime -> (a, c) -> O (SF (a, c) (b, d)) :* (b, d)
sf'
  where sf' :: DTime -> (a, c) -> O (SF (a, c) (b, d)) :* (b, d)
sf' DTime
dt (a
a,c
c) = let (O (SF a b)
r1:* b
b) = DTime -> a -> O (SF a b) :* b
sf1 DTime
dt a
a
                           (O (SF c d)
r2:* d
d) = DTime -> c -> O (SF c d) :* d
sf2 DTime
dt c
c
                       in (forall a. a -> O a
delay (forall a b c d. SF a b -> SF c d -> SF (a, c) (b, d)
parSplitPrim (forall a. O a -> a
adv O (SF a b)
r1) (forall a. O a -> a
adv O (SF c d)
r2))forall a b. a -> b -> a :* b
:* (b
b,d
d))

{-# ANN parFanOutPrim AllowLazyData #-}
-- | Apply two signal functions in parallel on the same input.
parFanOutPrim :: SF a b -> SF a c -> SF a (b, c)
parFanOutPrim :: forall a b c. SF a b -> SF a c -> SF a (b, c)
parFanOutPrim (SF DTime -> a -> O (SF a b) :* b
sf1) (SF DTime -> a -> O (SF a c) :* c
sf2) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF DTime -> a -> O (SF a (b, c)) :* (b, c)
sf'
  where sf' :: DTime -> a -> O (SF a (b, c)) :* (b, c)
sf' DTime
dt a
a = let (O (SF a b)
r1:* b
b) = DTime -> a -> O (SF a b) :* b
sf1 DTime
dt a
a
                       (O (SF a c)
r2:* c
c) = DTime -> a -> O (SF a c) :* c
sf2 DTime
dt a
a
                   in (forall a. a -> O a
delay (forall a b c. SF a b -> SF a c -> SF a (b, c)
parFanOutPrim (forall a. O a -> a
adv O (SF a b)
r1) (forall a. O a -> a
adv O (SF a c)
r2))forall a b. a -> b -> a :* b
:* (b
b,c
c))

instance Category SF where
  id :: forall a. SF a a
id = forall a. SF a a
identity
  . :: forall b c a. SF b c -> SF a b -> SF a c
(.) = forall b c a. SF b c -> SF a b -> SF a c
compose

instance Arrow SF where
  arrBox :: forall a b. Box (a -> b) -> SF a b
arrBox = forall a b. Box (a -> b) -> SF a b
arrPrim
  first :: forall a b c. SF a b -> SF (a, c) (b, c)
first = forall a b c. SF a b -> SF (a, c) (b, c)
firstPrim
  second :: forall a b c. SF a b -> SF (c, a) (c, b)
second = forall a b c. SF a b -> SF (c, a) (c, b)
secondPrim
  *** :: forall a b c d. SF a b -> SF c d -> SF (a, c) (b, d)
(***) = forall a b c d. SF a b -> SF c d -> SF (a, c) (b, d)
parSplitPrim
  &&& :: forall a b c. SF a b -> SF a c -> SF a (b, c)
(&&&) = forall a b c. SF a b -> SF a c -> SF a (b, c)
parFanOutPrim


-- | Loop with an initial value for the signal being fed back.
-- 
-- Note: The type of @loopPre@ is different from Yampa's as we need
-- the @O@ type here.
loopPre :: c -> SF (a:*c) (b:*O c) -> SF a b
loopPre :: forall c a b. c -> SF (a :* c) (b :* O c) -> SF a b
loopPre c
c (SF DTime -> (a :* c) -> O (SF (a :* c) (b :* O c)) :* (b :* O c)
sf) = forall a b. (DTime -> a -> O (SF a b) :* b) -> SF a b
SF DTime -> a -> O (SF a b) :* b
sf'
  where sf' :: DTime -> a -> O (SF a b) :* b
sf' DTime
d a
a = let (O (SF (a :* c) (b :* O c))
r:* (b
b:*O c
c')) = DTime -> (a :* c) -> O (SF (a :* c) (b :* O c)) :* (b :* O c)
sf DTime
d (a
aforall a b. a -> b -> a :* b
:*c
c)
                  in (forall a. a -> O a
delay (forall c a b. c -> SF (a :* c) (b :* O c) -> SF a b
loopPre (forall a. O a -> a
adv O c
c') (forall a. O a -> a
adv O (SF (a :* c) (b :* O c))
r))forall a b. a -> b -> a :* b
:* b
b)


-- | Precomposition with a pure function.
(^>>) :: Arrow a => Box (b -> c) -> a c d -> a b d
Box (b -> c)
f ^>> :: forall (a :: * -> * -> *) b c d.
Arrow a =>
Box (b -> c) -> a c d -> a b d
^>> a c d
a = forall (a :: * -> * -> *) b c. Arrow a => Box (b -> c) -> a b c
arrBox Box (b -> c)
f forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> a c d
a

-- | Postcomposition with a pure function.
(>>^) :: Arrow a => a b c -> Box (c -> d) -> a b d
a b c
a >>^ :: forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> Box (c -> d) -> a b d
>>^ Box (c -> d)
f = a b c
a forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> forall (a :: * -> * -> *) b c. Arrow a => Box (b -> c) -> a b c
arrBox Box (c -> d)
f

-- | Precomposition with a pure function (right-to-left variant).
(<<^) :: Arrow a => a c d -> Box (b -> c) -> a b d
a c d
a <<^ :: forall (a :: * -> * -> *) c d b.
Arrow a =>
a c d -> Box (b -> c) -> a b d
<<^ Box (b -> c)
f = a c d
a forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
<<< forall (a :: * -> * -> *) b c. Arrow a => Box (b -> c) -> a b c
arrBox Box (b -> c)
f

-- | Postcomposition with a pure function (right-to-left variant).
(^<<) :: Arrow a => Box (c -> d) -> a b c -> a b d
Box (c -> d)
f ^<< :: forall (a :: * -> * -> *) c d b.
Arrow a =>
Box (c -> d) -> a b c -> a b d
^<< a b c
a = forall (a :: * -> * -> *) b c. Arrow a => Box (b -> c) -> a b c
arrBox Box (c -> d)
f forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
<<< a b c
a