{-# LANGUAGE Arrows, RankNTypes #-}
module FRP.Moe.Core where

import Prelude hiding ((.))
import Data.Monoid
import Data.Functor
import Control.Category
import Control.Applicative
import Control.Arrow

type DTime = Double

-- | SF a b can be seen as a function from [x1, x2, ...] to [y1, y2, ...], in which

-- | x1, x2, ... have the type of a, and y1, y2, ... have the type of b. Notice that 

-- | each value is in fact a sample of a continuous temporal function of values in 

-- | some certain time. The time stamps are omitted in most of the comments.
newtype SF a b = MkSF {runSF :: DTime -> a -> (b, SF a b)}

sfId :: SF a a
-- ^ Input:  [x1, x2, ...]

-- ^ Output: [x1, x2, ...]
sfId = MkSF (\dt x -> (x, sfId))

sfArr :: (a -> b) -> SF a b
-- ^ Input: [x1, x2, ...]

-- ^ Output: [f x1, f x2, ...]
sfArr f = sf where sf = MkSF (\dt x -> (f x, sf))

sfComp :: SF b c -> SF a b -> SF a c
-- ^ Input of sf1:  [y1, y2, ...]

-- ^ Output of sf1: [z1, z2, ...]

-- ^ Input of sf2:  [x1, x2, ...]

-- ^ Output of sf2: [y1, y2, ...]

-- ^ Input:         [x1, x2, ...]

-- ^ Output:        [z1, z2, ...]
sfComp sf1 sf2 = MkSF (\dt x -> let (y, sf2') = runSF sf2 dt x
                                    (z, sf1') = runSF sf1 dt y
                                in (z, sfComp sf1' sf2'))

sfLoop :: SF (a, c) (b, c) -> SF a b
-- ^ Returns a signal function that:

-- ^ Input:        [x1, x2, ...]

-- ^ Output:       [y1, y2, ...]

-- ^ whenever the following property holds for sf:

-- ^ Input of sf:  [(x1, z1), (x2, z2), ...]

-- ^ Output of sf: [(y1, z1), (y2, z2), ...]
sfLoop sf = MkSF (\dt x -> let ((y, z), sf') = runSF sf dt (x, z)
                           in (y, sfLoop sf'))

sfFirst :: SF a b -> SF (a, c) (b, c)
-- ^ Input of sf:  [x1, x2, ...]

-- ^ Output of sf: [y1, y2, ...]

-- ^ Input:        [(x1, z1), (x2, z2), ...]

-- ^ Output:       [(y1, z1), (y2, z2), ...]
sfFirst sf = MkSF (\dt (x, y) -> let (z, sf') = runSF sf dt x
                                 in ((z, y), sfFirst sf'))
  
sfSecond :: SF a b -> SF (c, a) (c, b)
-- ^ Input of sf:  [x1, x2, ...]

-- ^ Output of sf: [y1, y2, ...]

-- ^ Input:        [(z1, x1), (z2, x2), ...]

-- ^ Output:       [(z1, y1), (z2, y2), ...]
sfSecond sf = MkSF (\dt (x, y) -> let (z, sf') = runSF sf dt y
                                  in ((x, z), sfSecond sf'))

sfLeft :: SF a b -> SF (Either a c) (Either b c)              
-- ^ Input of sf:  [x1, x2, ...]

-- ^ Output of sf: [y1, y2, ...]

-- ^ Input:        [Left x1, Right z1, Left x2, ...]

-- ^ Output:       [Left y1, Right z1, Left y2, ...]
sfLeft sf = MkSF (\dt xy -> case xy of
                     Left x -> let (z, sf') = runSF sf dt x
                               in (Left z, sfLeft sf')
                     Right y -> (Right y, sfLeft sf))
            
sfRight :: SF a b -> SF (Either c a) (Either c b)            
-- ^ Input of sf:  [x1, x2, ...]

-- ^ Output of sf: [y1, y2, ...]

-- ^ Input:        [Right x1, Left z1, Right x2, ...]

-- ^ Output:       [Right y1, Left z2, Right y2, ...]
sfRight sf = MkSF (\dt xy -> case xy of
                      Left x -> (Left x, sfRight sf)
                      Right y -> let (z, sf') = runSF sf dt y
                                 in (Right z, sfRight sf'))

sfDelay :: a -> SF a a
-- ^ Input of sf:  [x1, x2, ...]

-- ^ Output of sf: [x0, x1, x2, ...]
sfDelay x0 = MkSF (\dt x -> (x0, sfDelay x))

dTime :: SF () DTime
-- ^ Input of sf:  [x1, x2, ...]

-- ^ Output of sf: [dt1, dt2, ...] (dt_i is the time interval between x_(i-1) and x_i)
dTime = MkSF (\dt x -> (dt, dTime))

instance Monoid b => Monoid (SF a b) where
  mempty = sfArr (\_ -> mempty)
  mappend sf1 sf2 = mappend <$> sf1 <*> sf2

instance Functor (SF a) where
  -- ^ Input of sf:  [x1, x2, ...]

  -- ^ Output of sf: [y1, y2, ...]

  -- ^ Input:        [x1, x2, ...]

  -- ^ Output:       [f y1, f y2, ...]
  fmap f sf = MkSF (\dt x -> let (y, sf') = runSF sf dt x
                             in (f y, fmap f sf'))

instance Applicative (SF a) where
  -- ^ Input:  [i1, i2, ...] (ignored)

  -- ^ Output: [x,  x,  ...]         
  pure x = sfArr (const x)
  -- ^ Input of sff:  [x1, x2, ...]

  -- ^ Output of sff: [f1, f2, ...]

  -- ^ Input of sfy:  [x1, x2, ...]

  -- ^ Output of sfy: [y1, y2, ...]

  -- ^ Input:         [x1, x2, ...]

  -- ^ Output:        [f1 y1, f2 y2, ...]
  sff <*> sfy = MkSF (\dt x -> let (f, sff') = runSF sff dt x
                                   (y, sfy') = runSF sfy dt x
                               in (f y, sff' <*> sfy'))

instance Category SF where
  id = sfId
  (.) = sfComp

-- ^ second, (&&&) and (***) are not necessary. They exist only for optimization
instance Arrow SF where  
  arr = sfArr
  first = sfFirst
  second = sfSecond
  -- ^ Input of sf1:  [x1, x2, ...]

  -- ^ Output of sf1: [y1, y2, ...]

  -- ^ Input of sf2:  [x1, x2, ...]

  -- ^ Output of sf2: [z1, z2, ...]

  -- ^ Input:         [x1, x2, ...]

  -- ^ Output:        [(y1, z1), (y2, z2), ...]
  sf1 &&& sf2 = MkSF (\dt x -> let (y, sf1') = runSF sf1 dt x
                                   (z, sf2') = runSF sf2 dt x
                               in ((y, z), sf1' &&& sf2'))
  -- ^ Input of sf1:  [x1, x2, ...]

  -- ^ Output of sf1: [y1, y2, ...]

  -- ^ Input of sf2:  [u1, u2, ...]

  -- ^ Output of sf2: [z1, z2, ...]

  -- ^ Input:         [(x1, u1), (x2, u2), ...]

  -- ^ Output:        [(y1, z1), (y2, z2), ...]
  sf1 *** sf2 = MkSF (\dt (x, y) -> let (z, sf1') = runSF sf1 dt x
                                        (u, sf2') = runSF sf2 dt y
                                    in ((z, u), sf1' *** sf2'))
  
instance ArrowLoop SF where  
  loop = sfLoop
  
class Arrow a => ArrowDelay a where
  delay :: b -> a b b
  
instance ArrowDelay SF where
  delay = sfDelay
  
-- ^ right, (|||) and (+++) are not necessary. They exist only for optimization
instance ArrowChoice SF where  
  left = sfLeft
  right = sfRight
  -- ^ Input of sf1:  [x1, x2, ...]

  -- ^ Output of sf1: [y1, y2, ...]

  -- ^ Input of sf2:  [u1, u2, ...]

  -- ^ Output of sf2: [z1, z2, ...]

  -- ^ Input:         [Left x1, Right u1, Left x2, ...]

  -- ^ Output:        [y1, z1, y2, ...]
  sf1 ||| sf2 = MkSF (\dt xy -> case xy of
                         Left x -> let (z, sf1') = runSF sf1 dt x
                                   in (z, sf1' ||| sf2)
                         Right y -> let (z, sf2') = runSF sf2 dt y
                                    in (z, sf1 ||| sf2'))
  -- ^ Input of sf1:  [x1, x2, ...]

  -- ^ Output of sf1: [y1, y2, ...]

  -- ^ Input of sf2:  [u1, u2, ...]

  -- ^ Output of sf2: [z1, z2, ...]

  -- ^ Input:         [Left x1, Right u1, Left x2, ...]

  -- ^ Output:        [Left y1, Right z1, Left y2, ...]
  sf1 +++ sf2 = MkSF (\dt xy -> case xy of
                         Left x -> let (z, sf1') = runSF sf1 dt x
                                   in (Left z, sf1' +++ sf2)
                         Right y -> let (z, sf2') = runSF sf2 dt y
                                    in (Right z, sf1 +++ sf2'))

data Event a = Event a | NoEvent

switch :: SF b (c, Event d) -> (d -> SF b c) -> SF b c
-- ^ Input of sf:       [x1, x2, x3, ...]

-- ^ Output of sf:      [(y1, NoEvent), (y2, Event e1), (y3, NoEvent), ...]

-- ^ Input of (gen e):  [x2, x3, ...]

-- ^ Output of (gen e): [y(e1)1, y(e1)2, ...]

-- ^ Input:             [x1, x2, x3, ...]

-- ^ Output:            [y1, y(e1)1, y(e1)2, ...]
switch sf gen = MkSF (\dt x -> let ((y, ev), sf') = runSF sf dt x
                               in case ev of
                                 NoEvent -> (y, switch sf' gen)
                                 Event e -> runSF (gen e) dt x)

dswitch :: SF b (c, Event d) -> (d -> SF b c) -> SF b c
-- ^ Input of sf:       [x1, x2, x3, ...]

-- ^ Output of sf:      [(y1, NoEvent), (y2, Event e1), (y3, NoEvent), ...]

-- ^ Input of (gen e):  [x2, x3, ...]

-- ^ Output of (gen e): [y(e1)1, y(e1)2, ...]

-- ^ Input:             [x1, x2, x3, ...]

-- ^ Output:            [y1, y2, y(e1)2, ...]
dswitch sf gen = MkSF (\dt x -> let ((y, ev), sf') = runSF sf dt x
                                in case ev of
                                  NoEvent -> (y, dswitch sf' gen)
                                  Event e -> (y, snd (runSF (gen e) dt x)))
                 
kswitch :: SF a b -> SF (a, b) (Event c) -> (SF a b -> c -> SF a b) -> SF a b
-- ^ Input of sf:            [x1, x2, x3, ...]

-- ^ Output of sf:           [y1, y2, y3, ...]

-- ^ Input of sfe:           [(x1, y1), (x2, y2), (x3, y3), ...]

-- ^ Output of sfe:          [NoEvent, Event e1, NoEvent, ...]

-- ^ Input of (gen sf3 e1):  [x2, x3, ...]

-- ^ Output of (gen sf3 e1): [y(sf3,e1)1, y(sf3,e1)2, ...]

-- ^ Input:                  [x1, x2, x3, ...]

-- ^ Output:                 [y1, y(sf3,e1)1, y(sf3,e2)2, ...]

-- ^ where sf1 = [x1, x2, x3, ...] -- ^> [y1, y2, y3, ...]

-- ^       sf2 =     [x2, x3, ...] -- ^>     [y2, y3, ...]

-- ^       sf3 =         [x3, ...] -- ^>         [y3, ...]

-- ^       ...         
kswitch sf sfe gen = MkSF (\dt x -> let (y, sf') = runSF sf dt x
                                        (ev, sfe') = runSF sfe dt (x, y)
                                    in case ev of
                                      NoEvent -> (y, kswitch sf' sfe' gen)
                                      Event e -> runSF (gen sf' e) dt x)

dkswitch :: SF a b -> SF (a, b) (Event c) -> (SF a b -> c -> SF a b) -> SF a b
-- ^ Input of sf:            [x1, x2, x3, ...]

-- ^ Output of sf:           [y1, y2, y3, ...]

-- ^ Input of sfe:           [(x1, y1), (x2, y2), (x3, y3), ...]

-- ^ Output of sfe:          [NoEvent, Event e1, NoEvent, ...]

-- ^ Input of (gen sf3 e1):  [x2, x3, ...]

-- ^ Output of (gen sf3 e1): [y(sf3,e1)1, y(sf3,e1)2, ...]

-- ^ Input:                  [x1, x2, x3, ...]

-- ^ Output:                 [y1, y2, y(sf3,e2)2, ...]

-- ^ where sf1 = [x1, x2, x3, ...] -- ^> [y1, y2, y3, ...]

-- ^       sf2 =     [x2, x3, ...] -- ^>     [y2, y3, ...]

-- ^       sf3 =         [x3, ...] -- ^>         [y3, ...]

-- ^       ...         
dkswitch sf sfe gen = MkSF (\dt x -> let (y, sf') = runSF sf dt x
                                         (ev, sfe') = runSF sfe dt (x, y)
                                     in case ev of
                                       NoEvent -> (y, dkswitch sf' sfe' gen)
                                       Event e -> (y, snd (runSF (gen sf' e) dt x)))
                      
pswitch :: Functor col => 
           (forall sf. a -> col sf -> col (b, sf)) -> 
           col (SF b c) -> 
           SF (a, col c) (Event d) -> 
           (col (SF b c) -> d -> SF a (col c)) ->
           SF a (col c)
pswitch route sfs sfe gen = MkSF (\dt x -> let ysfs = route x sfs
                                               zssfs = fmap (\(y, sf) -> runSF sf dt y) ysfs
                                               zs = fmap fst zssfs
                                               sfs' = fmap snd zssfs
                                               (e, sfe') = runSF sfe dt (x, zs)
                                           in case e of
                                             NoEvent -> (zs, pswitch route sfs' sfe' gen)
                                             Event ev -> runSF (gen sfs' ev) dt x)
                      
dpswitch :: Functor col => 
           (forall sf. a -> col sf -> col (b, sf)) -> 
           col (SF b c) -> 
           SF (a, col c) (Event d) -> 
           (col (SF b c) -> d -> SF a (col c)) ->
           SF a (col c)
dpswitch route sfs sfe gen = MkSF (\dt x -> let ysfs = route x sfs
                                                zssfs = fmap (\(y, sf) -> runSF sf dt y) ysfs
                                                zs = fmap fst zssfs
                                                sfs' = fmap snd zssfs
                                                (e, sfe') = runSF sfe dt (x, zs)
                                            in case e of
                                              NoEvent -> (zs, dpswitch route sfs' sfe' gen)
                                              Event ev -> (zs, snd (runSF (gen sfs' ev) dt x)))

nth :: Int -> DTime -> SF () a -> a
-- ^ Get the element (x_n) the output [y1, y2, ...] of sf.

-- ^ Time interval is fixed to (dt).
nth n dt sf = let f = runSF sf
                  (x, sf') = f dt ()
              in if n == 0 
                 then x 
                 else x `seq` nth (n - 1) dt sf'

gen :: SF a b -> [(DTime, a)] -> [b]
-- ^ Turn sf into a stream function.
gen sf [] = []
gen sf ((dt, x) : dtxs) = let f = runSF sf
                              (y, sf') = f dt x
                          in y `seq` (y : gen sf' dtxs)