{-# LANGUAGE Arrows #-}

module Main(main) where

import Control.Arrow
import Control.Category
import Prelude hiding (id, (.))

class ArrowLoop a => ArrowCircuit a where
        delay :: b -> a b b

-- stream map instance

data Stream a = Cons a (Stream a)

instance Functor Stream where
        fmap f ~(Cons a as) = Cons (f a) (fmap f as)

zipStream :: Stream a -> Stream b -> Stream (a,b)
zipStream ~(Cons a as) ~(Cons b bs) = Cons (a,b) (zipStream as bs)

unzipStream :: Stream (a,b) -> (Stream a, Stream b)
unzipStream abs = (fmap fst abs, fmap snd abs)

newtype StreamMap a b = StreamMap (Stream a -> Stream b)
unStreamMap (StreamMap f) = f

instance Category StreamMap where
        id = StreamMap id
        StreamMap f . StreamMap g = StreamMap (f . g)

instance Arrow StreamMap where
        arr f = StreamMap (fmap f)
        first (StreamMap f) =
                StreamMap (uncurry zipStream . first f . unzipStream)

instance ArrowLoop StreamMap where
        loop (StreamMap f) =
                StreamMap (loop (unzipStream . f . uncurry zipStream))

instance ArrowCircuit StreamMap where
        delay a = StreamMap (Cons a)

listToStream :: [a] -> Stream a
listToStream = foldr Cons undefined

streamToList :: Stream a -> [a]
streamToList (Cons a as) = a:streamToList as

runStreamMap :: StreamMap a b -> [a] -> [b]
runStreamMap (StreamMap f) as =
        take (length as) (streamToList (f (listToStream as)))

-- simple automaton instance

data Auto a b = Auto (a -> (b, Auto a b))

instance Category Auto where
        id = Auto $ \a -> (a, id)
        Auto f . Auto g = Auto $ \b ->
                                let     (c, g') = g b
                                        (d, f') = f c
                                in      (d, f' . g')

instance Arrow Auto where
        arr f = Auto $ \a -> (f a, arr f)
        first (Auto f) = Auto $ \(b,d) -> let (c,f') = f b in ((c,d), first f')

instance ArrowLoop Auto where
        loop (Auto f) = Auto $ \b ->
                let     (~(c,d), f') = f (b,d)
                in      (c, loop f')

instance ArrowCircuit Auto where
        delay a = Auto $ \a' -> (a, delay a')

runAuto :: Auto a b -> [a] -> [b]
runAuto (Auto f) [] = []
runAuto (Auto f) (a:as) = let (b, f') = f a in b:runAuto f' as

-- Some simple example circuits

-- A resettable counter (first example in several Hawk papers):

counter :: ArrowCircuit a => a Bool Int
counter = proc reset -> do
        rec     output <- returnA -< if reset then 0 else next
                next <- delay 0 -< output+1
        returnA -< output

-- Some other basic circuits from the Hawk library.

-- flush: when reset is True, return d for n ticks, otherwise copy value.
-- (a variation on the resettable counter)

flush :: ArrowCircuit a => Int -> b -> a (b, Bool) b
flush n d = proc (value, reset) -> do
        rec     count <- returnA -< if reset then n else max (next-1) 0
                next <- delay 0 -< count
        returnA -< if count > 0 then d else value

-- latch: on each tick, return the last value for which reset was True,
-- or init if there was none.
--
latch :: ArrowCircuit a => b -> a (b, Bool) b
latch init = proc (value, reset) -> do
        rec     out <- returnA -< if reset then value else last
                last <- delay init -< out
        returnA -< out

-- Some tests using the counter

test_input = [True, False, True, False, False, True, False, True]
test_input2 = zip [1..] test_input

-- A test of the resettable counter.

main = do
        print (runStreamMap counter test_input)
        print (runAuto counter test_input)
        print (runStreamMap (flush 2 0) test_input2)
        print (runAuto (flush 2 0) test_input2)
        print (runStreamMap (latch 0) test_input2)
        print (runAuto (latch 0) test_input2)

-- A step function (cf current in Lustre)

step :: ArrowCircuit a => b -> a (Either b c) b
step b = proc x -> do
                rec     last_b <- delay b -< getLeft last_b x
                returnA -< last_b
        where   getLeft _ (Left b) = b
                getLeft b (Right _) = b
