{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE TypeFamilies           #-}

{-# OPTIONS_HADDOCK show-extensions #-}

module Control.Category.FreeEff
  ( EffCategory (..)
  , FreeEffCat (..)
  , liftCat
  , foldNatLift
  , liftKleisli
  ) where

import Prelude hiding (id, (.))

import Control.Arrow (Kleisli (..))
import Control.Category (Category (..))
import Data.Functor.Identity (Identity (..))

import Control.Category.Free (Cat)
import Control.Algebra.Free2 (FreeAlgebra2 (..))
import Data.Algebra.Free (AlgebraType, AlgebraType0, proof)


-- | Categories which can lift monadic actions, i.e. effectful categories.
--
class Category c => EffCategory c m | c -> m where
  lift :: m (c a b) -> c a b

instance Monad m => EffCategory (Kleisli m) m where
  lift m = Kleisli (\a -> m >>= \(Kleisli f) -> f a)

instance EffCategory (->) Identity where
  lift = runIdentity

-- | Category transformer, which adds @'EffCategory'@ instance to the
-- underlying base category.
--
data FreeEffCat :: (* -> *) -> (k -> k -> *) -> k -> k -> * where
  Base :: c a b -> FreeEffCat m c a b
  Lift :: m (FreeEffCat m c a b) -> FreeEffCat m c a b

instance (Functor m, Category c) => Category (FreeEffCat m c) where
  id = Base id
  Base f  . Base g  = Base $ f . g
  f       . Lift mg = Lift $ (f .) <$> mg
  Lift mf . g       = Lift $ (. g) <$> mf

instance (Functor m, Category c) => EffCategory (FreeEffCat m c) m where
  lift = Lift

type instance AlgebraType0 (FreeEffCat m) c = (Monad m, Category c)
type instance AlgebraType  (FreeEffCat m) c  = EffCategory c m
instance Monad m => FreeAlgebra2 (FreeEffCat m) where
  liftFree2    = Base
  foldNatFree2 nat (Base cab)  = nat cab
  foldNatFree2 nat (Lift mcab) = lift $ foldNatFree2 nat <$> mcab

  codom2  = proof
  forget2 = proof

-- | Wrap a transition into a free category @'Cat'@ and then in
-- @'FreeEffCat'@
--
-- prop> liftCat tr = Base (tr :.: Id)
--
liftCat :: Monad m => tr a b -> FreeEffCat m (Cat tr) a b
liftCat = liftFree2 . liftFree2

-- | Fold @'FreeLifing'@ category based on a free category @'Cat' tr@ using
-- a functor @tr x y -> c x y@.
--
foldNatLift
  :: (Monad m, EffCategory c m)
  => (forall x y. tr x y -> c x y)
  -> FreeEffCat m (Cat tr) a b
  -> c a b
foldNatLift nat = foldNatFree2 (foldNatFree2 nat)

-- |  Functor from @'->'@ category to @'Kleisli' m@.  If @m@ is @Identity@ then
-- it will respect @'lift'@ i.e. @liftKleisli (lift ar) = lift (liftKleisli <$>
-- ar).
--
liftKleisli :: Applicative m => (a -> b) -> Kleisli m a b
liftKleisli f = Kleisli (pure . f)