{-# OPTIONS_GHC -farrows #-}

{---

   Code from Arrows and Computation by Ross Paterson.

---}

module ArrowComp where 

import Control.Arrow

--
-- Utility functions
--
(|*|) f g (a, b) = (f a, g b)
assoc ~(~(a,b),c) = (a,(b,c))
unassoc ~(a, ~(b,c)) = ((a,b),c)
zipMap h s = (fst h s, snd h s)
unzipMap h  = (fst . h, snd . h)

--
-- Some notions of computation
--

--
-- State transformer
--
newtype State s i o = St ((s,i) -> (s,o))
instance Arrow (State s) where
    pure f = St $ id |*| f
    St f >>> St g = St $ g . f
    first (St f) = St $ assoc . (f |*| id) . unassoc

--
-- Non-deterministic functions
--
newtype NonDet i o = ND (i -> [o])
instance Arrow NonDet where
    pure f = ND $ \b -> [f b]
    ND f >>> ND g = ND $ \b -> [d | c <- f b, d <- g c]
    first (ND f) = ND $ \(b,d) -> [(c,d) | c <- f b]

--
-- Map transformer
--
newtype MapTrans s i o = MT ((s -> i) -> (s -> o))
instance Arrow (MapTrans s) where
    pure f = MT $ (f .)
    MT f >>> MT g = MT $ g . f
    first (MT f) = MT $ zipMap . (f |*| id) . unzipMap

--
-- Automata
--
newtype Auto i o = Au (i -> (o, Auto i o))
instance Arrow Auto where
    pure f = Au $ \b -> (f b, pure f)
    Au f >>> Au g = Au $ \b ->
                        let (c, f') = f b
                            (d, g') = g c
                        in (d, f' >>> g')
    first (Au f) = Au $ \(b, d) ->
                        let (c, f') = f b
                        in ((c, d), first f')

--
-- General add for Arrows
--
addA :: Arrow a => a b Int -> a b Int -> a b Int
addA f g = f &&& g >>> pure (uncurry (+))


--
-- Reader and Writer
--
newtype Reader s i o = Rd ((s,i) -> o)
instance Arrow (Reader s) where
    pure f = Rd $ f . snd
    Rd f >>> Rd g = Rd $ \b -> g (fst b, f b)
    first (Rd f) = Rd $ (f |*| id) . unassoc

newtype Writer i o =  Wt (i -> (String, o))
instance Arrow Writer where
    pure f = Wt $ \b -> ("", f b)
    Wt f >>> Wt g = Wt $ \b ->
                        let (o1,c) = f b
                            (o2,d) = g c
                        in (o1 ++ o2, d) 
    first (Wt f) = Wt $ assoc . (f |*| id)


--
-- Why is ListMap not an arrow?
--
newtype ListMap i o = LM ([i] -> [o])
instance Arrow ListMap where
    pure f = LM $ map f
    LM f >>> LM g = LM $ g . f
    first (LM f) = LM $ (uncurry zip) . (f |*| id) . unzip
appLM (LM f) x = f x


--
-- Streams
--
data Stream a = Cons a (Stream a)
mapSt f (Cons x xs) = Cons (f x) (mapSt f xs)
unzipSt xs = (mapSt fst xs, mapSt snd xs)
zipSt (Cons x xs) (Cons y ys) = Cons (x,y) (zipSt xs ys)

newtype StreamMap i o = SM (Stream i -> Stream o)
instance Arrow StreamMap where
    pure f = SM $ mapSt f
    SM f >>> SM g = SM $ g . f
    first (SM f) = SM $ uncurry zipSt . (f |*| id) . unzipSt

--
-- Monads
--
mkPair :: Arrow a => b -> a g (b, g)
mkPair b = pure $ \c -> (b, c)

curryA :: Arrow a => a (b,g) d -> b -> a g d
curryA f b = mkPair b >>> f


newtype Kl m a b = Kl (a -> m b)
instance Monad m => Arrow (Kl m) where
    pure f = Kl $ \b -> return $ f b
    Kl f >>> Kl g = Kl $ \b -> f b >>= g
    first (Kl f) = Kl $ \(b, d) -> f b >>= (\c -> return (c, d))

instance Monad m => ArrowApply (Kl m) where
    app = Kl $ \(Kl f, x) -> f x

newtype ArrMon a b = ArrM (a () b)
instance ArrowApply a => Monad (ArrMon a) where
    return x = ArrM $ pure (const x)
    ArrM m >>= f = ArrM $ m >>> pure (\x -> let ArrM h = f x in (h, undefined)) >>> app


--
-- Homogeneous Functions
--
data BalTree a = Zero a | Succ (BalTree (a, a)) deriving Show

tree0 = Zero 1 
tree1 = Succ $ Zero (1,2) 
tree2 = Succ $ Succ $ Zero ((1,2),(3,4)) 
tree3 = Succ $ Succ $ Succ $ Zero (((1,2),(3,4)),((5,6),(7,8))) 
tree4 = Succ $ Succ $ Succ $ Succ $ Zero ((((1,2),(3,4)),((5,6),(7,8))),(((9,10),(11,12)),((13,14),(15,16))))


data Hom a b = (a -> b) :&: Hom (a, a) (b, b)

apply :: Hom a b -> BalTree a -> BalTree b
apply (f :&: _) (Zero x) = Zero (f x)
apply (_ :&: fs) (Succ t) = Succ (apply fs t)

swap (a,b) = (b,a)
transpose ((a, b), (c, d)) = ((a, c), (b, d))

instance Arrow Hom where
    pure f = f :&: pure (f |*| f)
    (f :&: fs) >>> (g :&: gs) = (g . f) :&: (fs >>> gs)
    first (f :&: fs) = (f |*| id) :&: (pure transpose >>> first fs >>> pure transpose)


rsh :: a -> Hom a a
rsh v = const v :&: proc (o, e) -> do
                      o' <- rsh v -< e
                      returnA -< (o', o)

rsh0 :: a -> Hom a a
rsh0 v  = const v :&: (
    (pure (\ (o,e) -> e) >>> rsh0 v) &&& returnA >>> 
    pure (\ (o',(o,e)) -> (o',o)) >>> 
    returnA
 ) 

rsh1 :: a -> Hom a a
rsh1 v = const v :&: (second (rsh1 v) >>> pure swap)


scan :: Num b => Hom b b
scan = id :&: proc (o, e) -> do
                e' <- scan -< o + e
                el <- rsh 0 -< e'
                returnA -< (el + o, e')

scan0 :: Num b => Hom b b
scan0 = id :&: (
    (pure (\ (o,e) -> o + e) >>> scan0) &&& returnA >>>
    (pure (\ (e',(o,e)) -> e') >>> rsh 0) &&& returnA >>>
    pure (\ (el,(e',(o,e))) -> (el + o, e')) >>>
    returnA 
 )

scan1 :: Num b => Hom b b
scan1 = id :&: (
    (pure (\ (x,y) -> x + y) >>> scan) &&& pure fst >>> 
    (first (rsh 0) >>> pure (\ (x,y) -> x + y)) &&& pure fst
 )


butterfly :: ((b,b) -> (b,b)) -> Hom b b
butterfly f = id :&: proc (o,e) -> do
                       o' <- butterfly f -< o
                       e' <- butterfly f -< e
                       returnA -< f (o', e')

butterfly0 :: ((b,b) -> (b,b)) -> Hom b b
butterfly0 f = id :&: (butterfly0 f *** butterfly0 f >>> pure f)

butterfly1 :: ((b,b) -> (b,b)) -> Hom b b
butterfly1 f = id :&: (first (butterfly0 f) >>> second (butterfly0 f) >>> pure f)

rev :: Hom b b
rev = butterfly swap

unriffle :: Hom (b,b) (b,b)
unriffle = butterfly transpose

bisort :: Ord b => Hom b b
bisort = butterfly cmp where cmp (x,y) = (min x y, max x y)


--
-- Conditionals
--
(|+|) :: (a -> a') -> (b -> b') -> Either a b -> Either a' b'
(f |+| _) (Left a) = Left $ f a
(_ |+| g) (Right b) = Right $ g b

assocSum :: Either (Either a b) c -> Either a (Either b c)
assocSum (Left (Left a)) = Left a
assocSum (Left (Right b)) = Right $ Left b
assocSum (Right c) = Right $ Right c

distr :: (Either a b, c) -> Either (a, c) (b, c)
distr (Left a, c) = Left (a, c)
distr (Right b, c) = Right (b, c)

instance ArrowChoice Auto where
    left (Au f) = Au lf where
        lf (Left b) = let (c, f') = f b
                      in (Left c, left f')
        lf (Right c) = (Right c, left $ Au f)

instance ArrowChoice (State s) where
    left (St f) = St lf where
        lf (s, Left i) = let (s', o) = f (s, i) in (s', Left o)
        lf (s, Right d) = (s, Right d)

instance ArrowChoice NonDet where
    left (ND f) = ND lf where
        lf (Left b) = map Left $ f b
        lf (Right c) = [Right c]

instance ArrowChoice StreamMap where
    left (SM f) = SM $ \x -> combine x (f $ justLeft x) where
        justLeft (Cons (Left x) xs) = Cons x $ justLeft xs
        justLeft (Cons (Right _) xs) = justLeft xs

        combine (Cons (Left _) xs) (Cons y ys) = Cons (Left y) $ combine xs ys
        combine (Cons (Right x) xs) ys = Cons (Right x) $ combine xs ys

newtype Except a b g = E (a b (Either String g))
instance ArrowChoice a => Arrow (Except a) where
    pure f = E $ pure $ Right . f
    E f >>> E g = E $ f >>> (pure Left ||| g)
    first (E f) = E $ first f >>> pure distr >>> left (pure fst)  -- (pure $ (fst |+| id) . distr)


--
-- Loops
--
trace :: ((b, d) -> (g, d)) -> b -> g
trace f b = let (c, d) = f (b, d) in c

instance ArrowLoop (State s) where
    loop (St f) = St $ trace $ unassoc . f . assoc

instance ArrowLoop (MapTrans s) where
    loop (MT f) = MT $ trace $ unzipMap . f . zipMap

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

instance ArrowLoop StreamMap where
    loop (SM f) = SM $ \bs -> 
        let (gs, ds) = unzipSt $ f $ zipSt bs ds
        in gs

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

instance ArrowCircuit StreamMap where
    delay x = SM $ Cons x

instance ArrowCircuit Auto where
    delay x = Au $ \c -> (x, delay c)

runAuto :: Auto b g -> [b] -> [g]
runAuto _ [] = []
runAuto (Au f) (x:xs) = let (y,f') = f x in y : runAuto f' xs

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