{-# LANGUAGE RankNTypes, ScopedTypeVariables, DeriveTraversable, PatternSynonyms, UndecidableInstances, FlexibleInstances, ViewPatterns, InstanceSigs #-}

module Fix where

import Control.Monad (ap, join, (<=<))
import Control.Applicative (empty, Alternative, (<|>))
import Control.Arrow
import Data.Functor.Compose

-- Free f a = Mu x. a + f x
-- Cofree f a = Nu x. a * f x
-- Mu < Fix  < Nu

newtype Fix f = Fix (f (Fix f))

unfix :: Fix f -> f (Fix f)
unfix (Fix f) = f

cata :: Functor f => (f b -> b) -> Fix f -> b
cata alg = alg . fmap (cata alg) . unfix

cataM :: (Traversable t, Monad f) => (t b -> f b) -> Fix t -> f b
cataM alg = alg <=< traverse (cataM' alg) . unfix

ana :: Functor f => (a -> f a) -> a -> Fix f
ana coalg = Fix . fmap (ana coalg) . coalg

anaM :: (Monad m, Traversable f) => (a -> m (f a)) -> a -> m (Fix f)
anaM f = fmap Fix . traverse (anaM f) <=< f

futu :: Functor f => (a -> f (Free f a)) -> a -> Fix f
futu coalg t = ana go (Free(Fix(ReturnF t))) where
    go (Free(Fix(ReturnF a))) = coalg a
    go (Free(Fix(BindF fa))) = fmap Free fa

futuM :: (Traversable f, Monad m) => (a -> m (f (Free f a))) -> a -> m (Fix f)
futuM coalg t = anaM go (Free(Fix(ReturnF t))) where
    go (Free(Fix(ReturnF a)))  = coalg a
    go (Free(Fix(BindF fa))) = return (fmap Free fa)

histo :: Functor f => (f (Cofree f a) -> a) -> Fix f -> a
histo h = unfix >>> fmap worker >>> h where
    worker t = Cofree(Fix(CoBindF (histo h t) (fmap (uncofree . worker) (unfix t))))

histoM :: (Traversable f, Monad m) => (m (f (Cofree f a)) -> a) -> Fix f -> m a
histoM f = ?

hylo :: Functor f => (f b -> b) -> (a -> f a) -> a -> b
hylo f g = h where h = f . fmap h . g

hyloM :: (Traversable t, Monad m) => (t b -> m b) -> (a -> m (t a)) -> a -> m b
hyloM f g = f <=< traverse (hyloM f g) <=< g

data CofreeF f a r = CoBindF a (f r) deriving (Functor, Foldable, Traversable)

data FreeF f a r = ReturnF a | BindF (f r) deriving (Functor, Foldable, Traversable)

newtype Free f a = Free(Fix(FreeF f a))

newtype Cofree f a = Cofree(Fix(CofreeF f a))

unfree :: Free f a -> Fix (FreeF f a)
unfree (Free v) = v

uncofree :: Cofree f a -> Fix (CofreeF f a)
uncofree (Cofree v) = v

extract :: Cofree f a -> a
extract (Cofree(Fix(CoBindF a _))) = a

instance Functor f => Functor (Cofree f) where
    fmap :: (a -> b) -> Cofree f a -> Cofree f b
    fmap f = Cofree . go . uncofree where
        go (Fix (CoBindF a as)) = Fix (CoBindF (f a) (fmap go as))

instance Alternative f => Monad (Cofree f) where
    return a = Cofree (Fix (CoBindF a empty))
    (Cofree(Fix(CoBindF a m))) >>= f = case f a of
        (Cofree(Fix(CoBindF b n))) -> Cofree(Fix(CoBindF b (fmap uncofree ((fmap Cofree n) <|> (fmap ((>>= f) . Cofree) m)))))

instance Alternative f => Applicative(Cofree f) where
    pure = return
    (<*>) = ap

liftF :: Functor f => f a -> Free f a
liftF c = Free (Fix (BindF $ fmap (unfree . Free . Fix . ReturnF) c))

instance Functor f => Functor (Free f) where
    fmap :: (a -> b) -> Free f a -> Free f b
    fmap f = Free . cata go . unfree where
        go (ReturnF a) = Fix (ReturnF (f a))
        go (BindF   a) = Fix (BindF a)

instance Functor f => Applicative(Free f) where
    pure = return
    (<*>) = ap

instance Functor f => Monad (Free f) where
    return a = Free (Fix (ReturnF a))
    x >>= f = Free $ go $ unfree x where
                go (Fix (ReturnF a)) = unfree $ f a
                go (Fix (BindF a)) = Fix . BindF $ fmap go a

retract :: (Monad f, Traversable f) => Free f b -> f b
retract = cataM alg . unfree where
    alg (ReturnF a) = return a
    alg (BindF as) = as