{-# LANGUAGE DeriveFunctor         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}

module Control.Recursion
    ( -- * Typeclasses
      Base
    , Recursive (..)
    , Corecursive (..)
    -- * Types
    , Fix (..)
    -- * Recursion schemes
    , cata
    , ana
    , hylo
    , prepro
    , postpro
    , mutu
    , zygo
    , para
    , apo
    , elgot
    , coelgot
    , micro
    , meta
    , meta'
    -- * Mendler-style recursion schemes
    , mhisto
    , mcata
    -- * Monadic recursion schemes
    , cataM
    , anaM
    , hyloM
    -- * Helper functions
    , lambek
    , colambek
    ) where

import           Control.Monad   ((<=<))
import           Numeric.Natural (Natural)

class Base t (f :: * -> *) where

class (Functor f, Base t f) => Recursive f t where
    project :: t -> f t

class (Functor f, Base t f) => Corecursive f t where
    embed :: f t -> t

data ListF a b = Cons a b
               | Nil
               deriving (Functor)

newtype Fix f = Fix { unFix :: f (Fix f) }

instance Base (Fix t) f where

instance Base Natural Maybe where

instance Recursive Maybe Natural where
    project 0 = Nothing
    project n = Just (n-1)

instance Corecursive Maybe Natural where
    embed Nothing  = 0
    embed (Just n) = n+1

instance Base b (ListF a) where

instance Recursive (ListF a) [a] where
    project []     = Nil
    project (x:xs) = Cons x xs

instance Corecursive (ListF a) [a] where
    embed Nil         = []
    embed (Cons x xs) = x : xs

-- | Catamorphism. Folds a structure. (see [here](http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.41.125&rep=rep1&type=pdf))
cata :: (Recursive f t) => (f a -> a) -> t -> a
cata f = c where c = f . fmap c . project

-- | Anamorphism, meant to build up a structure recursively.
ana :: (Corecursive f t) => (a -> f a) -> a -> t
ana g = a where a = embed . fmap a . g

-- | Hylomorphism; fold a structure while buildiung it up.
hylo :: Functor f => (f b -> b) -> (a -> f a) -> a -> b
hylo f g = h where h = f . fmap h . g

cataM :: (Recursive f t, Traversable f, Monad m) => (f a -> m a) -> t -> m a
cataM f = c where c = f <=< (traverse c . project)

anaM :: (Corecursive f t, Traversable f, Monad m) => (a -> m (f a)) -> a -> m t
anaM f = a where a = (fmap embed . traverse a) <=< f

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

lambek :: (Recursive f t, Corecursive f t) => (t -> f t)
lambek = cata (fmap embed)

colambek :: (Recursive f t, Corecursive f t) => (f t -> t)
colambek = ana (fmap project)

-- | Prepromorphism. Fold a structure while applying a natural transformation at each step.
prepro :: (Recursive f t, Corecursive f t) => (f t -> f t) -> (f a -> a) -> t -> a
prepro e f = c
    where c = f . fmap (c . cata (embed . e)) . project

-- | Postpromorphism. Build up a structure, applying a natural transformation along the way.
postpro :: (Recursive f t, Corecursive f t) => (f t -> f t) -> (a -> f a) -> a -> t
postpro e g = a'
    where a' = embed . fmap (ana (e . project) . a') . g

-- | A mutumorphism.
mutu :: (Recursive f t) => (f (a, a) -> a) -> (f (a, a) -> a) -> t -> a
mutu f g = g . fmap (\x -> (mutu g f x, mutu f g x)) . project

-- | Zygomorphism (see [here](http://www.iis.sinica.edu.tw/~scm/pub/mds.pdf) for a neat example)
zygo :: (Recursive f t) => (f b -> b) -> (f (b, a) -> a) -> t -> a
zygo f g = snd . cata (\x -> (f $ fmap fst x, g x))

-- | Paramorphism
para :: (Recursive f t, Corecursive f t) => (f (t, a) -> a) -> t -> a
para f = snd . cata (\x -> (embed $ fmap fst x, f x))

-- | Gibbons' metamorphism. Tear down a structure, transform it, and then build up a new structure
meta :: (Corecursive f t', Recursive g t) => (a -> f a) -> (b -> a) -> (g b -> b) -> t -> t'
meta f e g = ana f . e . cata g

-- | Erwig's metamorphism. Essentially a hylomorphism with a natural
-- transformation in between. This allows us to use more than one functor in a
-- hylomorphism.
meta' :: (Functor g) => (f a -> a) -> (forall c. g c -> f c) -> (b -> g b) -> b -> a
meta' h e k = g
    where g = h . e . fmap g . k

-- | Mendler's catamorphism
mcata :: (forall y. ((y -> c) -> f y -> c)) -> Fix f -> c
mcata psi = psi (mcata psi) . unFix

-- | Mendler's histomorphism
mhisto :: (forall y. ((y -> c) -> (y -> f y) -> f y -> c)) -> Fix f -> c
mhisto psi = psi (mhisto psi) unFix . unFix

-- | Elgot algebra (see [this paper](https://arxiv.org/abs/cs/0609040))
elgot :: Functor f => (f a -> a) -> (b -> Either a (f b)) -> b -> a
elgot phi psi = h where h = (id `either` (phi . fmap h)) . psi

-- | Anamorphism that allows shortcuts.
micro :: (Corecursive f a) => (b -> Either a (f b)) -> b -> a
micro = elgot embed

-- | Elgot coalgebra
coelgot :: Functor f => ((a, f b) -> b) -> (a -> f a) -> a -> b
coelgot phi psi = h where h = phi . (\x -> (x, (fmap h . psi) x))

-- | Apomorphism
apo :: (Corecursive f t) => (a -> f (Either t a)) -> a -> t
apo g = a where a = embed . fmap (either id a) . g