{-# LANGUAGE CPP                 #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE DefaultSignatures   #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE InstanceSigs        #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeApplications    #-}

module Data.Algebra.Free
    ( -- * Free algebra class
      FreeAlgebra (..)
      -- ** Type level witnesses
    , Proof (..)
    , -- ** Algebra types \/ constraints
      AlgebraType
    , AlgebraType0
      -- * Combinators
    , unFoldMapFree
    , foldFree
    , natFree
    , fmapFree
    , joinFree
    , bindFree
    , cataFree
    , foldrFree
    , foldrFree'
    , foldlFree
    , foldlFree'
      -- * General free type
    , Free (..)
    )
    where

import           Prelude

#if __GLASGOW_HASKELL__ < 808
import           Data.DList (DList)
#endif
import           Data.DList as DList
import           Data.Functor.Identity (Identity (..))
import           Data.Fix (Fix, cata)
import           Data.Group (Group (..))
import           Data.Kind (Constraint, Type)
import           Data.List.NonEmpty (NonEmpty (..))
import           Data.Monoid ( Endo (..)
#if __GLASGOW_HASKELL__ < 808
                             , Monoid (..)
#endif
                             , Dual (..))
#if __GLASGOW_HASKELL__ < 808
import           Data.Semigroup ( Semigroup
                                , (<>)
                                )
#endif
import           Data.Algebra.Pointed (Pointed (..))

--
-- Prerequisites for @'FreeAlgebra'@
--

-- |
-- Type family which for each free algebra @m@ returns a type level lambda from
-- types to constraints.  It is describe the class of algebras for which this
-- free algebra is free.
--
-- A lawful instance for this type family must guarantee
-- that the constraint @'AlgebraType0' m f@ is implied by the @'AlgebraType'
-- m f@ constraint.  This guarantees that there exists a forgetful functor from
-- the category of types of kind @* -> *@ which satisfy @'AlgebraType' m@
-- constrain to the category of types of kind @* -> *@ which satisfy the
-- @'AlgebraType0 m@ constraint.
type family AlgebraType  (f :: k) (a :: l) :: Constraint

-- |
-- Type family which limits Hask to its full subcategory which satisfies
-- a given constraints.  Some free algebras, like free groups, or free abelian
-- semigroups have additional constraints on on generators, like @Eq@ or @Ord@.
type family AlgebraType0 (f :: k) (a :: l) :: Constraint

-- |
-- A proof that constraint @c@ holds for type @a@.
data Proof (c :: Constraint) (a :: l) where
    Proof :: c => Proof c a

-- |
-- A lawful instance has to guarantee that @'unFoldFree'@ is an inverse of
-- @'foldMapFree'@ (in the category of algebras of type @'AlgebraType' m@).
--
-- This in turn guaranties that @m@ is a left adjoint functor from full
-- subcategory of Hask (of types constrained by @'AlgebraType0' m) to algebras
-- of type @'AlgebraType' m@.  The right adjoint is the forgetful functor.  The
-- composition of left adjoin and the right one is always a monad, this is why
-- we will be able to build monad instance for @m@.
class FreeAlgebra (m :: Type -> Type)  where

    {-# MINIMAL returnFree, foldMapFree #-}

    -- | Injective map that embeds generators @a@ into @m@.
    returnFree :: a -> m a

    -- | The freeness property.
    foldMapFree
        :: forall d a
         . ( AlgebraType m d
           , AlgebraType0 m a
           )
        => (a -> d)   -- ^ a mapping of generators of @m@ into @d@
        -> (m a -> d) -- ^ a homomorphism from @m a@ to @d@

    -- |
    -- Proof that @AlgebraType0 m a => m a@ is an algebra of type @AlgebraType m@.
    -- This proves that @m@ is a mapping from the full subcategory of @Hask@ of
    -- types satisfying @AlgebraType0 m a@ constraint to the full subcategory
    -- satisfying @AlgebraType m a@, @'fmapFree'@ below proves that it's a functor.
    -- (@'codom'@ from codomain)
    codom  :: forall a. AlgebraType0 m a => Proof (AlgebraType m (m a)) (m a)

    default codom :: forall a. AlgebraType m (m a)
                  => Proof (AlgebraType m (m a)) (m a)
    codom = Proof

    -- |
    -- Proof that the forgetful functor from types @a@ satisfying @AgelbraType
    -- m a@ to @AlgebraType0 m a@ is well defined.
    forget :: forall a. AlgebraType  m a => Proof (AlgebraType0 m a) (m a)

    default forget :: forall a. AlgebraType0 m a
                   => Proof (AlgebraType0 m a) (m a)
    forget = Proof

--
-- Free combinators
--

-- |
-- Inverse of @'foldMapFree'@
--
-- It is uniquelly determined by its universal property (by Yonneda lemma):
--
-- prop> unFoldMapFree id = returnFree
--
-- Note that @'unFoldMapFree' id@ is the unit of the
-- [unit](https://ncatlab.org/nlab/show/unit+of+an+adjunction) of the
-- adjunction imposed by the @'FreeAlgebra'@ constraint.
unFoldMapFree
    :: FreeAlgebra m
    => (m a -> d)
    -> (a -> d)
unFoldMapFree f = f . returnFree
{-# INLINABLE unFoldMapFree #-}

-- |
-- All types which satisfy @'FreeAlgebra'@ constraint are foldable.
--
-- prop> foldFree . returnFree == id
--
-- @foldFree@ is the
-- [unit](https://ncatlab.org/nlab/show/unit+of+an+adjunction) of the
-- adjunction imposed by @FreeAlgebra@ constraint.
--
-- Examples:
--
-- > foldFree @[] = foldMap id
-- >              = foldr (<>) mempty
-- > foldFree @NonEmpty
-- >              = foldr1 (<>)
--
-- Note that @foldFree@ replaces the abstract \/ free algebraic operation in
-- @m a@ to concrete one in @a@.
foldFree
    :: forall m a .
       ( FreeAlgebra  m
       , AlgebraType  m a
       )
    => m a
    -> a
foldFree ma = case forget @m @a of
    Proof -> foldMapFree id ma
{-# INLINABLE foldFree #-}

-- |
-- The canonical quotient map from a free algebra of a wider class to a free
-- algebra of a narrower class, e.g. from a free semigroup to
-- free monoid, or from a free monoid to free commutative monoid,
-- etc.
--
-- prop> natFree . natFree == natFree
-- prop> fmapFree f . natFree == hoistFree . fmapFree f
--
-- the constraints:
-- * the algebra @n a@ is of the same type as algebra @m@ (this is
--    always true, just GHC cannot prove it here)
-- * @m@ is a free algebra generated by @a@
-- * @n@ is a free algebra generated by @a@
natFree :: forall m n a .
           ( FreeAlgebra  m
           , FreeAlgebra  n
           , AlgebraType0 m a
           , AlgebraType  m (n a)
           )
        => m a
        -> n a
natFree = foldMapFree returnFree
{-# INLINABLE natFree #-}

-- |
-- All types which satisfy @'FreeAlgebra'@ constraint are functors.
-- The constraint @'AlgebraType' m (m b)@ is always satisfied.
fmapFree :: forall m a b .
            ( FreeAlgebra  m
            , AlgebraType0 m a
            , AlgebraType0 m b
            )
         => (a -> b)
         -> m a
         -> m b
fmapFree f ma = case codom @m @b of
    Proof -> foldMapFree (returnFree . f) ma
{-# INLINABLE fmapFree #-}

-- |
-- @'FreeAlgebra'@ constraint implies @Monad@ constrain.
joinFree :: forall m a .
          ( FreeAlgebra  m
          , AlgebraType0 m a
          )
         => m (m a)
         -> m a
joinFree mma = case codom @m @a of
    Proof -> foldFree mma
{-# INLINABLE joinFree #-}

-- |
-- The monadic @'bind'@ operator.  @'returnFree'@ is the corresponding
-- @'return'@ for this monad.  This just @'foldMapFree'@ in disguise.
bindFree :: forall m a b .
            ( FreeAlgebra  m
            , AlgebraType0 m a
            , AlgebraType0 m b
            )
         => m a
         -> (a -> m b)
         -> m b
bindFree ma f = case codom @m @b of
    Proof -> foldMapFree f ma
{-# INLINABLE bindFree #-}

-- |
-- @'Fix' m@ is the initial algebra in the category of algebras of type
-- @'AlgebraType' m@ (the initial algebra is a free algebra generated by empty
-- set of generators, e.g. the @Viod@ type).
--
-- Another way of putting this is observing that @'Fix' m@ is isomorphic to @m
-- Void@ where @m@ is the /free algebra/.  This isomorphisms is given by
-- @
--   fixToFree :: (FreeAlgebra m, AlgebraType m (m Void), Functor m) => Fix m -> m Void
--   fixToFree = cataFree
-- @
-- For monoids the inverse is given by @'Data.Fix.ana' (\_ -> [])@.
cataFree :: ( FreeAlgebra  m
            , AlgebraType  m a
            , Functor m
            )
         => Fix m
         -> a
cataFree = cata foldFree

-- |
-- A version of @'Data.Foldable.foldr'@, e.g. it can specialize to
--
-- * @foldrFree \@[] :: (a -> b -> b) -> [a] -> b -> b@
-- * @foldrFree \@'Data.List.NonEmpty.NonEmpty' :: (a -> b -> b) -> 'Data.List.NonEmpty.NonEmpty' a -> b -> b@
foldrFree
    :: forall m a b .
       ( FreeAlgebra  m
       , AlgebraType  m (Endo b)
       , AlgebraType0 m a
       )
    => (a -> b -> b)
    -> b
    -> m a
    -> b
foldrFree f z t = appEndo (foldMapFree (Endo . f) t) z

-- |
-- Like @'foldrFree'@ but strict.
foldrFree'
    :: forall m a b .
       ( FreeAlgebra  m
       , AlgebraType  m (Dual (Endo (b -> b)))
       , AlgebraType0 m a
       )
    => (a -> b -> b)
    -> m a
    -> b
    -> b
foldrFree' f xs z0 = foldlFree f' id xs z0
    where
    f' k x z = k $! f x z

-- |
-- Generalizes @'Data.Foldable.foldl'@, e.g. it can specialize to
--
-- * @foldlFree \@[] :: (b -> a -> b) -> b -> [a] -> b@
-- * @foldlFree \@'Data.List.NonEmpty.NonEmpty' :: (b -> a -> b) -> b -> 'Data.List.NonEmpty.NonEmpty' a -> b@
foldlFree
    :: forall m a b .
       ( FreeAlgebra  m
       , AlgebraType  m (Dual (Endo b))
       , AlgebraType0 m a
       )
    => (b -> a -> b)
    -> b
    -> m a
    -> b
foldlFree f z t = appEndo (getDual (foldMapFree (Dual . Endo . flip f) t)) z

-- |
-- Like @'foldlFree'@ but strict.
foldlFree'
    :: forall m a b .
       ( FreeAlgebra  m
       , AlgebraType  m (Endo (b -> b))
       , AlgebraType0 m a
       )
    => (b -> a -> b)
    -> b
    -> m a
    -> b
foldlFree' f z0 xs = foldrFree f' id xs z0
    where
    f' x k z = k $! f z x


--
-- Instances
--


type instance AlgebraType0 Identity a = ()
type instance AlgebraType  Identity a = ()
instance FreeAlgebra Identity where
    returnFree = Identity
    foldMapFree f = f . runIdentity

type instance AlgebraType0 NonEmpty a = ()
type instance AlgebraType  NonEmpty m = Semigroup m
-- |
-- @'NonEmpty'@ is the free semigroup in the class of semigroup which are
-- strict in the left argument.
instance FreeAlgebra NonEmpty where
    returnFree a = a :| []
    -- @'foldMap'@ requires @'Monoid' d@ constraint which we don't need to
    -- satisfy here
    foldMapFree f (a :| []) = f a
    foldMapFree f (a :| (b : bs)) = f a <> foldMapFree f (b :| bs)

type instance AlgebraType0 [] a = ()
type instance AlgebraType  [] m = Monoid m
-- | 
-- Note that @'[]'@ is a free monoid only for monoids which multiplication is
-- strict in the left argument
-- [ref](http://comonad.com/reader/2015/free-monoids-in-haskell/). Note that
-- being strict adds additional equation to the monoid laws:
--
-- prop> undefined <> a = undefined
--
-- Thus, expectedly we get an equational theory for left / right / two-sided
-- strict monoids.
--
-- Snoc lists are free monoids in the class of monoids which are strict in the
-- right argument, @'Free' Monoid@ and @'DList' are free in the class of all
-- Haskell monoids.
instance FreeAlgebra [] where
    returnFree a = [a]
    foldMapFree = foldMap

type instance AlgebraType0 Maybe a = ()
type instance AlgebraType  Maybe m = Pointed m
instance FreeAlgebra Maybe where
    returnFree = Just
    foldMapFree _ Nothing  = point
    foldMapFree f (Just a) = f a

-- |
-- @'Free' c a@ represents free algebra for a constraint @c@ generated by
-- type @a@.
newtype Free (c :: Type -> Constraint) a = Free {
          runFree :: forall r. c r => (a -> r) -> r
        }

instance Semigroup (Free Semigroup a) where
    Free f <> Free g = Free $ \k -> f k <> g k

type instance AlgebraType0 (Free Semigroup) a = ()
type instance AlgebraType  (Free Semigroup) a = Semigroup a
instance FreeAlgebra (Free Semigroup) where
    returnFree a = Free $ \k -> k a
    foldMapFree f (Free k) = k f

instance Semigroup (Free Monoid a) where
    Free f <> Free g = Free $ \k -> f k `mappend` g k

instance Monoid (Free Monoid a) where
    mempty = Free (const mempty)
#if __GLASGOW_HASKELL__ <= 802
    mappend = (<>)
#endif


type instance AlgebraType0 (Free Monoid) a = ()
type instance AlgebraType  (Free Monoid) a = Monoid a
instance FreeAlgebra (Free Monoid) where
    returnFree a = Free $ \k -> k a
    foldMapFree f (Free k) = k f

type instance AlgebraType0 DList a = ()
type instance AlgebraType  DList a = Monoid a
-- |
-- @'DList'@ is isomorphic to @'Free' Monoid@; it is free in the class of all
-- monoids.
instance FreeAlgebra DList where
    returnFree = DList.singleton
    foldMapFree = foldMap

instance Semigroup (Free Group a) where
    Free f <> Free g = Free $ \k -> f k `mappend` g k

instance Monoid (Free Group a) where
    mempty = Free (const mempty)
#if __GLASGOW_HASKELL__ <= 802
    mappend = (<>)
#endif

instance Group (Free Group a) where
    invert (Free k) = Free (k . invert)

type instance AlgebraType0 (Free Group) a = ()
type instance AlgebraType  (Free Group) a = Group a
instance FreeAlgebra (Free Group) where
    returnFree a = Free $ \k -> k a
    foldMapFree f (Free k) = k f