Adjoint folds, using regular Category from base (but custom Functor)

adjointFoldPlain.hs

-- http://www.cs.ox.ac.uk/ralf.hinze/SSGIP10/AdjointFolds.pdf
-- http://www.cs.ox.ac.uk/ralf.hinze/publications/SCP-78-11.pdf
-- https://www.cs.ox.ac.uk/people/nicolas.wu/papers/URS.pdf
-- https://arxiv.org/pdf/2202.13633.pdf

{-# LANGUAGE
    MultiParamTypeClasses
  , FunctionalDependencies
  , GADTs
  , PolyKinds
  , DataKinds
  , ConstraintKinds
  , StandaloneKindSignatures
  , QuantifiedConstraints
  , TypeFamilies
  , RankNTypes
  , FlexibleContexts
  , FlexibleInstances
  , UndecidableInstances
  , DeriveFunctor
  , DerivingVia
  , TypeOperators
  , TypeApplications
  , LambdaCase
  , ImplicitParams
  , AllowAmbiguousTypes
  , ScopedTypeVariables
  , UndecidableSuperClasses
  #-}
import Prelude hiding (Functor, (.), fmap, id, mapM, map)
import qualified Prelude as P

import Control.Category

import Control.Applicative (Const(..))
import Control.Arrow ((&&&), (***), (|||), (+++), Kleisli(..), arr)
import Control.Monad (join)
import Control.Monad.Free (Free(..), wrap)
import Control.Comonad (Comonad(..), Cokleisli(..), (=>>))
import Control.Comonad.Cofree (Cofree(..), unwrap)
import Data.Bifunctor (first)
import Data.Distributive
import Data.Traversable
import Data.Functor.Identity
import Data.Functor.Kan.Lan
import Data.Functor.Kan.Ran
import Data.Functor.Rep
import Data.Kind (Constraint, Type)
import Data.Profunctor.Unsafe ((#.), (.#))
import GHC.Base (IP(..))

class Any a
instance Any a

type family Obj (c :: k -> k -> Type) (a :: k) :: Constraint
type instance Obj (->) a = Any a
type instance Obj (Kleisli m) a = Any a
type instance Obj (Cokleisli w) a = Any a

class Obj cod o => Obj' cod o
instance Obj cod o => Obj' cod o

class (Functor f dom cod, forall o. Obj dom o => Obj' cod (f o)) => FunctorLiftingObj f dom cod | f -> dom cod
instance (Functor f dom cod, forall o. Obj dom o => Obj' cod (f o)) => FunctorLiftingObj f dom cod

class (Category dom, Category cod) => Functor f dom cod | f -> dom cod where
  fmap :: dom a b -> cod (f a) (f b)

class (Category k) => FunctorDefault (k :: x -> x -> Type) where
  type IsFunctor' k (f :: x -> x) :: Constraint
  map :: IsFunctor' k f => k a b -> k (f a) (f b)

class (IsFunctor' k f, forall o. Obj k o => Obj' k (f o)) => IsFunctor k f
instance (IsFunctor' k f, forall o. Obj k o => Obj' k (f o)) => IsFunctor k f


class FunctorDefault k => Fixable (k :: x -> x -> Type) where
  type Fix k (f :: x -> x) :: x
  fix :: k (f (Fix k f)) (Fix k f)
  unfix :: k (Fix k f) (f (Fix k f))


hylo :: (Fixable k, IsFunctor k f) => k (f b) b -> k a (f a) -> k a b
hylo alg coa = c where c = alg . map c . coa

cata :: (Fixable k, IsFunctor k f) => k (f b) b -> k (Fix k f) b
cata alg = hylo alg unfix

ana  :: (Fixable k, IsFunctor k f) => k a (f a) -> k a (Fix k f)
ana  coa = hylo fix coa


class (Functor f d c, Functor u c d) => Adjunction (f :: y -> x) (u :: x -> y) (c :: x -> x -> Type) (d :: y -> y -> Type) | f -> u c d, u -> f c d where
  leftAdjunct  :: Obj d a => c (f a) b -> d a (u b)
  rightAdjunct :: Obj c b => d a (u b) -> c (f a) b
  unit :: Obj d a => d a (u (f a))
  unit = leftAdjunct id
  counit :: Obj c b => c (f (u b)) b
  counit = rightAdjunct id


wrapCata
  :: (Adjunction l r c d, Obj c b, Obj d frb)
  => (d frb (r b) -> d fixf (r b))
  -> (c (l frb) b -> c (l fixf) b)
wrapCata f = rightAdjunct . f . leftAdjunct

wrapAna
  :: (Adjunction l r c d, Obj c fla, Obj d a)
  => (c (l a) fla -> c (l a) fixf)
  -> (d a (r fla) -> d a (r fixf))
wrapAna f = leftAdjunct . f . rightAdjunct

wrapHylo
  :: (Adjunction m r c d, Obj c b, Obj d frb,
      Adjunction l m d c, Obj d fla, Obj c a)
  => (d frb (r b) -> d (l a) fla -> d (l a) (r b))
  -> (c (m frb) b -> c a (m fla) -> c a b)
wrapHylo f alg coalg = rightAdjunct (f (leftAdjunct alg) (rightAdjunct coalg)) . unit


adjointFold :: (Adjunction l r c d, Fixable d, IsFunctor d f, Obj c b, Obj d (r b))
            => c (l (f (r b))) b -> c (l (Fix d f)) b
adjointFold = wrapCata cata

adjointUnfold :: (Adjunction l r d c, Fixable d, IsFunctor d f, Obj c a, Obj d (l a))
              => c a (r (f (l a))) -> c a (r (Fix d f))
adjointUnfold = wrapAna ana

adjointRefold :: (Adjunction l m d c, Adjunction m r c d, Fixable d, IsFunctor d f, Obj c a, Obj c b, Obj d (l a), Obj d (r b))
              => c (m (f (r b))) b -> c a (m (f (l a))) -> c a b
adjointRefold = wrapHylo hylo



newtype FixF f = In { out :: f (FixF f) }

instance Fixable (->) where
  type Fix (->) f = FixF f
  fix = In
  unfix = out

instance FunctorDefault (->) where
  type IsFunctor' (->) f = P.Functor f
  map = P.fmap


instance Functor Identity (->) (->) where
  fmap = P.fmap

instance Adjunction Identity Identity (->) (->) where
  leftAdjunct f = Identity #. f .# Identity
  rightAdjunct f = runIdentity #. f .# runIdentity

wrapFold
  :: ((fb -> c) -> fixf -> c)
  -> ((fb -> c) -> fixf -> c)
wrapFold cata' alg = wrapCata ((Identity #.) #. cata' .# (runIdentity #.)) (alg .# runIdentity) .# Identity

wrapUnfold
  :: ((c -> fa) -> c -> fixf)
  -> ((c -> fa) -> c -> fixf)
wrapUnfold ana' coalg = runIdentity #. wrapAna ((.# runIdentity) . ana' . (.# Identity)) (Identity #. coalg)

wrapRefold
  :: ((fb -> b) -> (a -> fa) -> a -> b)
  -> ((fb -> b) -> (a -> fa) -> a -> b)
wrapRefold hylo' alg coalg = wrapHylo hylo'' (alg .# runIdentity) (Identity #. coalg)
  where
    hylo'' alg' coalg' = Identity #. hylo' (runIdentity #. alg') (coalg' .# Identity) .# runIdentity

fold :: P.Functor f => (f b -> b) -> FixF f -> b
fold = wrapFold cata

unfold :: P.Functor f => (a -> f a) -> a -> FixF f
unfold = wrapUnfold ana

refold :: P.Functor f => (f b -> b) -> (a -> f a) -> a -> b
refold = wrapRefold hylo


instance Functor ((->) e) (->) (->) where
  fmap = P.fmap
instance Functor ((,) e) (->) (->) where
  fmap = P.fmap

instance Adjunction ((,) e) ((->) e) (->) (->) where
  leftAdjunct f a e = f (e, a)
  rightAdjunct f (e, a) = index (f a) e

wrapParamFold
  :: ((fb -> (a -> c)) -> fixf -> (a -> c))
  -> (a -> fb -> c) -> a -> fixf -> c
wrapParamFold cata' = curry . wrapCata cata' . uncurry

wrapParamUnfold
  :: (((a, b) -> fb) -> (a, b) -> fixf)
  -> ((b -> a -> fb) -> b -> a -> fixf)
wrapParamUnfold = wrapAna

paramFold :: P.Functor f => (a -> f (a -> b) -> b) -> a -> FixF f -> b
paramFold = wrapParamFold cata

paramUnfold :: P.Functor f => (b -> a -> f (a, b)) -> b -> a -> FixF f
paramUnfold = wrapParamUnfold ana


instance Monad m => Fixable (Kleisli m) where
  type Fix (Kleisli m) f = FixF f
  fix = arr In
  unfix = arr out

instance Monad m => FunctorDefault (Kleisli m) where
  type IsFunctor' (Kleisli m) f = Traversable f
  map = Kleisli . mapM . runKleisli

newtype KleisliAdjF (m :: Type -> Type) a = KAF { unKAF :: a }
instance Monad m => Functor (KleisliAdjF m) (->) (Kleisli m) where
  fmap f = Kleisli (return . KAF . f . unKAF)

newtype KleisliAdjG m a = KAG { unKAG :: m a }
instance Monad m => Functor (KleisliAdjG m) (Kleisli m) (->) where
  fmap (Kleisli f) = KAG . (>>= f) . unKAG

instance Monad m => Adjunction (KleisliAdjF m) (KleisliAdjG m) (Kleisli m) (->) where
  leftAdjunct (Kleisli f) = KAG . f . KAF
  rightAdjunct f = Kleisli $ unKAG . f . unKAF

wrapUnfoldM
  :: Monad m
  => (Kleisli m a fla -> Kleisli m a fixf)
  -> (a -> m fla) -> a -> m fixf
wrapUnfoldM ana' coalg = unKAG #. wrapAna ((.# unKAF) #. ana' .# (.# KAF)) (KAG #. coalg)

unfoldM :: (Traversable f, Monad m) => (b -> m (f b)) -> b -> m (FixF f)
unfoldM = wrapUnfoldM ana


instance Comonad w => Fixable (Cokleisli w) where
  type Fix (Cokleisli w) f = FixF f
  fix = arr In
  unfix = arr out

instance Comonad w => FunctorDefault (Cokleisli w) where
  type IsFunctor' (Cokleisli w) t = Distributive t
  map = Cokleisli #. cotraverse .# runCokleisli

newtype CokleisliAdjF (w :: Type -> Type) a = CoKAF { unCoKAF :: a }
instance Comonad w => Functor (CokleisliAdjF w) (->) (Cokleisli w) where
  fmap f = Cokleisli ((CoKAF #. f .# unCoKAF) . extract)

newtype CokleisliAdjG w a = CoKAG { unCoKAG :: w a }
instance Comonad w => Functor (CokleisliAdjG w) (Cokleisli w) (->) where
  fmap (Cokleisli f) = CoKAG #. (=>> f) .# unCoKAG

instance Comonad w => Adjunction (CokleisliAdjG w) (CokleisliAdjF w) (->) (Cokleisli w) where
  leftAdjunct f = Cokleisli $ CoKAF #. f .# CoKAG
  rightAdjunct (Cokleisli f) = unCoKAF #. f .# unCoKAG

foldW :: (Distributive f, Comonad w) => (w (f b) -> b) -> w (FixF f) -> b
foldW = wrapFoldW cata

wrapFoldW
  :: Comonad w
  => (Cokleisli w frb b -> Cokleisli w fixf b)
  -> (w frb -> b) -> w fixf -> b
wrapFoldW cata' alg = wrapCata ((CoKAF #.) #. cata' .# (unCoKAF #.)) (alg .# unCoKAG) .# CoKAG


newtype Nat f g = Nat { unNat :: forall a. f a -> g a }
type instance Obj Nat a = P.Functor a
instance Category Nat where
  id = Nat id
  Nat f . Nat g = Nat (f . g)

instance FunctorDefault Nat where
  type IsFunctor' Nat f = FunctorLiftingObj f Nat Nat
  map = fmap

instance Functor (Ran j) Nat Nat where
  fmap (Nat f) = Nat $ \(Ran h) -> Ran (f . h)
instance Functor (Lan j) Nat Nat where
  fmap (Nat f) = Nat $ \(Lan h g) -> Lan h (f g)

newtype Pre j f a = Pre { unPre :: f (j a) }
instance (P.Functor f, P.Functor j) => P.Functor (Pre j f) where
  fmap f (Pre fja) = Pre (P.fmap (P.fmap f) fja)
instance P.Functor j => Functor (Pre j) Nat Nat where
  fmap (Nat f) = Nat $ \(Pre fja) -> Pre (f fja)

instance P.Functor j => Adjunction (Pre j) (Ran j) Nat Nat where
  leftAdjunct (Nat n) = Nat $ \fb -> Ran $ \bjx -> n (Pre (bjx <$> fb))
  rightAdjunct (Nat n) = Nat $ \(Pre fja) -> runRan (n fja) id

instance P.Functor j => Adjunction (Lan j) (Pre j) Nat Nat where
  leftAdjunct (Nat n) = Nat $ Pre . n . Lan id
  rightAdjunct (Nat n) = Nat $ \(Lan jyx ay) -> jyx <$> unPre (n ay)

newtype FixH f a = InH { outH :: f (FixH f) a }
instance P.Functor (f (FixH f)) => P.Functor (FixH f) where
  fmap f = InH . P.fmap f . outH

instance Fixable Nat where
  type Fix Nat f = FixH f
  fix = Nat InH
  unfix = Nat outH

wrapGNestedFold
  :: (P.Functor frb, P.Functor b, P.Functor j)
  => (Nat frb (Ran j b) -> Nat fixf (Ran j b))
  -> (Nat (Pre j frb) b -> Nat (Pre j fixf) b)
wrapGNestedFold = wrapCata

wrapGNestedUnfold
  :: (P.Functor fla, P.Functor a, P.Functor j)
  => (Nat (Lan j a) fla -> Nat (Lan j a) fixf)
  -> (Nat a (Pre j fla) -> Nat a (Pre j fixf))
wrapGNestedUnfold = wrapAna

wrapGNestedRefold
  :: (P.Functor frb, P.Functor b, P.Functor fla, P.Functor a, P.Functor j)
  => (Nat frb (Ran j b) -> Nat (Lan j a) fla -> Nat (Lan j a) (Ran j b))
  -> (Nat (Pre j frb) b -> Nat a (Pre j fla) -> Nat a b)
wrapGNestedRefold = wrapHylo

gnestedFold :: (FunctorLiftingObj f Nat Nat, P.Functor b, P.Functor j)
            => (forall y. f (Ran j b) (j y) -> b y) -> FixH f (j x) -> b x
gnestedFold alg = unNat (wrapGNestedFold cata (Nat $ alg . unPre)) . Pre

gnestedUnfold :: (FunctorLiftingObj f Nat Nat, P.Functor a, P.Functor j)
              => (forall y. a y -> f (Lan j a) (j y)) -> a x -> FixH f (j x)
gnestedUnfold coalg = unPre . unNat (wrapGNestedUnfold ana (Nat $ Pre . coalg))

gnestedRefold :: (FunctorLiftingObj f Nat Nat, P.Functor a, P.Functor b, P.Functor j)
              => (forall y. f (Ran j b) (j y) -> b y) -> (forall y. a y -> f (Lan j a) (j y)) -> a x -> b x
gnestedRefold alg coalg = unNat (wrapGNestedRefold hylo (Nat $ alg . unPre) (Nat $ Pre . coalg))

newtype Rsh j b x = Rsh { unRsh :: (x -> j) -> b } deriving P.Functor
ran2rsh :: Nat (Ran (Const a) (Const b)) (Rsh a b)
ran2rsh = Nat $ \(Ran f) -> Rsh (\g -> getConst (f (Const . g)))

nestedFold :: FunctorLiftingObj f Nat Nat => (f (Rsh j b) j -> b) -> FixH f j -> b
nestedFold alg = getConst . gnestedFold (Const . alg . unNat (fmap ran2rsh) . P.fmap getConst) . P.fmap Const

data Lsh j a x = Lsh (j -> x) a deriving P.Functor
lsh2lan :: Nat (Lsh j a) (Lan (Const j) (Const a))
lsh2lan = Nat $ \(Lsh f a) -> Lan (f . getConst) (Const a)

nestedUnfold :: FunctorLiftingObj f Nat Nat => (a -> f (Lsh j a) j) -> a -> FixH f j
nestedUnfold coalg = P.fmap getConst . gnestedUnfold (P.fmap Const . unNat (fmap lsh2lan) . coalg . getConst) . Const

nestedRefold :: FunctorLiftingObj f Nat Nat => (f (Rsh j b) j -> b) -> (a -> f (Lsh j a) j) -> a -> b
nestedRefold alg coalg =
  getConst . gnestedRefold
      (Const . alg . unNat (fmap ran2rsh) . P.fmap getConst)
      (P.fmap Const . unNat (fmap lsh2lan) . coalg . getConst)
  . Const

data Pair a = Pair a a deriving P.Functor
data PerfectF r a = Zero a | Succ (r (Pair a)) deriving P.Functor
type Perfect = FixH PerfectF
instance Functor PerfectF Nat Nat where
  fmap n = Nat $ \case
    Zero a -> Zero a
    Succ t -> Succ $ unNat n t

instance Show a => Show (Perfect a) where
  show = nestedFold f . P.fmap show where
    f (Zero i) = i
    f (Succ (Rsh g)) = g (\(Pair a b) -> "(" ++ a ++ ", " ++ b ++ ")")

fromDepth :: Int -> Perfect Int
fromDepth = nestedUnfold f where
  f 0 = Zero 0
  f b = Succ (Lsh (\j -> Pair (2 * j) (2 * j + 1)) (b - 1))



class Comonad w => Coalgebra w a where
  coalgebra :: a -> w a
newtype CoEMMor (w :: Type -> Type) a b = CoEMMor { unCoEMMor :: a -> b }
  deriving Category via (->)
type instance Obj (CoEMMor w) a = Coalgebra w a

instance Fixable (CoEMMor w) where
  type Fix (CoEMMor w) f = FixF f
  fix = CoEMMor In
  unfix = CoEMMor out

class P.Functor f => Codist w f where
  codistMap :: (a -> w b) -> f a -> w (f b)
instance {-# INCOHERENT #-} (Coalgebra w a, Codist w f) => Coalgebra w (f a) where
  coalgebra = codistMap coalgebra
instance Comonad w => Coalgebra w (w a) where
  coalgebra = duplicate
instance FunctorDefault (CoEMMor w) where
  type IsFunctor' (CoEMMor w) f = Codist w f
  map (CoEMMor f) = CoEMMor (map f)

newtype CofreeCoalg w a = CofreeCoalg { unCofreeCoalg :: w a } deriving P.Functor
instance Comonad w => Codist w (CofreeCoalg w) where
  codistMap coalg (CofreeCoalg wa) = CofreeCoalg . coalg <$> wa
instance Comonad w => Coalgebra w (CofreeCoalg w a) where
  coalgebra (CofreeCoalg wa) = extend CofreeCoalg wa
instance Comonad w => Functor (CofreeCoalg w) (->) (CoEMMor w) where
  fmap f = CoEMMor (P.fmap f)

newtype ForgetCoalg w a = ForgetCoalg { unForgetCoalg :: a } deriving P.Functor
instance Comonad w => Functor (ForgetCoalg w) (CoEMMor w) (->) where
  fmap (CoEMMor f) = P.fmap f

instance Comonad w => Adjunction (ForgetCoalg w) (CofreeCoalg w) (->) (CoEMMor w) where
  leftAdjunct f = CoEMMor $ CofreeCoalg . P.fmap (f . ForgetCoalg) . coalgebra
  rightAdjunct (CoEMMor f) = extract . unCofreeCoalg . f . unForgetCoalg

wrapGCata
  :: (Comonad w, Coalgebra w frb)
  => ((frb -> w b) -> fixf -> w b)
  -> (frb -> b) -> fixf -> b
wrapGCata cata' alg = wrapCata (CoEMMor #. (CofreeCoalg #.) #. cata' .# (unCofreeCoalg #.) .# unCoEMMor) (alg .# unForgetCoalg) .# ForgetCoalg

gcata :: (Comonad w, Codist w f) => (f (w b) -> b) -> FixF f -> b
gcata = wrapGCata cata

instance P.Functor f => Codist (Cofree f) f where
  codistMap f fa = (extract . f <$> fa) :< (codistMap id . unwrap . f <$> fa)

wrapHisto
  :: (P.Functor f, Coalgebra (Cofree f) frb)
  => ((frb -> Cofree f b) -> fixf -> Cofree f b)
  -> (frb -> b) -> fixf -> b
wrapHisto = wrapGCata

histo :: P.Functor f => (f (Cofree f b) -> b) -> FixF f -> b
histo = wrapHisto cata


class Monad m => Algebra m a where
  algebra :: m a -> a
newtype EMMor (m :: Type -> Type) a b = EMMor { unEMMor :: a -> b }
  deriving Category via (->)
type instance Obj (EMMor m) a = Algebra m a

instance Fixable (EMMor m) where
  type Fix (EMMor m) f = FixF f
  fix = EMMor In
  unfix = EMMor out

class P.Functor f => Dist m f where
  distMap :: (m a -> b) -> m (f a) -> f b
instance {-# INCOHERENT #-} (Algebra m a, Dist m f) => Algebra m (f a) where
  algebra = distMap algebra
instance Monad m => Algebra m (m a) where
  algebra = join
instance FunctorDefault (EMMor m) where
  type IsFunctor' (EMMor m) f = Dist m f
  map (EMMor f) = EMMor (map f)

newtype FreeAlg m a = FreeAlg { unFreeAlg :: m a } deriving P.Functor
instance Monad m => Dist m (FreeAlg m) where
  distMap mab mfa = FreeAlg $ (\(FreeAlg ma) -> mab ma) <$> mfa
instance Monad m => Algebra m (FreeAlg m a) where
  algebra ma = FreeAlg $ ma >>= unFreeAlg
instance Monad m => Functor (FreeAlg m) (->) (EMMor m) where
  fmap f = EMMor (P.fmap f)

newtype ForgetAlg m a = ForgetAlg { unForgetAlg :: a } deriving P.Functor
instance Functor (ForgetAlg m) (EMMor m) (->) where
  fmap (EMMor f) = P.fmap f

instance Monad m => Adjunction (FreeAlg m) (ForgetAlg m) (EMMor m) (->) where
  leftAdjunct (EMMor f) = ForgetAlg . f . FreeAlg . return
  rightAdjunct f = EMMor $ algebra . P.fmap (unForgetAlg . f) . unFreeAlg

wrapGAna
  :: (Monad m, Algebra m fla)
  => ((m a -> fla) -> m a -> fixf)
  -> (a -> fla) -> a -> fixf
wrapGAna ana' coalg = unForgetAlg #. wrapAna (EMMor #. (.# unFreeAlg) #. ana' .# (.# FreeAlg) .# unEMMor) (ForgetAlg #. coalg)

gana :: (Monad m, Dist m f) => (b -> f (m b)) -> b -> FixF f
gana = wrapGAna ana

instance P.Functor f => Dist (Free f) f where
  distMap f (Pure fa) = f . return <$> fa
  distMap f (Free fmfa) = f . wrap . distMap id <$> fmfa

futu :: P.Functor f => (b -> f (Free f b)) -> b -> FixF f
futu = gana



newtype BiEMMor (m :: Type -> Type) (w :: Type -> Type) a b = BiEMMor { unBiEMMor :: a -> b }
  deriving Category via (->)
type instance Obj (BiEMMor m w) a = (Algebra m a, Coalgebra w a)

instance Fixable (BiEMMor m w) where
  type Fix (BiEMMor m w) f = FixF f
  fix = BiEMMor In
  unfix = BiEMMor out

instance FunctorDefault (BiEMMor m w) where
  type IsFunctor' (BiEMMor m w) f = (Dist m f, Codist w f)
  map (BiEMMor f) = BiEMMor (map f)

newtype FreeBialg m w a = FreeBialg { unFreeBialg :: m a } deriving P.Functor
instance Monad m => Dist m (FreeBialg m w) where
  distMap mab mfa = FreeBialg $ (\(FreeBialg ma) -> mab ma) <$> mfa
instance Monad m => Algebra m (FreeBialg m w a) where
  algebra ma = FreeBialg $ ma >>= unFreeBialg
instance Comonad w => Coalgebra w (FreeBialg m w a) where
  coalgebra = error "not needed?"
instance (Monad m, Comonad w) => Functor (FreeBialg m w) (->) (BiEMMor m w) where
  fmap f = BiEMMor (P.fmap f)

newtype CofreeBialg m w a = CofreeBialg { unCofreeBialg :: w a } deriving P.Functor
instance Comonad w => Codist w (CofreeBialg m w) where
  codistMap coalg (CofreeBialg wa) = CofreeBialg . coalg <$> wa
instance Comonad w => Coalgebra w (CofreeBialg m w a) where
  coalgebra (CofreeBialg wa) = extend CofreeBialg wa
instance Monad m => Algebra m (CofreeBialg m w a) where
  algebra = error "not needed?"
instance (Monad m, Comonad w) => Functor (CofreeBialg m w) (->) (BiEMMor m w) where
  fmap f = BiEMMor (P.fmap f)

newtype ForgetBialg m w a = ForgetBialg { unForgetBialg :: a } deriving P.Functor
instance Functor (ForgetBialg m w) (BiEMMor m w) (->) where
  fmap (BiEMMor f) = P.fmap f

instance (Monad m, Comonad w) => Adjunction (FreeBialg m w) (ForgetBialg m w) (BiEMMor m w) (->) where
  leftAdjunct (BiEMMor f) = ForgetBialg . f . FreeBialg . return
  rightAdjunct f = BiEMMor $ algebra . P.fmap (unForgetBialg . f) . unFreeBialg

instance (Monad m, Comonad w) => Adjunction (ForgetBialg m w) (CofreeBialg m w) (->) (BiEMMor m w) where
  leftAdjunct f = BiEMMor $ CofreeBialg . P.fmap (f . ForgetBialg) . coalgebra
  rightAdjunct (BiEMMor f) = extract . unCofreeBialg . f . unForgetBialg

-- This is essentially the same as ghylo from recursion-schemes, but the unneeded and
-- impossible Coalgebra w (FreeBialg m w a) and Algebra m (CofreeBialg m w a)
-- show that the theory is not quite right
ghylo :: forall m f w a b. (Monad m, Dist m f, Comonad w, Codist w f)
      => (f (w b) -> b) -> (a -> f (m a)) -> a -> b
ghylo alg coalg = adjointRefold alg' coalg'
  where
    alg' :: ForgetBialg m w (f (CofreeBialg m w b)) -> b
    alg' = alg . P.fmap unCofreeBialg . unForgetBialg
    coalg' :: a -> ForgetBialg m w (f (FreeBialg m w a))
    coalg' = ForgetBialg . P.fmap FreeBialg . coalg

-- wrapGHylo hylo doesn't work
wrapGHylo
  :: forall m w fla frb a b. (Algebra m fla, Coalgebra w fla, Algebra m frb, Coalgebra w frb)
  => ((frb -> w b) -> (m a -> fla) -> m a -> w b)
  -> (frb -> b) -> (a -> fla) -> a -> b
wrapGHylo cata' alg coalg = wrapHylo cata'' (alg .# unForgetBialg) (ForgetBialg #. coalg)
  where
    cata'' (BiEMMor alg') (BiEMMor coalg') =
      BiEMMor @m @w $ CofreeBialg @m @w #. cata' (unCofreeBialg #. alg') (coalg' .# FreeBialg @m @w) .# unFreeBialg


chrono :: P.Functor f => (f (Cofree f b) -> b) -> (a -> f (Free f a)) -> a -> b
chrono = ghylo

dyna :: P.Functor f => (f (Cofree f a) -> a) -> (c -> f c) -> c -> a
dyna a c = chrono a (P.fmap Pure . c)


data ListF a as = Nil | Cons a as
instance P.Functor (ListF Int) where
  fmap _ Nil = Nil
  fmap f (Cons a as) = Cons a (f as)

natural :: Int -> ListF Int Int
natural 0 = Nil
natural n = Cons n (n - 1)

takeCofree :: Int -> Cofree (ListF x) a -> [a]
takeCofree 0 _ = []
takeCofree _ (a :< Nil) = [a]
takeCofree n (a :< Cons _ c) = a : takeCofree (n - 1) c

catalan :: Int -> Integer
catalan = dyna coa natural where
    coa :: ListF Int (Cofree (ListF Int) Integer) -> Integer
    coa Nil = 1
    coa (Cons x table) = sum $ zipWith (*) xs (reverse xs)
        where xs = takeCofree x table


data Two = P0 | P1
type (:**:) :: (x -> x -> Type) -> (y -> y -> Type) -> (Two -> x) -> (Two -> y) -> Type
data (:**:) k0 k1 a b where
  (:**:) :: k0 (a 'P0) (b 'P0) -> k1 (a 'P1) (b 'P1) -> (k0 :**: k1) a b

type instance Obj (k0 :**: k1) a = (Obj k0 (a 'P0), Obj k1 (a 'P1))
instance (Category k0, Category k1) => Category (k0 :**: k1) where
  id = id :**: id
  (f0 :**: f1) . (g0 :**: g1) = (f0 . g0) :**: (f1 . g1)

newtype Prod a = Prod { unProd :: (a 'P0, a 'P1) }
instance Functor Prod ((->) :**: (->)) (->) where
  fmap (f :**: g) = Prod . (f *** g) . unProd

newtype Coprod a = Coprod { unCoprod :: Either (a 'P0) (a 'P1) }
instance Functor Coprod ((->) :**: (->)) (->) where
  fmap (f :**: g) = Coprod . (f +++ g) . unCoprod

newtype Diag a (b :: Two) = Diag { unDiag :: a }
instance Functor Diag (->) ((->) :**: (->)) where
  fmap f = (\(Diag a) -> Diag (f a)) :**: (\(Diag a) -> Diag (f a))

instance Adjunction Diag Prod ((->) :**: (->)) (->) where
  leftAdjunct (fab1 :**: fab2) = Prod . (fab1 . Diag &&& fab2 . Diag)
  rightAdjunct f = (fst . unProd . f . unDiag) :**: (snd . unProd . f . unDiag)

instance Adjunction Coprod Diag (->) ((->) :**: (->)) where
  leftAdjunct f = (Diag . f . Coprod . Left) :**: (Diag . f . Coprod . Right)
  rightAdjunct (fab1 :**: fab2) = (unDiag . fab1 ||| unDiag . fab2) . unCoprod

data Proj :: Type -> Type -> Two -> Type where
  Proj0 :: { unProj0 :: a } -> Proj a b 'P0
  Proj1 :: { unProj1 :: b } -> Proj a b 'P1


wrapMutu
  :: ((frb -> (a, b)) -> fixf -> (a, b))
  -> (frb -> a, frb -> b) -> (fixf -> a, fixf -> b)
wrapMutu cata' (alg0, alg1) =
  let f :**: g = wrapCata
                  (((Prod #. (Proj0 *** Proj1)) .) . cata' . (((unProj0 *** unProj1) .# unProd) .))
                  ((Proj0 . alg0 .# unDiag) :**: (Proj1 . alg1 .# unDiag))
  in (unProj0 . f .# Diag, unProj1 . g .# Diag)

mutu :: forall f a b. P.Functor f => (f (a, b) -> a, f (a, b) -> b) -> (FixF f -> a, FixF f -> b)
mutu = wrapMutu cata

mutuHist :: P.Functor f => (f (Cofree f (a, b)) -> a, f (Cofree f (a, b)) -> b) -> (FixF f -> a, FixF f -> b)
mutuHist = wrapMutu (wrapHisto cata)


wrapComutu
  :: ((Either a b -> fla) -> Either a b -> fixf)
  -> (a -> fla, b -> fla) -> (a -> fixf, b -> fixf)
wrapComutu ana' (coalg0, coalg1) =
  let f :**: g = wrapAna
                  ((. ((unProj0 +++ unProj1) . unCoprod)) . ana' . (. (Coprod . (Proj0 +++ Proj1))))
                  ((Diag #. coalg0 . unProj0) :**: (Diag #. coalg1 . unProj1))
  in (unDiag #. f . Proj0, unDiag #. g . Proj1)

comutu :: P.Functor f => (a -> f (Either a b), b -> f (Either a b)) -> (a -> FixF f, b -> FixF f)
comutu = wrapComutu ana

wrapApo
  :: (P.Functor f, fla ~ f (Either x fixf), fixf ~ FixF f)
  => ((Either a fixf -> fla) -> Either a fixf -> fixf)
  -> (a -> fla) -> a -> fixf
wrapApo ana' alg = fst $ wrapComutu ana' (alg, P.fmap Right . out)

apo :: P.Functor f => (a -> f (Either a (FixF f))) -> a -> FixF f
apo = wrapApo ana



newtype Slice y a b = Slice { unSlice :: a -> b }
  deriving Category via (->)
type instance Obj (Slice y) a = ?slice :: a -> y

newtype ForgetSlice y a = ForgetSlice { unForgetSlice :: a }
instance Functor (ForgetSlice y) (Slice y) (->) where
  fmap (Slice f) = ForgetSlice #. f .# unForgetSlice

newtype PairSlice y a = PairSlice { unPairSlice :: (a, y) }
instance Functor (PairSlice y) (->) (Slice y) where
  fmap f = Slice $ PairSlice #. first f .# unPairSlice

instance Adjunction (ForgetSlice y) (PairSlice y) (->) (Slice y) where
  leftAdjunct f = Slice $ PairSlice #. (f .# ForgetSlice &&& ?slice)
  rightAdjunct (Slice f) = (fst . (unPairSlice #. f)) .# unForgetSlice

wrapZygo
  :: forall nm frb y fixf b. IP nm (frb -> y)
  => ((frb -> (b, y)) -> fixf -> (b, y))
  -> (frb -> b) -> fixf -> b
wrapZygo cata' alg = let ?slice = ip @nm in
  wrapCata (Slice #. (PairSlice #.) #. cata' .# (unPairSlice #.) .# unSlice) (alg .# unForgetSlice) .# ForgetSlice

zygo :: P.Functor f => (f b -> b) -> (f (a, b) -> a) -> FixF f -> a
zygo alg = let ?alg = alg . P.fmap snd in wrapZygo @"alg" cata

zygo2 :: P.Functor f => (f b -> b) -> (f c -> c) -> (f ((a, b), c) -> a) -> FixF f -> a
zygo2 algb algc =
  let ?algb = algb . P.fmap (snd . fst)
      ?algc = algc . P.fmap snd
  in wrapZygo @"algb" $ wrapZygo @"algc" cata

data TreeF a t = Empty | Node t a t deriving P.Functor
isPerfect :: FixF (TreeF a) -> Bool
isPerfect = zygo d p
  where
    p :: TreeF a (Bool, Integer) -> Bool
    p Empty = True
    p (Node (pl, dl) _ (pr, dr)) = pl && pr && (dl == dr)
    d :: TreeF a Integer -> Integer
    d Empty = 0
    d (Node dl _ dr) = 1 + max dl dr


para :: P.Functor f => (f (a, FixF f) -> a) -> FixF f -> a
para = zygo In

wrapPara
  :: (P.Functor f, frb ~ f (x, fixf), fixf ~ FixF f)
  => ((frb -> (a, fixf)) -> fixf -> (a, fixf))
  -> (frb -> a) -> fixf -> a
wrapPara = let ?alg = In . P.fmap snd in wrapZygo @"alg"