Another go at implementing polynomial functors a la David Spivak

Endofunctors.hs

{-# LANGUAGE GHC2021, GADTs, DataKinds #-}
{-# LANGUAGE TypeData #-}
module Endofunctors where

import Control.Comonad (Comonad(..))
import Control.Comonad.Cofree (Cofree(..))
import Control.Monad (join, ap, void)
import Control.Monad.Free (Free(..))
import Data.Bifunctor (first, second, bimap)
import Data.Functor.Day
import GHC.Generics hiding (R)



infixr 0 ~>
type p ~> q = forall y. p y -> q y

type K = Rec0
type Y = Par1


-- A monomial a*y^b
data Mono a b y = Mono a (b -> y) deriving Functor

bimapMono :: (a -> a') -> (b' -> b) -> Mono a b ~> Mono a' b'
bimapMono fa fb (Mono a by) = Mono (fa a) (by . fb)


-- Cofree comonad

cofreeHead :: Cofree p y -> y
cofreeHead (y :< _) = y

cofreeTail :: Cofree p y -> p (Cofree p y)
cofreeTail (_ :< pc) = pc

cofreeDuplicate :: Functor p => Cofree p y -> Cofree p (Cofree p y)
cofreeDuplicate c@(_ :< pc) = c :< fmap cofreeDuplicate pc

mapCofree :: Functor p => (p ~> q) -> Cofree p ~> Cofree q
mapCofree f (y :< pc) = y :< f (mapCofree f <$> pc)

runMultiple :: Comonad c => (c ~> p) -> c ~> Cofree p
runMultiple f c = extract c :< f (extend (runMultiple f) c)

microservice :: Free p `Day` Cofree q ~> Free (p `Day` q)
microservice (Day (Pure a) (b :< _) aby) = Pure (aby a b)
microservice (Day (Free pm) (_ :< qc) aby) = Free (Day pm qc $ \m c -> microservice $ Day m c aby)

swapCofree :: (Functor p, Functor q) => (q :.: p ~> p :.: q) -> q :.: Cofree p ~> Cofree p :.: q
swapCofree swap = Comp1 . go . unComp1
  where
    go qc =
      fmap cofreeHead qc :<
      fmap go (unComp1 . swap . Comp1 $ fmap cofreeTail qc)

spooling :: (Functor p, Functor q) => (q :.: p ~> p :.: q) -> Cofree (p :.: q) ~> Cofree p :.: Free q
spooling swap = Comp1 . go
  where
    go (y :< Comp1 pqc) =
      Pure y :<
      fmap (fmap Free . unComp1 . swapCofree swap . Comp1 . fmap go) pqc

spoolingTrav :: (Traversable t, Applicative f) => Cofree (f :.: t) ~> Cofree f :.: Free t
spoolingTrav = spooling (Comp1 . sequenceA . unComp1)



-- Sum and product helper functions

eitherSum :: (p y -> r) -> (q y -> r) -> (p :+: q) y -> r
eitherSum l _ (L1 p) = l p
eitherSum _ r (R1 p) = r p

bimapSum :: (p ~> p') -> (q ~> q') -> p :+: q ~> p' :+: q'
bimapSum fp _ (L1 p) = L1 (fp p)
bimapSum _ fq (R1 q) = R1 (fq q)

prodFst :: p :*: q ~> p
prodFst (p :*: _) = p

prodSnd :: p :*: q ~> q
prodSnd (_ :*: q) = q

bimapProd :: (p ~> p') -> (q ~> q') -> p :*: q ~> p' :*: q'
bimapProd fp fq (p :*: q) = fp p :*: fq q



-- Parallel or Dirichlet product as Day convolution of product

swapDay :: p `Day` q ~> q `Day` p
swapDay (Day p q f) = Day q p (flip f)

bimapDay :: (p ~> p') -> (q ~> q') -> p `Day` q ~> p' `Day` q'
bimapDay fp fq (Day p q f) = Day (fp p) (fq q) f

dualDayMono :: Mono a b `Day` Mono b a ~> Y
dualDayMono (Day (Mono a by) (Mono b ay) f) = Par1 (f (by b) (ay a))



-- Composition

bimapComp :: Functor p' => (p ~> p') -> (q ~> q') -> p :.: q ~> p' :.: q'
bimapComp fp fq = Comp1 . fmap fq . fp . unComp1

dayAsComp :: (Functor p, Functor q) => p `Day` q ~> p :.: q
dayAsComp (Day p q f) = Comp1 (fmap (\py -> fmap (f py) q) p)

duoidal :: (p :.: p') `Day` (q :.: q') ~> (p `Day` q) :.: (p' `Day` q')
duoidal (Day (Comp1 pp) (Comp1 qq) f) = Comp1 (Day pp qq (\p q -> Day p q f))


-- Endofunctor exponentiation

newtype (q ^ p) y where
  Exp :: { unExp :: forall x. (y -> x) -> p x -> q x } -> (q ^ p) y

deriving instance Functor (q ^ p)

dimapExp :: (p' ~> p) -> (q ~> q') -> q ^ p ~> q' ^ p'
dimapExp fp fq (Exp f) = Exp (\ya -> fq . f ya . fp)

applyExp :: ((q ^ p) :*: p) ~> q
applyExp = uncurryExp id

unapplyExp :: Functor q => q ~> ((q :*: p) ^ p)
unapplyExp = curryExp id

curryExp :: Functor o => (o :*: p) ~> q -> o ~> (q ^ p)
curryExp f o = Exp (\ya p -> f (fmap ya o :*: p))

uncurryExp :: o ~> (q ^ p) -> (o :*: p) ~> q
uncurryExp f (o :*: p) = case f o of Exp g -> g id p

composeExp :: (q ^ p) :*: (r ^ q) ~> r ^ p
composeExp (Exp pq :*: Exp qr) = Exp (\yx -> qr yx . pq yx)

liftProd0 :: Monad t => U1 ~> t ^ Y
liftProd0 U1 = Exp (\_ (Par1 a) -> pure a)

liftProd2 :: (Monad t, Functor p, Functor q) => (t ^ p) :*: (t ^ q) ~> t ^ (p :.: q)
liftProd2 (Exp pt :*: Exp qt) = Exp
  (\yx (Comp1 pq) -> join $ pt (pure . yx) $ fmap (qt yx) pq)

liftComp0 :: Monad t => Y ~> t ^ U1
liftComp0 (Par1 y) = Exp (\f U1 -> pure (f y))

liftComp2 :: (Monad t, Functor p, Functor q) => (t ^ p) :.: (t ^ q) ~> t ^ (p :*: q)
liftComp2 (Comp1 (Exp k)) = Exp
  (\yx (p :*: q) -> join $ k (\(Exp f) -> f yx q) (fmap pure p))

instance (Monad t, Functor p) => Applicative (t ^ p) where
  pure = dimapExp (const U1) id . liftComp0 . Par1
  (<*>) = ap
instance (Monad t, Functor p) => Monad (t ^ p) where
  ta >>= atb = (dimapExp (\p -> p :*: p) id . liftComp2 . Comp1) (fmap atb ta)

squared :: (t ^ K Bool) a -> (t a, t a)
squared (Exp k) = (k id (K1 False), k id (K1 True))

unsquared :: Functor t => (t a, t a) -> (t ^ K Bool) a
unsquared (l, r) = Exp (\f (K1 b) -> if b then fmap f r else fmap f l)

-- join for exponentiating by a constant is mapping join over the diagonal
testSquared :: (Monad t, Eq (t a)) => (t (t a, t a), t (t a, t a)) -> Bool
testSquared inp =
  squared (join $ unsquared (bimap (fmap unsquared) (fmap unsquared) inp))
  ==
  bimap (fst =<<) (snd =<<) inp



-- Dirichlet closure [p, q] = DClo p q

data DClo p q y where
  DClo :: (forall a b. (y -> a -> b) -> p a -> q b) -> DClo p q y
  deriving Functor

dimapDClo :: (p' ~> p) -> (q ~> q') -> DClo p q ~> DClo p' q'
dimapDClo fp fq (DClo f) = DClo (\k -> fq . f k . fp)

applyDClo :: (DClo p q `Day` p) ~> q
applyDClo = uncurryDClo id

unapplyDClo :: Functor q => q ~> DClo p (q `Day` p)
unapplyDClo = curryDClo id

curryDClo :: Functor o => (o `Day` p) ~> q -> o ~> DClo p q
curryDClo f o = DClo (\k p -> f (Day o p k))

uncurryDClo :: o ~> DClo p q -> (o `Day` p) ~> q
uncurryDClo f (Day o p g) = case f o of DClo h -> h g p

idDClo :: Functor p => Y ~> DClo p p
idDClo (Par1 y) = DClo (\f p -> f y <$> p)

composeDClo :: DClo q r `Day` DClo p q ~> DClo p r
composeDClo (Day (DClo qr) (DClo pq) f) = DClo (\k -> qr (\b (c, a) -> k (f b c) a) . pq (,))

multDClo :: DClo p q `Day` DClo p' q' ~> DClo (p `Day` p') (q `Day` q')
multDClo (Day (DClo pq) (DClo pq') f) = DClo (\k (Day p p' g) -> Day (pq (,) p) (pq' (,) p') (\(b, b') (c, c') -> k (f b c) (g b' c')))

dCloAsComp :: Functor q => (->) a :.: q ~> DClo ((,) a) q
dCloAsComp (Comp1 f) = DClo (\k (a, b) -> (\y -> k y b) <$> f a)

compAsDClo :: DClo ((,) a) q ~> (->) a :.: q
compAsDClo (DClo f) = Comp1 (\a -> f const (a, ()))

dCloAsComp' :: Functor q => (,) a :.: q ~> DClo ((->) a) q
dCloAsComp' (Comp1 (a, q)) = DClo (\k f -> (\y -> k y (f a)) <$> q)


data HomCoalg p q where
  Hom :: (s -> DClo p q s) -> HomCoalg p q

liftHomCoalg :: Functor q => (p ~> q) -> HomCoalg p q
liftHomCoalg pq = Hom (\() -> DClo (\k -> fmap (k ()) . pq))

idHomCoalg :: Functor p => HomCoalg p p
idHomCoalg = liftHomCoalg id

compHomCoalg :: HomCoalg q r -> HomCoalg p q -> HomCoalg p r
compHomCoalg (Hom qr) (Hom pq) = Hom (\(s, t) -> composeDClo (Day (qr s) (pq t) (,)))

multHomCoalg :: HomCoalg p q -> HomCoalg p' q' -> HomCoalg (Day p p') (Day q q')
multHomCoalg (Hom pq) (Hom pq') = Hom (\(s, t) -> multDClo (Day (pq s) (pq' t) (,)))

lensHomCoalg :: ((s, p) -> (a, b -> (t, p))) -> HomCoalg (Mono s t) (Mono a b)
lensHomCoalg l = Hom (\p -> DClo (\f (Mono s ty) -> case l (s, p) of (a, btp) -> Mono a (\b -> let (t, p') = btp b in f p' (ty t))))

homCoalgLens :: HomCoalg (Mono s t) (Mono a b) -> (forall p. ((s, p) -> (a, b -> (t, p))) -> r) -> r
homCoalgLens (Hom f) k = k (\(s, p) -> case f p of DClo k' -> case k' (flip (,)) (Mono s id) of Mono a ty -> (a, ty))



-- Composition left coclosure, aka left kan extension
--  Lan q p   or   ⎡q⎤
--                 ⎣p⎦

data Lan q p y where
  Lan :: p b -> (q b -> y) -> Lan q p y

deriving instance (Functor p, Functor q) => Functor (Lan q p)

dimapLan :: (q' ~> q) -> (p ~> p') -> Lan q p ~> Lan q' p'
dimapLan fq fp (Lan p k) = Lan (fp p) (k . fq)

applyLan :: Functor p => Lan q (p :.: q) ~> p
applyLan = cocurryLan id

unapplyLan :: p ~> (Lan q p :.: q)
unapplyLan = uncocurryLan id

cocurryLan :: Functor o => p ~> (o :.: q) -> Lan q p ~> o
cocurryLan f (Lan p k) = k <$> unComp1 (f p)

uncocurryLan :: Lan q p ~> o -> p ~> (o :.: q)
uncocurryLan f p = Comp1 (f (Lan p id))

-- Lawvere theory of a monad
type Law m = Lan ([] :.: m) []
-- (forall y. [y] -> o [m y]) -> Law m y -> o y
-- (forall y. [y] -> [m y]) -> Law m y -> y
instance Monad m => Comonad (Law m) where
  extract = unPar1 . cocurryLan (Comp1 . Par1 . Comp1 . map return)
  duplicate = unComp1 . cocurryLan (\ys ->
    Comp1 (Comp1 (
      (\(Comp1 mys) ->
          (\(Comp1 mms) ->
            Comp1 (fmap join mms)
          ) <$> unComp1 (unapplyLan mys)
      ) <$> unComp1 (unapplyLan ys)
    )))


-- Adjunctions

-- Mono _ x -| _(x)
adjAtTo :: (Mono a x ~> p) -> a -> p x
adjAtTo f a = f (Mono a id)

adjAtFrom :: Functor p => (a -> p x) -> Mono a x ~> p
adjAtFrom f (Mono a xy) = xy <$> f a

-- _(1) -| K
adjKTo :: Functor p => (p () -> a) -> p ~> K a
adjKTo f p = K1 (f (void p))

adjKFrom :: (p ~> K a) -> p () -> a
adjKFrom f p = unK1 (f p)

-- Gamma -| Mono ()
newtype Gamma p = Gamma { unGamma :: forall a. p a -> a }

adjGammaTo :: (p ~> Mono () b) -> b -> Gamma p
adjGammaTo f b = Gamma $ \p -> case f p of Mono () by -> by b

adjGammaFrom :: (b -> Gamma p) -> p ~> Mono () b
adjGammaFrom f p = Mono () (\b -> unGamma (f b) p)

-- _(1) x Gamma -| Mono
adjMonoTo :: (p ~> Mono a b) -> (p () -> a, b -> Gamma p)
adjMonoTo f = (\p -> case f p of Mono a _ -> a, \b -> Gamma $ \p -> case f p of Mono _ by -> by b)

adjMonoFrom :: Functor p => (p () -> a, b -> Gamma p) -> p ~> Mono a b
adjMonoFrom (f, g) p = Mono (f (void p)) (\b -> unGamma (g b) p)


-- Systems

-- also: p-coalgebra
type System s p = Mono s s ~> p

fromSystem :: System s p -> s -> p s
fromSystem f s = f (Mono s id)

toSystem :: Functor p => (s -> p s) -> System s p
toSystem f (Mono s sy) = sy <$> f s

data Stream a = Str a (Stream a) deriving Show

runMono :: System s (Mono a b) -> s -> Stream b -> Stream a
runMono sys s (Str b bs) = case fromSystem sys s of
  Mono a by -> Str a (runMono sys (by b) bs)

seqMono :: System s0 (Mono a b) -> System s1 (Mono b c) -> System (s0, s1) (Mono a c)
seqMono s0 s1 = combineDayMonoSeq . bimapDay s0 s1 . splitDayMono

parMono :: System s0 (Mono a b) -> System s1 (Mono c d) -> System (s0, s1) (Mono (a, c) (b, d))
parMono s0 s1 = combineDayMono . bimapDay s0 s1 . splitDayMono

altMono :: System s0 (Mono a b) -> System s1 (Mono c d) -> System (s0, s1) (Mono (a, c) (Either b d))
altMono sys0 sys1 (Mono (s0, s1) f) = combineProdMono (sys0 (Mono s0 (\s0' -> f (s0', s1))) :*: sys1 (Mono s1 (\s1' -> f (s0, s1'))))

flipflop :: System Bool (Mono Bool ())
flipflop = toSystem (\b -> Mono b (const (not b)))

counter :: System Int (Mono Int Bool)
counter = toSystem (\i -> Mono i (\b -> if b then i + 1 else i))

adder :: System Int (Mono String Int)
adder = toSystem (\i -> Mono (show i) (i +))


-- State machines with internal state

data SM p = forall s. SM { transition :: System s p, initialState :: s }

mapSM :: p ~> q -> SM p -> SM q
mapSM pq (SM f s) = SM (pq . f) s

runSM :: SM p -> p (SM p)
runSM = runMapSM id

runMapSM :: (SM p -> a) -> SM p -> p a
runMapSM g (SM f s) = f (Mono s (g . SM f))

runMultipleSM :: SM p -> Cofree p (SM p)
runMultipleSM sm = sm :< runMapSM runMultipleSM sm

basicSM :: Functor p => (s -> p s) -> s -> SM p
basicSM f = SM (\(Mono s g) -> g <$> f s)

seqSM :: Functor p => SM p -> SM q -> SM (p :.: q)
seqSM (SM fp sp) (SM fq sq) = SM (bimapComp fp fq . splitCompMono) (sp, sq)

parSM :: SM p -> SM q -> SM (p `Day` q)
parSM (SM fp sp) (SM fq sq) = SM (bimapDay fp fq . splitDayMono) (sp, sq)

altSM :: SM p -> SM q -> SM (p :*: q)
altSM (SM fp sp) (SM fq sq) = SM (bimapProd fp fq . splitProdMono) (sp, sq)


newtype Mealy i o y = Mealy { unMealy :: i -> (o, y) }
  deriving Functor

combineCompMealy :: Mealy a b :.: Mealy b c ~> Mealy a c
combineCompMealy (Comp1 (Mealy abm)) = Mealy $ \a -> case abm a of (b, Mealy bcy) -> bcy b

combineDayMealy :: Mealy a b `Day` Mealy c d ~> Mealy (a, c) (b, d)
combineDayMealy (Day (Mealy aby) (Mealy cdy) f) = Mealy $
  \(a, c) -> case (aby a, cdy c) of
    ((b, ya), (d, yc)) -> ((b, d), f ya yc)

combineProdMealy :: Mealy a b :*: Mealy c d ~> Mealy (Either a c) (Either b d)
combineProdMealy (Mealy aby :*: Mealy cdy) = Mealy $ either (first Left . aby) (first Right . cdy)


basicMeally :: (s -> i -> (o, s)) -> s -> SM (Mealy i o)
basicMeally f = basicSM (Mealy . f)

runMealy :: SM (Mealy i o) -> i -> (o, SM (Mealy i o))
runMealy = unMealy . runSM

runMultipleMealy :: (Monoid o, Foldable f) => SM (Mealy i o) -> f i -> (o, SM (Mealy i o))
runMultipleMealy c = second extract . foldl (\(o, _ :< Mealy ioc) -> first (o <>) . ioc) (mempty, runMultipleSM c)

seqMealy :: SM (Mealy a b) -> SM (Mealy b c) -> SM (Mealy a c)
seqMealy s0 s1 = mapSM combineCompMealy $ seqSM s0 s1

parMealy :: SM (Mealy a b) -> SM (Mealy c d) -> SM (Mealy (a, c) (b, d))
parMealy s0 s1 = mapSM combineDayMealy $ parSM s0 s1

altMealy :: SM (Mealy a b) -> SM (Mealy c d) -> SM (Mealy (Either a c) (Either b d))
altMealy s0 s1 = mapSM combineProdMealy $ altSM s0 s1



data MultiDay ps y where
  MDNil :: y -> MultiDay '[] y
  MDCons :: p x -> MultiDay ps (x -> y) -> MultiDay (p ': ps) y

data TA
data TB
data TC
type Plant = Mono TC (TA, TB)
type Controller = Mono TB TC
type TheSystem = Mono TC TA
wiring :: MultiDay '[Plant, Controller] ~> TheSystem
wiring (Mono tc taby `MDCons` (Mono tb tcy `MDCons` MDNil f)) = Mono tc (\ta -> f (tcy tc) (taby (ta, tb)))



combineSumMono :: Mono a b :+: Mono c d ~> Mono (Either a c) (b, d)
combineSumMono (L1 (Mono a by)) = Mono (Left a) (by . fst)
combineSumMono (R1 (Mono c dy)) = Mono (Right c) (dy . snd)

combineProdMono :: Mono a b :*: Mono c d ~> Mono (a, c) (Either b d)
combineProdMono (Mono a by :*: Mono c dy) = Mono (a, c) (either by dy)

uncombineProdMono :: Mono (a, c) (Either b d) ~> Mono a b :*: Mono c d
uncombineProdMono (Mono (a, c) bdy) = Mono a (bdy . Left) :*: Mono c (bdy . Right)

combineDayMono :: Mono a b `Day` Mono c d ~> Mono (a, c) (b, d)
combineDayMono (Day (Mono a by) (Mono c dy) f) = Mono (a, c) (\(b, d) -> f (by b) (dy d))

combineDayMonoSeq :: Mono a b `Day` Mono b c ~> Mono a c
combineDayMonoSeq (Day (Mono a by) (Mono b cy) f) = Mono a (f (by b) . cy)

combineCompMono :: Mono a b :.: Mono c d ~> Mono (a, b -> c) (b, d)
combineCompMono (Comp1 (Mono a bcd)) = Mono (a, \b -> case bcd b of Mono c _ -> c) (\(b, d) -> case bcd b of Mono _ dy -> dy d)

combineDCloMono :: a -> Mono a b `DClo` Mono c d ~> Mono (c, d -> b) d
combineDCloMono a (DClo f) = case f (,) (Mono a id) of Mono c dy -> Mono (c, snd . dy) (fst . dy)

combineLanMono :: Mono c b `Lan` Mono a b ~> Mono a c
combineLanMono (Lan (Mono a by) k) = Mono a (\c -> k (Mono c by))

splitCompMono :: Mono (a, c) (b, d) ~> Mono a b :.: Mono c d
splitCompMono (Mono (a, c) bdy) = Comp1 (Mono a (\b -> Mono c (\d -> bdy (b, d))))

splitDayMono :: Mono (a, c) (b, d) ~> Mono a b `Day` Mono c d
splitDayMono (Mono (a, c) f) = Day (Mono a id) (Mono c id) (curry f)

splitProdMono :: Mono (a, b) (a, b) ~> Mono a a :*: Mono b b
splitProdMono (Mono (a, b) f) = Mono a (\a' -> f (a', b)) :*: Mono b (\b' -> f (a, b'))


-- l f ~> g -> f ~> r g
-- l f ~> g -> f ~> RightAdjoint l g

data RightAdjoint l g a where
  RightAdjoint :: Functor f => l f ~> g -> f a -> RightAdjoint l g a

toRightAdjoint :: Functor f => l f ~> g -> f ~> RightAdjoint l g
toRightAdjoint lfg = RightAdjoint lfg

-- f ~> r g -> l f ~> g
-- LeftAdjoint r f a -> f ~> r g -> g a

newtype LeftAdjoint r f a = LeftAdjoint {
  getLeftAdjoint :: forall g. Functor g => f ~> r g -> g a
}

fromLeftAdjoint :: Functor g => f ~> r g -> LeftAdjoint r f ~> g
fromLeftAdjoint frg (LeftAdjoint k) = k frg


type Lens s t a b = Mono s t ~> Mono a b
lens :: (s -> a) -> (s -> b -> t) -> Lens s t a b
lens sa sbt (Mono s ty) = Mono (sa s) (ty . sbt s)

get :: Lens s t a b -> s -> a
get f s = case f (Mono s id) of Mono a _ -> a

set :: Lens s t a b -> s -> b -> t
set f s b = case f (Mono s id) of Mono _ ty -> ty b

lensCoprod :: Lens s t a b -> Lens s' t' a' b' -> (Mono s t :+: Mono s' t') ~> (Mono a b :+: Mono a' b')
lensCoprod l r = bimapSum l r

get' :: (Mono s t :+: Mono s' t') ~> (Mono a b :+: Mono a' b') -> (Either s s' -> Either a a')
get' f ss = case f (either (L1 . (`Mono` Left)) (R1 . (`Mono` Right)) ss) of
  L1 (Mono a _) -> Left a
  R1 (Mono a' _) -> Right a'

IsCountable.hs

{-# LANGUAGE GHC2021 #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
module IsCountable (IsCountable(..), pattern S, pattern Z, enumAll, absurd, Nat(..), pred) where

import Data.Kind (Type)
import Data.Void (Void)
import GHC.TypeNats (type (+))
import Prelude hiding (pred)
import Numeric.Natural (Natural)

class IsCountable s where
  -- | A countable set is either empty: @Pred s ~ 'Nothing@ or one bigger than another countable set, its predecessor: @Pred s ~ 'Just p@
  type Pred s :: Maybe Type
  fromPredCase :: PredCase s -> s
  -- | pattern match on elements of @s@
  toPredCase :: s -> PredCase s
  -- | An "if" based on the type of @Pred s@.
  tIf :: (Pred s ~ 'Nothing => r)
      -> (forall p. (Pred s ~ 'Just p, IsCountable p) => r)
      -> r

absurd :: (IsCountable s, Pred s ~ 'Nothing) => s -> x
absurd = \case {}

enumAll :: forall s. IsCountable s => [s]
enumAll = tIf @s [] (Z : (S <$> enumAll))

-- | Type level size, this is SLOW on larger types.
type Size s = SIZE (Pred s)
type family SIZE ps where
  SIZE 'Nothing = 0
  SIZE ('Just p) = 1 + Size p


data PredCase s where
  PredZ :: (Pred s ~ 'Just p, IsCountable p) => PredCase s
  PredS :: (Pred s ~ 'Just p, IsCountable p) => p -> PredCase s

convPredCase :: Pred a ~ Pred b => PredCase a -> PredCase b
convPredCase PredZ = PredZ
convPredCase (PredS p) = PredS p

pred :: IsCountable s
  => (forall p. (Pred s ~ 'Just p, IsCountable p) => r) -- ^ the extra element case
  -> (forall p. (Pred s ~ 'Just p, IsCountable p) => p -> r) -- ^ the lifted elements case
  -> s -> r
pred z _ Z = z
pred _ s (S p) = s p

pattern Z :: IsCountable s => (Pred s ~ 'Just p, IsCountable p) => s
pattern Z <- (toPredCase -> PredZ) where Z = fromPredCase PredZ

pattern S :: IsCountable s => (Pred s ~ 'Just p, IsCountable p) => p -> s
pattern S p <- (toPredCase -> PredS p) where S p = fromPredCase (PredS p)

{-# COMPLETE Z, S #-}


data Nat = NZ | NS Nat
data Fin n where
  Fz :: Fin (NS n)
  Fs :: Fin n -> Fin (NS n)

instance IsCountable (Fin NZ) where
  type Pred (Fin NZ) = 'Nothing
  fromPredCase = \case {}
  toPredCase = \case {}
  tIf z _ = z

instance IsCountable (Fin n) => IsCountable (Fin (NS n)) where
  type Pred (Fin (NS n)) = 'Just (Fin n)
  fromPredCase PredZ = Fz
  fromPredCase (PredS n) = Fs n
  toPredCase Fz = PredZ
  toPredCase (Fs n) = PredS n
  tIf _ s = s


instance IsCountable Void where
  type Pred Void = 'Nothing
  fromPredCase = \case {}
  toPredCase = \case {}
  tIf z _ = z

instance IsCountable () where
  type Pred () = 'Just Void
  fromPredCase PredZ = ()
  fromPredCase (PredS v) = case v of
  toPredCase () = PredZ
  tIf _ s = s

instance IsCountable Bool where
  type Pred Bool = 'Just ()
  fromPredCase PredZ = False
  fromPredCase (PredS ()) = True
  toPredCase False = PredZ
  toPredCase True = PredS ()
  tIf _ s = s

instance IsCountable Natural where
  type Pred Natural = 'Just Natural
  fromPredCase PredZ = 0
  fromPredCase (PredS n) = n + 1
  toPredCase 0 = PredZ
  toPredCase n = PredS (n - 1)
  tIf _ s = s

instance IsCountable a => IsCountable (Maybe a) where
  type Pred (Maybe a) = 'Just a
  fromPredCase PredZ = Nothing
  fromPredCase (PredS a) = Just a
  toPredCase = maybe PredZ PredS
  tIf _ s = s


type family SUM pa b where
  SUM 'Nothing b = Pred b
  SUM ('Just pa) b = 'Just (Either pa b)

instance (IsCountable a, IsCountable b) => IsCountable (Either a b) where
  type Pred (Either a b) = SUM (Pred a) b
  fromPredCase PredZ = tIf @a (Right Z) (Left Z)
  fromPredCase (PredS p) = tIf @a (Right . S) (either (Left . S) Right) $ p
  toPredCase (Left a) = case toPredCase a of PredZ -> PredZ; PredS pa -> PredS (Left pa)
  toPredCase (Right b) = tIf @a (convPredCase $ toPredCase b) (PredS $ Right b)
  tIf z s = tIf @a (tIf @b z s) s


-- type Fin2 = Bool
-- type Fin3 = Maybe Bool
-- type Fin4 = (Fin2, Fin2)
-- type Fin16 = (Fin4, Fin4)
-- type Fin256 = Fin4 -> Fin4
-- type Fin65536 = (Fin256, Fin256)
-- type Fin4294967296 = (Fin65536, Fin65536)

IsPoly.hs

{-# LANGUAGE GHC2021, LambdaCase, DataKinds, TypeFamilies, UndecidableInstances, AllowAmbiguousTypes, PatternSynonyms, ViewPatterns #-}
{-# LANGUAGE BlockArguments #-}
module IsPoly where

import Control.Comonad.Cofree (Cofree(..))
import Control.Monad.Free (Free(..))
import Data.Functor.Day
import Data.Kind (Type, Constraint)
import Data.Void (Void)
import Data.Type.Ord
import GHC.Generics
import GHC.TypeLits (Symbol, AppendSymbol, ConsSymbol, )
import qualified GHC.TypeLits as N
import Prelude hiding (pred)

import IsCountable
import Endofunctors


-- A polynomial is either empty, or it is a sum of a monomial a*y^b and another polynomial
data PolyCase p y where
  PCHead :: (Uncons p ~ 'Just '(a, b, q), IsPoly q) => a -> (b -> y) -> PolyCase p y
  PCTail :: (Uncons p ~ 'Just '(a, b, q), IsPoly q) => q y -> PolyCase p y

convCase :: Uncons p ~ Uncons q => PolyCase p y -> PolyCase q y
convCase (PCHead a by) = PCHead a by
convCase (PCTail qy) = PCTail qy


class Functor p => IsPoly p where
  type Uncons p :: Maybe (Type, Type, Type -> Type)
  constr :: PolyCase p y -> p y
  destr :: p y -> PolyCase p y
  tyRep :: TyRep p


pattern Head :: IsPoly p => (Uncons p ~ 'Just '(a, b, q), IsPoly q) => a -> (b -> y) -> p y
pattern Head a by <- (destr -> PCHead a by)
  where
    Head a by = constr (PCHead a by)

pattern Tail :: IsPoly p => (Uncons p ~ 'Just '(a, b, q), IsPoly q) => q y -> p y
pattern Tail qy <- (destr -> PCTail qy)
  where
    Tail qy = constr (PCTail qy)

{-# COMPLETE Head, Tail #-}


data TyRep p where
  TyN :: Uncons p ~ 'Nothing => TyRep p
  TyJ :: forall a b p' p. (Uncons p ~ 'Just '(a, b, p'), IsPoly p') => TyRep p

tyCase
  :: forall p r. IsPoly p
  => (Uncons p ~ 'Nothing => r)
  -> (forall a b p'. (Uncons p ~ 'Just '(a, b, p'), IsPoly p') => r)
  -> r
tyCase n j = case tyRep @p of
  TyN -> n
  TyJ -> j

convRep :: Uncons p ~ Uncons q => TyRep p -> TyRep q
convRep TyN = TyN
convRep TyJ = TyJ


-- Morphisms between polynomial functors, aka dependent lenses
data Poly p q where
  PNil :: (Uncons p ~ 'Nothing, IsPoly p, IsPoly q) => Poly p q
  P :: (Uncons p ~ 'Just '(a, b, p'), IsPoly p, IsPoly q) => (a -> q b) -> Poly p' q -> Poly p q

toPoly :: forall p q. (IsPoly p, IsPoly q) => p ~> q -> Poly p q
toPoly f = tyCase @p PNil (P (\a -> f (Head a id)) (toPoly (f . Tail)))

fromPoly :: Poly p q -> p ~> q
fromPoly PNil = absurd1
fromPoly (P aqb p'q) = \case
  Head a by -> by <$> aqb a
  Tail p' -> fromPoly p'q p'



type AllA c p = ALLA c (Uncons p)
type family ALLA (c :: Type -> Constraint) p :: Constraint where
  ALLA c 'Nothing = ()
  ALLA c ('Just '(a, b, p)) = (c a, AllA c p)

type AllB c p = ALLB c (Uncons p)
type family ALLB (c :: Type -> Constraint) p :: Constraint where
  ALLB c 'Nothing = ()
  ALLB c ('Just '(a, b, p)) = (c b, AllB c p)


fmapDefault :: IsPoly p => (x -> y) -> p x -> p y
fmapDefault f (Head a by) = Head a (f . by)
fmapDefault f (Tail qy) = Tail (fmapDefault f qy)

absurd1 :: (IsPoly p, Uncons p ~ 'Nothing) => p y -> r
absurd1 p = case destr p of



instance IsPoly (Mono a b) where
  type Uncons (Mono a b) = 'Just '(a, b, V1)
  constr (PCHead a by) = Mono a by
  constr (PCTail q) = case q of {}
  destr (Mono a by) = PCHead a by
  tyRep = TyJ


instance IsPoly V1 where
  type Uncons V1 = 'Nothing
  constr = \case {}
  destr = \case {}
  tyRep = TyN

instance IsPoly U1 where
  type Uncons U1 = 'Just '((), Void, V1)
  constr (PCHead () _) = U1
  constr (PCTail q) = case q of {}
  destr U1 = PCHead () (\case {})
  tyRep = TyJ

instance IsPoly (Rec0 a) where
  type Uncons (Rec0 a) = 'Just '(a, Void, V1)
  constr (PCHead a _) = K1 a
  constr (PCTail q) = case q of {}
  destr (K1 a) = PCHead a (\case {})
  tyRep = TyJ

instance IsPoly Par1 where
  type Uncons Par1 = 'Just '((), (), V1)
  constr (PCHead () by) = Par1 (by ())
  constr (PCTail q) = case q of {}
  destr (Par1 y) = PCHead () (const y)
  tyRep = TyJ

instance IsPoly Maybe where
  type Uncons Maybe = 'Just '((), Void, Par1)
  constr (PCHead () _) = Nothing
  constr (PCTail (Par1 y)) = Just y
  destr = maybe (PCHead () absurd) (PCTail . Par1)
  tyRep = TyJ

instance IsPoly (Either a) where
  type Uncons (Either a) = 'Just '(a, Void, Par1)
  constr (PCHead a _) = Left a
  constr (PCTail (Par1 y)) = Right y
  destr = either (\a -> PCHead a absurd) (PCTail . Par1)
  tyRep = TyJ

instance IsPoly ((,) a) where
  type Uncons ((,) a) = 'Just '(a, (), V1)
  constr (PCHead a by) = (a, by ())
  constr (PCTail q) = case q of {}
  destr (a, y) = PCHead a (const y)
  tyRep = TyJ

instance IsPoly [] where
  type Uncons [] = 'Just '((), Void, Y :*: [])
  constr (PCHead () _) = []
  constr (PCTail (Par1 y :*: ys)) = y : ys
  destr [] = PCHead () absurd
  destr (y:ys) = PCTail (Par1 y :*: ys)
  tyRep = TyJ

instance (IsPoly p, AllB IsCountable p) => IsPoly (Free p) where
  type Uncons (Free p) = 'Just '((), (), p :.: Free p)
  constr (PCHead () by) = Pure $ by ()
  constr (PCTail (Comp1 f)) = Free f
  destr (Pure y) = PCHead () (const y)
  destr (Free p) = PCTail (Comp1 p)
  tyRep = TyJ


(!!!) :: Cofree (Mono a b) y -> [b] -> y
(y :< _) !!! [] = y
(_ :< pcp) !!! (b:bs) = case pcp of
  Head _ bcp -> bcp b !!! bs
  Tail v -> case v of

instance IsPoly (Cofree (Mono a b)) where
  type Uncons (Cofree (Mono a b)) = 'Just '([b] -> a, [b], V1)
  constr (PCHead as bys) = bys [] :< Head (as []) (\b -> Head (as . (b:)) (bys . (b:)))
  constr (PCTail q) = case q of {}
  destr c@(_ :< pcp) = case pcp of
    Head a _ -> PCHead (const a) (c !!!)
    Tail v -> case v of
  tyRep = TyJ

instance (IsPoly p, AllB IsCountable p) => IsPoly (Cofree (K a :+: p)) where
  type Uncons (Cofree (K a :+: p)) = 'Just '(a, (), Y :*: (p :.: Cofree (K a :+: p)))
  constr (PCHead a by) = by () :< Head a absurd
  constr (PCTail (Par1 y :*: Comp1 p')) = y :< Tail (R1 p')
  destr (y :< L1 (K1 a)) = PCHead a (const y)
  destr (y :< R1 p) = PCTail (Par1 y :*: Comp1 p)
  tyRep = TyJ




type family SUM p q where
  SUM 'Nothing q = Uncons q
  SUM ('Just '(a, b, p)) q = 'Just '(a, b, p :+: q)

instance (IsPoly p, IsPoly q) => IsPoly (p :+: q) where
  type Uncons (p :+: q) = SUM (Uncons p) q
  constr (PCHead a by) = tyCase @p (R1 $ Head a by) (L1 $ Head a by)
  constr (PCTail p') = tyCase @p (R1 . Tail) (eitherSum (L1 . Tail) R1) p'
  destr (L1 p) = case p of
    Head a by -> PCHead a by
    Tail q -> PCTail (L1 q)
  destr (R1 q) = tyCase @p (convCase $ destr q) (PCTail (R1 q))
  tyRep = tyCase @p (tyCase @q TyN TyJ) TyJ



data Scale a (p :: Type -> Type) y = Scale a (p y) deriving Functor

type family SCALE a p where
  SCALE a 'Nothing = 'Nothing
  SCALE a1 ('Just '(a2, b, p)) = 'Just '((a1, a2), b, Scale a1 p)

instance IsPoly p => IsPoly (Scale a p) where
  type Uncons (Scale a p) = SCALE a (Uncons p)
  constr pc = tyCase @p (case pc of {}) $ case pc of
    PCHead (a0, a1) by -> Scale a0 (Head a1 by)
    PCTail (Scale a q') -> Scale a (Tail q')
  destr (Scale a0 p) = case p of
    Head a1 by -> PCHead (a0, a1) by
    Tail qy -> PCTail (Scale a0 qy)
  tyRep = tyCase @p TyN TyJ

mapScale :: (p ~> q) -> Scale a p ~> Scale a q
mapScale pq (Scale a p) = Scale a (pq p)


data IncDeg b p y = IncDeg (b -> y) (p y) deriving Functor

type family INCDEG b p where
  INCDEG a 'Nothing = 'Nothing
  INCDEG b1 ('Just '(a, b2, p)) = 'Just '(a, Either b1 b2, IncDeg b1 p)

instance IsPoly p => IsPoly (IncDeg b p) where
  type Uncons (IncDeg b p) = INCDEG b (Uncons p)
  constr pc = tyCase @p (case pc of {}) $ case pc of
    PCHead a by -> IncDeg (by . Left) (Head a (by . Right))
    PCTail (IncDeg by q') -> IncDeg by (Tail q')
  destr (IncDeg by0 p) = case p of
    Head a by1 -> PCHead a (either by0 by1)
    Tail qy -> PCTail (IncDeg by0 qy)
  tyRep = tyCase @p TyN TyJ

mapIncDeg :: (p ~> q) -> IncDeg a p ~> IncDeg a q
mapIncDeg pq (IncDeg by p) = IncDeg by (pq p)


newtype MultDeg b p y = MultDeg (p (b -> y)) deriving Functor

type family MULTDEG b p where
  MULTDEG a 'Nothing = 'Nothing
  MULTDEG b1 ('Just '(a, b2, p)) = 'Just '(a, (b1, b2), MultDeg b1 p)

instance IsPoly p => IsPoly (MultDeg b p) where
  type Uncons (MultDeg b p) = MULTDEG b (Uncons p)
  constr pc = tyCase @p (case pc of {}) $ case pc of
    PCHead a by -> MultDeg (Head a (\b0 b1 -> by (b1, b0)))
    PCTail (MultDeg q'by) -> MultDeg (Tail q'by)
  destr (MultDeg p) = case p of
    Head a by -> PCHead a (\(b0, b1) -> by b1 b0)
    Tail qy -> PCTail (MultDeg qy)
  tyRep = tyCase @p TyN TyJ

mapMultDeg :: (p ~> q) -> MultDeg a p ~> MultDeg a q
mapMultDeg pq (MultDeg p) = MultDeg (pq p)



type ProdStep a b p' q = Scale a (IncDeg b q) :+: (p' :*: q)

type family PROD p q where
  PROD 'Nothing q = 'Nothing
  PROD ('Just '(a, b, p')) q = Uncons (ProdStep a b p' q)

fromProdStep :: (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p') => ProdStep a b p' q y -> (p :*: q) y
fromProdStep (L1 (Scale a (IncDeg by q))) = Head a by :*: q
fromProdStep (R1 (p' :*: q)) = Tail p' :*: q

toProdStep :: (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p') => (p :*: q) y -> ProdStep a b p' q y
toProdStep (p :*: q) = case p of
  Head a by -> L1 (Scale a (IncDeg by q))
  Tail p' -> R1 (p' :*: q)

instance (IsPoly p, IsPoly q) => IsPoly (p :*: q) where
  type Uncons (p :*: q) = PROD (Uncons p) q
  constr = tyCase @p (\case {}) $ fromProdStep . constr . convCase
  destr = tyCase @p (\case (p :*: _) -> case p of {}) $ convCase . destr . toProdStep
  tyRep = case tyRep @p of
    TyN -> TyN
    TyJ @a @b @p' -> convRep $ tyRep @(ProdStep a b p' q)



-- Parallel or Dirichlet product as Day convolution of product

type DayStep a b p' q = Scale a (MultDeg b q) :+: (p' `Day` q)

type family DAY p q where
  DAY 'Nothing q = 'Nothing
  DAY ('Just '(a, b, p')) q = Uncons (DayStep a b p' q)

fromDayStep :: forall p q a b p' y. (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p') => DayStep a b p' q y -> (p `Day` q) y
fromDayStep (L1 (Scale pa (MultDeg qby))) = Day (Head pa id) qby (\b by -> by b)
fromDayStep (R1 (Day p q f)) = Day (Tail p) q f

toDayStep :: forall p q a b p' y. (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p') => (p `Day` q) y -> DayStep a b p' q y
toDayStep (Day p q f) = case p of
  Head pa by -> L1 (Scale pa (MultDeg $ fmap (\qy b -> f (by b) qy) q))
  Tail p' -> R1 (Day p' q f)

instance (IsPoly p, IsPoly q) => IsPoly (p `Day` q) where
  type Uncons (p `Day` q) = DAY (Uncons p) q
  constr = tyCase @p (\case {}) $ fromDayStep . constr . convCase
  destr = tyCase @p (\case {Day p _ _ -> case p of}) $ convCase . destr . toDayStep
  tyRep = case tyRep @p of
    TyN -> TyN
    TyJ @a @b @p' -> convRep $ tyRep @(DayStep a b p' q)




-- Exponentiation by a scalar

newtype (p ^. b) y = ExpS { getExpS :: b -> p y } deriving Functor

type ExpSStep p b = p :*: (p ^. b)

type family EXPS p b where
  EXPS p 'Nothing = 'Just '((), Void, V1)
  EXPS p ('Just b) = Uncons (ExpSStep p b)

fromExpSStep :: (IsPoly p, IsCountable b, Pred b ~ 'Just b') => ExpSStep p b' y -> (p ^. b) y
fromExpSStep (p :*: b'p) = ExpS (\case Z -> p; S b' -> getExpS b'p b')

toExpSStep :: (IsPoly p, IsCountable b, Pred b ~ 'Just b', IsCountable b') => (p ^. b) y -> ExpSStep p b' y
toExpSStep (ExpS bp) = bp Z :*: ExpS (bp . S)

instance (IsPoly p, IsCountable b) => IsPoly (p ^. b) where
  type Uncons (p ^. b) = EXPS p (Pred b)
  constr = tIf @b (\case PCHead{} -> ExpS absurd; PCTail q -> absurd1 q) $ fromExpSStep . constr . convCase
  destr = tIf @b (const $ PCHead () (\case {})) $ convCase . destr . toExpSStep
  tyRep = tIf @b TyJ $ tyRep'
    where
      tyRep' :: forall b'. (Pred b ~ 'Just b', IsCountable b') => TyRep (p ^. b)
      tyRep' = convRep $ tyRep @(ExpSStep p b')

mapExpS :: (p ~> p') -> p ^. b ~> p' ^. b
mapExpS f (ExpS bp) = ExpS (f . bp)



-- Composition

type CompStep a b p q = Scale a (q ^. b) :+: (p :.: q)

type family COMP p q where
  COMP 'Nothing q = 'Nothing
  COMP ('Just '(a, b, p)) q = Uncons (CompStep a b p q)

fromCompStep :: (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p') => CompStep a b p' q ~> (p :.: q)
fromCompStep (L1 (Scale a (ExpS bqy))) = Comp1 (Head a bqy)
fromCompStep (R1 (Comp1 p'qy)) = Comp1 (Tail p'qy)

toCompStep :: (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p') => (p :.: q) ~> CompStep a b p' q
toCompStep (Comp1 pqy) = case pqy of
  Head a bqy -> L1 (Scale a (ExpS bqy))
  Tail p'qy -> R1 (Comp1 p'qy)

instance (IsPoly p, IsPoly q, AllB IsCountable p) => IsPoly (p :.: q) where
  type Uncons (p :.: q) = COMP (Uncons p) q
  constr = tyCase @p (\case {}) $ fromCompStep . constr . convCase
  destr = tyCase @p (\case {}) $ convCase . destr . toCompStep
  tyRep = case tyRep @p of
    TyN -> TyN
    TyJ @a @b @p' -> convRep $ tyRep @(CompStep a b p' q)

bimapComp :: Functor p' => (p ~> p') -> (q ~> q') -> p :.: q ~> p' :.: q'
bimapComp fp fq = Comp1 . fmap fq . fp . unComp1

dayAsComp :: (Functor p, Functor q) => p `Day` q ~> p :.: q
dayAsComp (Day p q f) = Comp1 (fmap (\py -> fmap (f py) q) p)

duoidal :: (p :.: p') `Day` (q :.: q') ~> (p `Day` q) :.: (p' `Day` q')
duoidal (Day (Comp1 pp) (Comp1 qq) f) = Comp1 (Day pp qq (\p q -> Day p q f))


-- Exponentiation by a polynomial

type ExpStep a b p' q = ((q :.: Either b) ^. a) :*: (q ^ p')

type family EXP p q where
  EXP 'Nothing q = 'Just '((), Void, V1)
  EXP ('Just '(a, b, p')) q = Uncons (ExpStep a b p' q)

fromExpStep :: forall p q a b p'. (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p', IsCountable a, AllB IsCountable q) => ExpStep a b p' q ~> (q ^ p)
fromExpStep (ExpS aq :*: Exp f) = Exp \yx p -> case p of
  Head a bx -> either bx yx <$> unComp1 (aq a)
  Tail p' -> f yx p'

toExpStep :: forall p q a b p'. (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p', IsCountable a, AllB IsCountable q) => (q ^ p) ~> ExpStep a b p' q
toExpStep (Exp f) = ExpS (\a -> Comp1 (f Right (Head a Left))) :*: Exp \yx p' -> f yx (Tail p')

instance (IsPoly p, IsPoly q, AllA IsCountable p, AllB IsCountable q) => IsPoly (q ^ p) where
  type Uncons (q ^ p) = EXP (Uncons p) q
  constr = tyCase @p (\case PCHead{} -> Exp (const absurd1); PCTail q -> absurd1 q) $ fromExpStep . constr . convCase
  destr = tyCase @p (const $ PCHead () absurd) $ convCase . destr . toExpStep
  tyRep = case tyRep @p of
    TyN -> TyJ
    TyJ @a @b @p' -> convRep $ tyRep @(ExpStep a b p' q)



-- Dirichlet closure [p, q] = DClo p q

type DCloStep a b p' q = ((q :.: (,) b) ^. a) :*: DClo p' q

type family DCLO p q where
  DCLO 'Nothing q = 'Just '((), Void, V1)
  DCLO ('Just '(a, b, p')) q = Uncons (DCloStep a b p' q)

fromDCloStep :: forall p q a b p'. (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p', IsCountable a, AllB IsCountable q) => DCloStep a b p' q ~> DClo p q
fromDCloStep (ExpS aq :*: DClo k) = DClo \f p ->
  case p of
    Head a by -> fmap (\(b, py) -> f py (by b)) (unComp1 (aq a))
    Tail p' -> k f p'

toDCloStep :: forall p q a b p'. (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p', IsCountable a, AllB IsCountable q) => DClo p q ~> DCloStep a b p' q
toDCloStep (DClo k) =
  ExpS (\a -> Comp1 (k (\y bb -> (either id id bb, y)) (Head a Left)))
  :*: DClo \f p' -> k f (Tail p')

instance (IsPoly p, IsPoly q, AllA IsCountable p, AllB IsCountable q) => IsPoly (DClo p q) where
  type Uncons (DClo p q) = DCLO (Uncons p) q
  constr = tyCase @p (\case PCHead{} -> DClo (const absurd1); PCTail q -> absurd1 q) $ fromDCloStep . constr . convCase
  destr = tyCase @p (const $ PCHead () absurd) $ convCase . destr . toDCloStep
  tyRep = case tyRep @p of
    TyN -> TyJ
    TyJ @a @b @p' -> convRep $ tyRep @(DCloStep a b p' q)



-- Composition left coclosure, aka left kan extension
--  Lan q p   or   ⎡q⎤
--                 ⎣p⎦

type CoCloStep a b p' q = Mono a (q b) :+: Lan q p'

type family COCLO p q where
  COCLO 'Nothing q = 'Nothing
  COCLO ('Just '(a, b, p')) q = Uncons (CoCloStep a b p' q)

fromCoCloStep :: (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p') => CoCloStep a b p' q ~> Lan q p
fromCoCloStep (L1 (Mono a f)) = Lan (Head a id) f
fromCoCloStep (R1 (Lan p k)) = Lan (Tail p) k

toCoCloStep :: forall p q a b p'. (IsPoly p, IsPoly q, Uncons p ~ 'Just '(a, b, p'), IsPoly p') => Lan q p ~> CoCloStep a b p' q
toCoCloStep (Lan p k) = case p of
  Head a by -> L1 (Mono a (k . fmap by))
  Tail p' -> R1 (Lan p' k)

instance (IsPoly p, IsPoly q) => IsPoly (Lan q p) where
  type Uncons (Lan q p) = COCLO (Uncons p) q
  constr = tyCase @p (\case {}) $ fromCoCloStep . constr . convCase
  destr = tyCase @p (\(Lan p k) -> absurd1 p) $ convCase . destr . toCoCloStep
  tyRep = case tyRep @p of
    TyN -> TyN
    TyJ @a @b @p' -> convRep $ tyRep @(CoCloStep a b p' q)



-- Derivative of a polynomial functor

data Deriv p y where
  DerivJ :: (Uncons p ~ 'Just '(a, b, p'), IsPoly p', Pred b ~ 'Just b', IsCountable b') => Mono (a, b) b' y -> Deriv p y
  DerivN :: forall a b p' p y. (Uncons p ~ 'Just '(a, b, p'), IsPoly p') => Deriv p' y -> Deriv p y

deriving instance Functor p => Functor (Deriv p)

type family DERIV p where
  DERIV 'Nothing = 'Nothing
  DERIV ('Just '(a, b, p')) = 'Just '((a, b), DERIV' (Pred b), Deriv p')

type family DERIV' b' where
  DERIV' ('Just b') = b'
  DERIV' 'Nothing = Void

instance (IsPoly p, AllB IsCountable p) => IsPoly (Deriv p) where
  type Uncons (Deriv p) = DERIV (Uncons p)
  constr (PCHead a by) = case tyRep @p of
    TyJ @a @b @p' -> tIf @b (case a of (_, b') -> case b' of) $ DerivJ $ Mono a by
  constr (PCTail q) = case tyRep @p of
    TyJ @a @b @p' -> tIf @b DerivN DerivN q
  destr (DerivJ (Mono (a, b) b'y)) = PCHead (a, b) b'y
  destr (DerivN @a @b p) = tIf @b (PCTail p) (PCTail p)
  tyRep = tyCase @p TyN TyJ

b'2b :: (IsCountable b, Eq b, Pred b ~ 'Just b', IsCountable b') => b -> b' -> b
b'2b b b' = if S b' == b then Z else S b'

b'y2by :: (IsCountable b, Eq b, Pred b ~ 'Just b') => (b, y) -> (b' -> y) -> b -> y
b'y2by (b, y) b'y b' = if b' == b then y else pred (pred y b'y b) b'y b'

by2b'y :: (IsCountable b, Eq b, Pred b ~ 'Just b', IsCountable b') => b -> (b -> y) -> b' -> y
by2b'y b by b' = by (b'2b b b')

plug :: (IsPoly p, AllB IsCountable p, AllB Eq p) => Deriv p y -> y -> p y
plug (DerivJ (Mono (a, b) b'y)) y = Head a (b'y2by (b, y) b'y)
plug (DerivN p') y = Tail $ plug p' y

allDerivatives :: (IsPoly p, AllB IsCountable p, AllB Eq p) => p y -> p (Deriv p y, y)
allDerivatives (Head @_ @_ @_ @b a by) = tIf @b undefined $ Head a $ \b -> (DerivJ (Mono (a, b) (by2b'y b by)), by b)
allDerivatives (Tail py) = Tail $ fmap (\(d, y) -> (DerivN d, y)) (allDerivatives py)

around :: (IsPoly p, AllB IsCountable p, AllB Eq p) => Deriv p y -> y -> Deriv p (Deriv p y, y)
around (DerivJ @_ @_ @b (Mono (a, b) b'y)) y = DerivJ (Mono (a, b) (\b' -> let newb = b'2b b b' in (DerivJ (Mono (a, newb) (by2b'y newb $ b'y2by (b, y) b'y)), b'y b')))
around (DerivN d) y = DerivN ((\(d', y') -> (DerivN d', y')) <$> around d y)





instance IsPoly (Mealy i o) where
  type Uncons (Mealy i o) = 'Just '(i -> o, i, V1)
  constr (PCHead io iy) = Mealy $ \i -> (io i, iy i)
  constr (PCTail q) = case q of {}
  destr (Mealy ioy) = PCHead (fst . ioy) (snd . ioy)
  tyRep = TyJ



type ShowPoly p = SHOW (ToTMap (Uncons p))
type family SHOW (m :: TMap) :: Symbol where
  SHOW '[] = "0"
  SHOW '[ba] = ShowPair ba
  SHOW (ba ': m) = ShowPair ba `AppendSymbol` " + " `AppendSymbol` SHOW m

type family ShowPair ba :: Symbol where
  ShowPair '(b, a) = ShowNat a `AppendSymbol` "y^" `AppendSymbol` ShowNat b
type family ShowNat (n :: N.Nat) :: Symbol where
  ShowNat n = N.NatToChar (n N.+ 48) `ConsSymbol` ""

type family Simp t :: N.Nat where
  Simp () = 1
  Simp Void = 0
  Simp (a, b) = Simp a N.* Simp b
  Simp (Either a b) = Simp a N.+ Simp b

type TMap = [(N.Nat, N.Nat)]

type family InsertTMap (ba :: (N.Nat, N.Nat)) (m :: TMap) :: TMap where
  InsertTMap ba '[] = '[ba]
  InsertTMap '(b, a) ('(b', a') ': m) =
    OrdCond (Compare b b')
      ('(b, a) : '(b', a') : m)
      ('(b, a N.+ a') : m)
      ('(b', a') : InsertTMap '(b, a) m)

type family ToTMap p :: TMap where
  ToTMap 'Nothing = '[]
  ToTMap ('Just '(a, b, q)) = InsertTMap '(Simp b, Simp a) (ToTMap (Uncons q))