{-# 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.
  stepSF :: DTime -> a -> (O(SF a b), b)}

-- | The identity signal function that does nothing.
identity :: SF a a
identity = SF (\ _ x -> (delay identity,x))

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

-- | Compute the integral of a signal. The first argument is the
-- offset.
integral :: (Stable a, VectorSpace a s) => a -> SF a a
integral acc = SF sf'
  where sf' t a = let acc' = acc ^+^ (realToFrac t *^ a)
                  in (delay (integral acc'), 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 (SF sf) f = SF sf'
  where sf' t a = let (nxt, (b,c')) =  sf t a
                  in case c' of
                       Just' c -> stepSF (unbox f c) t a
                       Nothing' -> (delay (switch (adv nxt) f), 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 (SF sf) = SF sf'
  where sf' t (a,m) = case m of
                        Just' (SF newSf) ->
                           let (nxt, b) = newSf t a
                           in (delay (rSwitch (adv nxt)),b)
                        Nothing' -> let (nxt, b) = sf t a
                                   in (delay (rSwitch (adv nxt)),b)

-- | Constant signal function.
constant :: Stable b => b -> SF a b
constant x = run
  where run = SF (\ _ _ -> (delay run,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 --> (SF sf) = SF sf'
  where sf' d x = (fst (sf d x),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 -:> sf = SF sf'
  where sf' _d _x = (sf,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 >-- (SF sf) = SF sf'
  where sf' d _a = sf d a


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

-- | Override initial value of input signal.

initially :: a -> SF a a
initially = (--> 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 f = run where
  run = SF (\ _d a -> (delay run, unbox f a ))


-- | Apply a signal function to the first component.
firstPrim :: SF a b -> SF (a,c) (b,c)
firstPrim (SF sf) = SF sf'
  where sf' d (a,c) = let (r, b) = sf d a
                      in (delay (firstPrim (adv r)), (b,c))

-- | Apply a signal function to the second component.
secondPrim :: SF a b -> SF (c,a) (c,b)
secondPrim (SF sf) = SF sf'
  where sf' d (c,a) = let (r, b) = sf d a
                      in (delay (secondPrim (adv r)), (c,b))


-- | Apply two signal functions in parallel.
parSplitPrim :: SF a b -> SF c d  -> SF (a,c) (b,d)
parSplitPrim (SF sf1) (SF sf2) = SF sf'
  where sf' dt (a,c) = let (r1, b) = sf1 dt a
                           (r2, d) = sf2 dt c
                       in (delay (parSplitPrim (adv r1) (adv r2)), (b,d))

-- | Apply two signal functions in parallel on the same input.
parFanOutPrim :: SF a b -> SF a c -> SF a (b, c)
parFanOutPrim (SF sf1) (SF sf2) = SF sf'
  where sf' dt a = let (r1, b) = sf1 dt a
                       (r2, c) = sf2 dt a
                   in (delay (parFanOutPrim (adv r1) (adv r2)), (b,c))

instance Category SF where
  id = identity
  (.) = compose

instance Arrow SF where
  arrBox = arrPrim
  first = firstPrim
  second = secondPrim
  (***) = parSplitPrim
  (&&&) = 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 c (SF sf) = SF sf'
  where sf' d a = let (r, (b,c')) = sf d (a,c)
                  in (delay (loopPre (adv c') (adv r)), b)


-- | Precomposition with a pure function.
(^>>) :: Arrow a => Box (b -> c) -> a c d -> a b d
f ^>> a = arrBox f >>> a

-- | Postcomposition with a pure function.
(>>^) :: Arrow a => a b c -> Box (c -> d) -> a b d
a >>^ f = a >>> arrBox f

-- | Precomposition with a pure function (right-to-left variant).
(<<^) :: Arrow a => a c d -> Box (b -> c) -> a b d
a <<^ f = a <<< arrBox f

-- | Postcomposition with a pure function (right-to-left variant).
(^<<) :: Arrow a => Box (c -> d) -> a b c -> a b d
f ^<< a = arrBox f <<< a