{-# LANGUAGE DeriveFoldable            #-}
{-# LANGUAGE DeriveFunctor             #-}
{-# LANGUAGE DeriveTraversable         #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE TypeFamilies              #-}

module Control.Recursion
    ( -- * Typeclasses
      Base
    , Recursive (..)
    , Corecursive (..)
    -- * Types
    , Fix (..)
    , Mu (..)
    , Nu (..)
    , ListF (..)
    , NonEmptyF (..)
    -- * Recursion schemes
    , hylo
    , prepro
    , postpro
    , mutu
    , zygo
    , para
    , apo
    , elgot
    , coelgot
    , micro
    , meta
    , meta'
    , dicata
    , cata
    , ana
    -- * Mendler-style recursion schemes
    , mhisto
    , mcata
    -- * Monadic recursion schemes
    , cataM
    , anaM
    , hyloM
    , zygoM
    , zygoM'
    -- * Helper functions
    , lambek
    , colambek
    , hoist
    , refix
    -- * Additional types
    , Trans
    , Lens
    ) where

import           Control.Arrow       ((&&&))
import           Control.Composition ((.*), (.**))
import           Control.Monad       ((<=<))
import           Data.Foldable       (toList)
import           Data.List.NonEmpty  (NonEmpty (..))
import qualified Data.List.NonEmpty  as NE
import           Data.Traversable    (Traversable (..))
import           Numeric.Natural     (Natural)

type family Base t :: * -> *

class (Functor (Base t)) => Recursive t where

    project :: t -> Base t t

class (Functor (Base t)) => Corecursive t where

    embed :: Base t t -> t

-- | A map of \\( F \\)-algebras (pseudoprism)
type Trans s a = forall f. Functor f => (f a -> a) -> f s -> s

-- | A map of \\( F \\)-coalgebras
type Lens s a = forall f. Functor f => (a -> f a) -> s -> f s

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

data NonEmptyF a b = NonEmptyF a (Maybe b)
    deriving (Functor, Foldable, Traversable)

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

data Nu f = forall a. Nu (a -> f a) a

newtype Mu f = Mu (forall a. (f a -> a) -> a)

type instance Base (Fix f) = f

type instance Base (Fix f) = f

type instance Base (Mu f) = f

type instance Base (Nu f) = f

type instance Base Natural = Maybe

type instance Base [a] = ListF a

type instance Base (NonEmpty a) = NonEmptyF a

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

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

instance Functor f => Recursive (Nu f) where
    project (Nu f a) = Nu f <$> f a

instance Functor f => Corecursive (Nu f) where
    embed = colambek

instance Functor f => Recursive (Mu f) where
    project = lambek

instance Functor f => Corecursive (Mu f) where
    embed m = Mu (\f -> f (fmap (cata f) m))

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

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

instance Recursive (NonEmpty a) where
    project (x :| []) = NonEmptyF x Nothing
    project (x :| xs) = NonEmptyF x (Just (NE.fromList xs))

instance Corecursive (NonEmpty a) where
    embed (NonEmptyF x Nothing)   = x :| []
    embed (NonEmptyF x (Just xs)) = x :| toList xs

instance Functor f => Recursive (Fix f) where
    project = unFix

instance Functor f => Corecursive (Fix f) where
    embed = Fix

-- | 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 t) => (Base t a -> a) -> t -> a
cata f = c where c = f . fmap c . project
{-# NOINLINE [0] cata #-}

{-# RULES
  "cata/Mu" forall f (g :: forall a. (f a -> a) -> a). cata f (Mu g) = g f;
     #-}

-- | Anamorphism, meant to build up a structure recursively.
ana :: (Corecursive t) => (a -> Base t a) -> a -> t
ana g = a where a = embed . fmap a . g
{-# NOINLINE [0] ana #-}

{-# RULES
   "ana/Nu" forall (f :: a -> f a). ana f = Nu f;
      #-}

-- | Base functor for a list of type @[a]@.
-- | 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
{-# NOINLINE [0] hylo #-}

{-# RULES
  "ana/cata/hylo"  forall f g x. cata f (ana g x) = hylo f g x;
     #-}

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

zygoM :: (Recursive t, Traversable (Base t), Monad m) => (Base t b -> m b) -> (Base t (b, a) -> m a) -> t -> m a
zygoM f g = fmap snd . cataM (\x -> (,) <$> f (fmap fst x) <*> g x)

zygoM' :: (Recursive t, Traversable (Base t), Monad m) => (Base t b -> m b) -> (Base t (b, a) -> m a) -> t -> m a
zygoM' f g = fmap snd . cataM (\x -> do { a <- g x; b <- f (fmap fst x); pure (b, a) })

anaM :: (Corecursive t, Traversable (Base t), Monad m) => (a -> m (Base t 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 t, Corecursive t) => (t -> Base t t)
lambek = cata (fmap embed)

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

-- | Prepromorphism. Fold a structure while applying a natural transformation at each step.
prepro :: (Recursive t, Corecursive t) => (Base t t -> Base t t) -> (Base t 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 t, Corecursive t) => (Base t t -> Base t t) -> (a -> Base t a) -> a -> t
postpro e g = a'
    where a' = embed . fmap (ana (e . project) . a') . g

-- | A mutumorphism.
mutu :: (Recursive t) => (Base t (a, a) -> a) -> (Base t (a, a) -> a) -> t -> a
mutu f g =  snd . cata (f &&& g)

-- | Catamorphism collapsing along two data types simultaneously. Basically a fancy zygomorphism.
dicata :: (Recursive t) => (Base t (a, t) -> a) -> (Base t (a, t) -> t) -> t -> a
dicata = fst .** (cata .* (&&&))

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

-- | Paramorphism
para :: (Recursive t, Corecursive t) => (Base t (t, a) -> a) -> t -> a
para f = snd . cata (((,) . embed . fmap fst) <*> f)

-- | Gibbons' metamorphism. Tear down a structure, transform it, and then build up a new structure
meta :: (Corecursive t', Recursive t) => (a -> Base t' a) -> (b -> a) -> (Base t 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 = mc where mc = psi mc . unFix

-- | Mendler's histomorphism
mhisto :: (forall y. ((y -> c) -> (y -> f y) -> f y -> c)) -> Fix f -> c
mhisto psi = mh where mh = psi mh 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 = either id (phi . fmap h) . psi

-- | Anamorphism allowing shortcuts. Compare 'apo'
micro :: (Corecursive a) => (b -> Either a (Base a 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 . ((,) <*> (fmap h . psi))

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

hoist :: (Recursive s, Corecursive t)
      => (forall a. Base s a -> Base t a)
      -> s
      -> t
hoist = cata . (embed .)
{-# NOINLINE [0] hoist #-}

hoistMu :: (forall a. f a -> g a) -> Mu f -> Mu g
hoistMu eta (Mu f) = Mu (f . (. eta))

hoistNu :: (forall a. f a -> g a) -> Nu f -> Nu g
hoistNu n (Nu f x) = Nu (n . f) x

{-# RULES
  "hoist/hoistMu" forall (eta :: forall a. f a -> f a) (f :: forall a. (f a -> a) -> a). hoist eta (Mu f) = hoistMu eta (Mu f);
     #-}

{-# RULES
  "hoist/hoistNu" forall (eta :: forall a. f a -> f a) (f :: a -> f a) x. hoist eta (Nu f x) = hoistNu eta (Nu f x);
     #-}

refix :: (Recursive s, Corecursive t, Base s ~ Base t) => s -> t
refix = cata embed