Implementation of the Poly category in Haskell

PolyRep.hs

{-# LANGUAGE GHC2021, PatternSynonyms, LambdaCase, DataKinds, PolyKinds, TypeFamilies, NoStarIsType, UndecidableInstances, UndecidableSuperClasses, AllowAmbiguousTypes, DefaultSignatures #-}

import Data.Functor (void)
import Data.Bifunctor (bimap)
import Data.Kind (Type, Constraint)
import Data.Void
-- from fin
import Data.Fin (Fin(..))
import qualified Data.Fin as Fin
import Data.Type.Nat


prev :: forall n. SNat ('S n) -> SNat n
prev SS = snat


type PolyKind = [(Type, Type)]

data Prf (p :: PolyKind) where
  PrfNil :: Prf '[]
  PrfCons :: Prf p -> Prf ('(a, b) ': p)

class IsP (p :: PolyKind) where
  prf :: Prf p
instance IsP '[] where
  prf = PrfNil
instance (ab ~ '(a, b), IsP p) => IsP (ab ': p) where
  prf = PrfCons prf

infixr 0 //
(//) :: Prf p -> (IsP p => r) -> r
PrfNil // r = r
PrfCons p // r = p // r

infixr 0 ///
(///) :: (q -> Prf p) -> (IsP p => q -> r) -> q -> r
f /// g = \q -> f q // g q


type family AllA (ca :: Type -> Constraint) (p :: PolyKind) :: Constraint where
  AllA ca '[] = ()
  AllA ca ('(a, b) ': p) = (ca a, AllA ca p)

type family AllB (cb :: Type -> Constraint) (p :: PolyKind) :: Constraint where
  AllB cb '[] = ()
  AllB cb ('(a, b) ': p) = (cb b, AllB cb p)

type All ca cb p = (AllA ca p, AllB cb p)


-- | Polynomial functors as a sum of monomials @a * y ^ b@
data PolyF (p :: PolyKind) y where
  Y :: IsP p => a -> (b -> y) -> PolyF ('(a, b) ': p) y
  N :: PolyF p y -> PolyF ('(a, b) ': p) y

instance Functor (PolyF p) where
  fmap f (Y a by) = Y a (f . by)
  fmap f (N p) = N (fmap f p)

-- | Natural transformations between polynomial functors
type p ~> q = forall y. PolyF p y -> PolyF q y


-- | Concrete representation of natural transformations between polynomial functors.
-- For each monomial @a * y ^ b@ we have for each @a@ a @q@ polynomial full of @b@s which specify which @y@ was used from the monomial.
data Poly (p :: PolyKind) (q :: PolyKind) where
  PNil :: IsP q => Poly '[] q
  P :: (a -> PolyF q b) -> Poly p q -> Poly ('(a, b) ': p) q

-- | @Y@ in the shape of the argument to @P@
py :: IsP p => (b -> y) -> a -> PolyF ('(a, b) ': p) y
py = flip Y

fromMonoF :: PolyF (Mono a b) y -> (a, b -> y)
fromMonoF (Y a by) = (a, by)
fromMonoF (N p) = case p of {}


-- Conversions between the two representations.
fromPoly :: Poly p q -> p ~> q
fromPoly PNil p = case p of {}
fromPoly (P qs _) (Y a ys) = fmap ys (qs a)
fromPoly (P _ pq) (N p) = fromPoly pq p

toPoly :: forall p q. (IsP p, IsP q) => (p ~> q) -> Poly p q
toPoly f = case prf @p of
  PrfNil -> PNil
  PrfCons p -> p // P (f . py id) (toPoly (f . N))


prfF :: PolyF p y -> Prf p
prfF (Y _ _) = prf
prfF (N p) = PrfCons (prfF p)

prfP :: Poly p q -> Prf p
prfP PNil = PrfNil
prfP (P _ pq) = PrfCons (prfP pq)

prfQ :: Poly p q -> Prf q
prfQ PNil = prf
prfQ (P _ pq) = prfQ pq

prfTail :: Prf ('(a, b) : p) -> Prf p
prfTail (PrfCons p) = p



mapP :: IsP r => (q ~> r) -> Poly p q -> Poly p r
mapP _ PNil = PNil
mapP f (P qs pq) = P (f . qs) (mapP f pq)

liftQ :: Poly p q -> Poly p ('(a, b) ': q)
liftQ PNil = PNil
liftQ (P qs pq) = P (N . qs) (liftQ pq)


-- Poly is a category
idP :: IsP p => Poly p p
idP = toPoly id

compP :: Poly q r -> Poly p q -> Poly p r
compP q = prfQ q // mapP (fromPoly q)


type Zero = '[]
type Mono a b = '[ '(a, b) ]
type Duo a b c d = '[ '(a, b), '(c, d) ]

type I = Mono () Void
type Y = Mono () ()
mkY :: y -> PolyF Y y
mkY y = Y () (const y)

-- type List = ListN 'Z
-- type ListN n = '((), Fin n) ': ListN ('S n)


-- Some adjunctions

-- p(1) |- Ay (x = ()) and p(0) |- A (x = Void)
adj1 :: IsP p => (a -> PolyF p x) -> Poly (Mono a x) p
adj1 f = P f PNil

-- A |- p(1)
adj2 :: IsP p => (PolyF p () -> a) -> Poly p (Mono a Void)
adj2 f = toPoly (\p1 -> Y (f $ void p1) absurd)

data Gamma (p :: PolyKind) where
  GNil :: Gamma '[]
  G :: b -> Gamma p -> Gamma ('(a, b) ': p)

-- y^A |- Gamma(p)
adj3 :: forall p a. IsP p => Poly (Mono () a) p -> (Gamma p -> a)
adj3 (P f PNil) g = loop (f ()) g where
  loop :: PolyF q a -> (Gamma q -> a)
  loop (Y _ ba) (G b _) = ba b
  loop (N p) (G _ g') = loop p g'


-- Coproducts

type family (++) (p :: PolyKind) (q :: PolyKind) :: PolyKind where
  '[] ++ q = q
  ('(a, b) ': p) ++ q = '(a, b) ': (p ++ q)

coprodPrf :: Prf p -> Prf q -> Prf (p ++ q)
coprodPrf PrfNil q = q
coprodPrf (PrfCons p) q = PrfCons (coprodPrf p q)

inj1P :: Prf p -> Prf q -> Poly p (p ++ q)
inj1P PrfNil q = q // PNil
inj1P (PrfCons p) q = (prfQ /// P (py id) . liftQ) (inj1P p q)

inj2P :: Prf p -> Prf q -> Poly q (p ++ q)
inj2P PrfNil q = q // idP
inj2P (PrfCons p) q = p // liftQ (inj2P p q)

(|||) :: Poly p r -> Poly q r -> Poly (p ++ q) r
PNil   ||| qr = qr
P r pr ||| qr = P r (pr ||| qr)

coprodSwap :: Prf p -> Prf q -> Poly (p ++ q) (q ++ p)
coprodSwap p q = inj2P q p ||| inj1P q p

coprodBimap :: Poly p0 q0 -> Poly p1 q1 -> Poly (p0 ++ p1) (q0 ++ q1)
coprodBimap pq0 pq1 = (inj1P q0 q1 `compP` pq0) ||| (inj2P q0 q1 `compP` pq1)
  where
    q0 = prfQ pq0
    q1 = prfQ pq1

inj1 :: forall p q. (IsP p, IsP q) => Poly p (p ++ q)
inj1 = inj1P @p @q prf prf

inj2 :: forall p q. (IsP p, IsP q) => Poly q (p ++ q)
inj2 = inj2P @p @q prf prf

inj1F :: Prf p -> Prf q -> p ~> (p ++ q)
inj1F p q = fromPoly (inj1P p q)

inj2F :: Prf p -> Prf q -> q ~> (p ++ q)
inj2F p q = fromPoly (inj2P p q)

coprodF :: Prf p -> (PolyF p y -> r) -> (PolyF q y -> r) -> PolyF (p ++ q) y -> r
coprodF PrfNil _ qr r = qr r
coprodF (PrfCons p) pr _ (Y a by) = p // pr (Y a by)
coprodF (PrfCons p) pr qr (N pq) = coprodF p (pr . N) qr pq


-- Apply Hask functors to the parameters of a polynomial

type family ApplyA (f :: Type -> Type) (q :: PolyKind) :: PolyKind where
  ApplyA f '[] = '[]
  ApplyA f ('(qa, qb) ': q) = '(f qa, qb) ': ApplyA f q

applyAPrf :: forall f q. Prf q -> Prf (ApplyA f q)
applyAPrf PrfNil = PrfNil
applyAPrf (PrfCons q) = PrfCons (applyAPrf @f q)

applyAP :: forall f p y. (forall a. a -> f a) -> PolyF p y -> PolyF (ApplyA f p) y
applyAP f (Y pa pby) = prfTail (applyAPrf @f (prf @p)) // Y (f pa) pby
applyAP f (N p) = N (applyAP f p)

type family ApplyB (f :: Type -> Type) (q :: PolyKind) :: PolyKind where
  ApplyB f '[] = '[]
  ApplyB f ('(qa, qb) ': q) = '(qa, f qb) ': ApplyB f q

applyBPrf :: forall f q. Prf q -> Prf (ApplyB f q)
applyBPrf PrfNil = PrfNil
applyBPrf (PrfCons q) = PrfCons (applyBPrf @f q)

applyBP :: forall f p y0 y1. (forall b. (b -> y0) -> f b -> y1) -> PolyF p y0 -> PolyF (ApplyB f p) y1
applyBP f (Y pa pby) = prfTail (applyBPrf @f (prf @p)) // Y pa (f pby)
applyBP f (N p) = N (applyBP f p)

type Apply f g q = ApplyA f (ApplyB g q)


-- Multiply with a monomial

-- a * y^b * q
type family (.*) (ab :: (Type, Type)) (q :: PolyKind) :: PolyKind where
  '(a, b) .* q = Apply ((,) a) (Either b) q

scaleAPrf :: forall a b p q. Prf ('(a, b) ': p) -> Prf q -> Prf (ApplyA ((,) a) q)
scaleAPrf _ = applyAPrf @((,) a)

scalePrf :: forall a b p q. Prf ('(a, b) ': p) -> Prf q -> Prf ('(a, b) .* q)
scalePrf p q = scaleAPrf p (applyBPrf @(Either b) q)

scaleP :: forall a b p y. a -> (b -> y) -> PolyF p y -> PolyF ('(a, b) .* p) y
scaleP a by = scaleAP a . applyBP (either by)

scaleAP :: forall a p y. a -> PolyF p y -> PolyF (ApplyA ((,) a) p) y
scaleAP a = applyAP ((,) a)

getScaleAP :: Prf p -> PolyF (ApplyA ((,) a) p) y -> a
getScaleAP PrfNil p = case p of {}
getScaleAP (PrfCons _) (Y (a, _) _) = a
getScaleAP (PrfCons p) (N ap) = getScaleAP p ap

getUnscaledAP :: Prf ('(a, b) ': p) -> Prf q -> PolyF (ApplyA ((,) a) q) y -> PolyF q y
getUnscaledAP _ PrfNil q = case q of {}
getUnscaledAP _ (PrfCons q) (Y (_, a) by) = q // Y a by
getUnscaledAP p (PrfCons q) (N aq) = N (getUnscaledAP p q aq)


-- Products

type family (*) (p :: PolyKind) (q :: PolyKind) :: PolyKind where
  '[] * q = '[]
  (ab ': p) * q = (ab .* q) ++ (p * q)

prodPrf :: Prf p -> Prf q -> Prf (p * q)
prodPrf PrfNil _ = PrfNil
prodPrf abp@(PrfCons p) q = coprodPrf (scalePrf abp q) (prodPrf p q)

proj1P :: forall p q. Prf p -> Prf q -> Poly (p * q) p
proj1P PrfNil _ = PNil
proj1P (PrfCons p) q = (prfP /// (||| liftQ (proj1P p q))) (p // step q)
  where
    step :: forall p' q' a b. IsP p' => Prf q' -> Poly ('(a, b) .* q') ('(a, b) ': p')
    step PrfNil = PNil
    step (PrfCons q') = P (py Left . fst) (step q')

proj2P :: forall p q. Prf p -> Prf q -> Poly (p * q) q
proj2P PrfNil q = q // PNil
proj2P abp@(PrfCons p) q = step abp q ||| proj2P p q
  where
    step :: forall p' q' a b. Prf ('(a, b) ': p') -> Prf q' -> Poly ('(a, b) .* q') q'
    step _ PrfNil = PNil
    step p' (PrfCons q') = q' // P (py Right . snd) (liftQ $ step p' q')

(&&&) :: Poly r p -> Poly r q -> Poly r (p * q)
rp@PNil &&& rq@PNil = prodPrf (prfQ rp) (prfQ rq) // PNil
P ab rp &&& P ab' rq = P (\a -> prodF (ab a) (ab' a)) (rp &&& rq)

prodSwap :: Prf p -> Prf q -> Poly (p * q) (q * p)
prodSwap p q = proj2P p q &&& proj1P p q

prodBimap :: Poly p0 q0 -> Poly p1 q1 -> Poly (p0 * p1) (q0 * q1)
prodBimap pq0 pq1 = (pq0 `compP` proj1P p0 p1) &&& (pq1 `compP` proj2P p0 p1)
  where
    p0 = prfP pq0
    p1 = prfP pq1

proj1 :: forall p q. (IsP p, IsP q) => Poly (p * q) p
proj1 = proj1P @p @q prf prf

proj2 :: forall p q. (IsP p, IsP q) => Poly (p * q) q
proj2 = proj2P @p @q prf prf

proj1F :: Prf p -> Prf q -> (p * q) ~> p
proj1F p q = fromPoly (proj1P p q)

proj2F :: Prf p -> Prf q -> (p * q) ~> q
proj2F p q = fromPoly (proj2P p q)

prodF :: PolyF p y -> PolyF q y -> PolyF (p * q) y
prodF p@(Y pa pby) q = let pq = scaleP pa pby q in inj1F (prfF pq) (prodPrf (prfTail (prfF p)) (prfF q)) pq
prodF abp@(N p) q = let pq = prodF p q in inj2F (scalePrf (prfF abp) (prfF q)) (prfF pq) pq


-- Parallel or Dirichlet product

type ParMono a b q = Apply ((,) a) ((,) b) q

type family X (p :: PolyKind) (q :: PolyKind) :: PolyKind where
  '[] `X` q = '[]
  ('(a, b) ': p) `X` q = ParMono a b q ++ (p `X` q)

parMonoPrf :: forall a b p q. Prf ('(a, b) ': p) -> Prf q -> Prf (ParMono a b q)
parMonoPrf p q = scaleAPrf p (applyBPrf @((,) b) q)

parProdPrf :: Prf p -> Prf q -> Prf (p `X` q)
parProdPrf PrfNil _ = PrfNil
parProdPrf abp@(PrfCons p) q = coprodPrf (parMonoPrf abp q) (parProdPrf p q)

parProdF :: PolyF p py -> PolyF q qy -> PolyF (p `X` q) (py, qy)
parProdF p@(Y pa pby) q = let pq = scaleAP pa (applyBP (bimap pby) q) in inj1F (prfF pq) (parProdPrf (prfTail (prfF p)) (prfF q)) pq
parProdF abp@(N p) q = let pq = parProdF p q in inj2F (parMonoPrf (prfF abp) (prfF q)) (prfF pq) pq

parProdBimap :: Poly p0 q0 -> Poly p1 q1 -> Poly (p0 `X` p1) (q0 `X` q1)
parProdBimap pq0 pq1 = parProdPrf (prfQ pq0) (prfQ pq1) // parProdBimap' pq0 pq1
  where
    parProdBimap' :: IsP (q0 `X` q1) => Poly p0 q0 -> Poly p1 q1 -> Poly (p0 `X` p1) (q0 `X` q1)
    parProdBimap' PNil _ = PNil
    parProdBimap' (P ab pq0') pq1' = step ab pq1' ||| parProdBimap' pq0' pq1'
    step :: IsP (q0 `X` q1) => (a -> PolyF q0 b) -> Poly p1 q1 -> Poly (ParMono a b p1) (q0 `X` q1)
    step _ PNil = PNil
    step ab0 (P ab1 pq1') = P (\(a0, a1) -> parProdF (ab0 a0) (ab1 a1)) (step ab0 pq1')

dualParMono :: Poly (Mono a b `X` Mono b a) Y
dualParMono = P (\(a, b) -> mkY (b, a)) PNil

dualParDuo :: Poly (Duo a b c d `X` Mono (b, d) (Either a c)) Y
dualParDuo
  = P (\(a, (b, _)) -> mkY (b, Left a))
  $ P (\(c, (_, d)) -> mkY (d, Right c))
  $ PNil

dualParDuo' :: Poly (Mono (b, d) (Either a c) `X` Duo a b c d) Y
dualParDuo'
  = P (\((b, _), a) -> mkY (Left a, b))
  $ P (\((_, d), c) -> mkY (Right c, d))
  $ PNil


-- Composition

class CanExp b where
  type (p :: PolyKind) ^. b :: PolyKind
  expPrf :: Prf p -> Prf (p ^. b)
  expF :: (b -> PolyF p y) -> PolyF (p ^. b) y
  indexF :: Prf p -> PolyF (p ^. b) y -> b -> PolyF p y

instance CanExp (Fin 'Z) where
  type p ^. Fin 'Z = I
  expPrf _ = prf
  expF _ = Y () absurd
  indexF _ _ = Fin.absurd

instance CanExp (Fin ('S 'Z)) where
  type p ^. Fin ('S 'Z) = p
  expPrf p = p
  expF f = f FZ
  indexF _ p _ = p

instance CanExp (Fin ('S n)) => CanExp (Fin ('S ('S n))) where
  type p ^. Fin ('S ('S n)) = p * (p ^. Fin ('S n))
  expPrf p = prodPrf p (expPrf @(Fin ('S n)) p)
  expF f = prodF (f FZ) (expF (f . FS))
  indexF p pb FZ = proj1F p (expPrf @(Fin ('S n)) p) pb
  indexF p pb (FS n) = indexF p (proj2F p (expPrf @(Fin ('S n)) p) pb) n

instance CanExp Void where
  type p ^. Void = I
  expPrf _ = prf
  expF _ = Y () absurd
  indexF _ _ = absurd

instance CanExp () where
  type p ^. () = p
  expPrf p = p
  expF f = f ()
  indexF _ p _ = p

instance (CanExp a, CanExp b) => CanExp (Either a b) where
  type p ^. Either a b = (p ^. a) * (p ^. b)
  expPrf p = prodPrf (expPrf @a p) (expPrf @b p)
  expF f = prodF (expF (f . Left)) (expF (f . Right))
  indexF p pe (Left a) = indexF p (proj1F (expPrf @a p) (expPrf @b p) pe) a
  indexF p pe (Right b) = indexF p (proj2F (expPrf @a p) (expPrf @b p) pe) b

instance (CanExp a, CanExp b) => CanExp (a, b) where
  type p ^. (a, b) = (p ^. a) ^. b
  expPrf p = expPrf @b (expPrf @a p)
  expF f = expF (\b -> expF (\a -> f (a, b)))
  indexF p pp (a, b) = indexF p (indexF (expPrf @a p) pp b) a

type ExpMono a b q = ApplyA ((,) a) (q ^. b)

type family (<|) (p :: PolyKind) (q :: PolyKind) :: PolyKind where
  '[] <| q = '[]
  ('(a, b) ': p) <| q = ExpMono a b q ++ (p <| q)

type AllCanExp p = AllB CanExp p


expAPrf :: forall a b p q. CanExp a => Prf ('(a, b) ': p) -> Prf q -> Prf (q ^. a)
expAPrf _ = expPrf @a

expBPrf :: forall a b p q. CanExp b => Prf ('(a, b) ': p) -> Prf q -> Prf (q ^. b)
expBPrf _ = expPrf @b

expMonoPrf :: CanExp b => Prf ('(a, b) ': p) -> Prf q -> Prf (ExpMono a b q)
expMonoPrf abp q = scaleAPrf abp (expBPrf abp q)

compPrf :: AllCanExp p => Prf p -> Prf q -> Prf (p <| q)
compPrf PrfNil _ = PrfNil
compPrf abp@(PrfCons p) q = coprodPrf (expMonoPrf abp q) (compPrf p q)

toCompF :: AllCanExp p => Prf p -> Prf q -> PolyF p (PolyF q y) -> PolyF (p <| q) y
toCompF PrfNil _ pq = case pq of {}
toCompF abp@(PrfCons p) q (Y a bqy) = inj1F (scaleAPrf abp (expBPrf abp q)) (compPrf p q) (scaleAP a (expF bqy))
toCompF abp@(PrfCons p) q (N pq) = inj2F (scaleAPrf abp (expBPrf abp q)) (compPrf p q) (toCompF p q pq)

fromCompF :: AllCanExp p => Prf p -> Prf q -> PolyF (p <| q) y -> PolyF p (PolyF q y)
fromCompF PrfNil _ pq = case pq of {}
fromCompF abp@(PrfCons p) q pq = coprodF (scaleAPrf abp (expBPrf abp q)) decompose (N . fromCompF p q) pq
  where
    decompose aqb = p // Y (getScaleAP (expBPrf abp q) aqb) (indexF q (getUnscaledAP abp (expBPrf abp q) aqb))


-- Cartesian closure

type family (^) (q :: PolyKind) (p :: PolyKind) :: PolyKind where
  q ^ '[] = I
  q ^ ('(a, b) ': p) = ((q <| Duo b Void () ()) ^. a) * (q ^ p)

fromDuoF :: PolyF (Duo a Void () ()) y -> Either y a
fromDuoF (Y a _) = Right a
fromDuoF (N (Y () by)) = Left (by ())
fromDuoF (N (N p)) = case p of {}

toDuoF :: Either y a -> PolyF (Duo a Void () ()) y
toDuoF (Right a) = Y a (\case {})
toDuoF (Left y) = N (Y () (\() -> y))

duoBPrf :: Prf ('(a, b) ': p) -> Prf (Duo b Void () ())
duoBPrf _ = PrfCons (PrfCons PrfNil)

cartClosurePrf :: (AllA CanExp p, AllB CanExp q) => Prf p -> Prf q -> Prf (q ^ p)
cartClosurePrf PrfNil _ = PrfCons PrfNil
cartClosurePrf abp@(PrfCons p) q = prodPrf (expAPrf abp (compPrf q (duoBPrf abp))) (cartClosurePrf p q)

fromCartClosure :: forall p q p'. (AllA CanExp p, AllB CanExp q) => Prf p -> Prf q -> Poly p' (q ^ p) -> Poly (p' * p) q
fromCartClosure PrfNil q PNil = q // PNil
fromCartClosure (PrfCons p) q PNil = cartClosurePrf p q // fromCartClosure p q PNil
fromCartClosure p q (P ab p'qp) = (q // step p ab) ||| fromCartClosure p q p'qp
  where
    step :: forall a b p1. (AllA CanExp p1, IsP q) => Prf p1 -> (a -> PolyF (q ^ p1) b) -> Poly ('(a, b) .* p1) q
    step PrfNil _ = PNil
    step abp@(PrfCons p1) ab1 =
      let duoB = duoBPrf abp
          qOfB = compPrf q duoB
          qOfBPowA = expAPrf abp qOfB
          cCloPQ = cartClosurePrf p1 q
      in P
        (\(a1, a2) -> fmap fromDuoF $ fromCompF q duoB $ indexF qOfB (proj1F qOfBPowA cCloPQ (ab1 a1)) a2)
        (step p1 (proj2F qOfBPowA cCloPQ . ab1))

toCartClosure :: forall p q p'. (AllA CanExp p, AllB CanExp q) => Prf p' -> Prf p -> Prf q -> Poly (p' * p) q -> Poly p' (q ^ p)
toCartClosure PrfNil p q PNil = cartClosurePrf p q // PNil
toCartClosure abp'@(PrfCons p') p q p'qp = P (step p $ compP p'qp $ inj1P (scalePrf abp' p) (prodPrf p' p)) (toCartClosure p' p q (compP p'qp (inj2P (scalePrf abp' p) (prodPrf p' p))))
  where
    step :: AllA CanExp p1 => Prf p1 -> Poly ('(a, b) .* p1) q -> a -> PolyF (q ^ p1) b
    step PrfNil _ _ = Y () absurd
    step abp@(PrfCons p1) (P ab pq) a1 = prodF
      (expF (\a2 -> toCompF q (duoBPrf abp) $ fmap toDuoF $ ab (a1, a2)))
      (step p1 pq a1)


-- Dirichlet closure [–,–]

type family DClo (q :: PolyKind) (p :: PolyKind) :: PolyKind where
  DClo q '[] = I
  DClo q ('(a, b) ': p) = ((q <| Mono b ()) ^. a) * DClo q p

monoBPrf :: Prf ('(a, b) ': p) -> Prf (Mono b ())
monoBPrf _ = PrfCons PrfNil

diriClosurePrf :: (AllA CanExp p, AllB CanExp q) => Prf p -> Prf q -> Prf (DClo q p)
diriClosurePrf PrfNil _ = PrfCons PrfNil
diriClosurePrf abp@(PrfCons p) q = prodPrf (expAPrf abp (compPrf q (monoBPrf abp))) (diriClosurePrf p q)

fromDiriClosure :: forall p q p'. (AllA CanExp p, AllB CanExp q) => Prf p -> Prf q -> Poly p' (DClo q p) -> Poly (p' `X` p) q
fromDiriClosure PrfNil q PNil = q // PNil
fromDiriClosure (PrfCons p) q PNil = diriClosurePrf p q // fromDiriClosure p q PNil
fromDiriClosure p q (P ab p'qp) = (q // step p ab) ||| fromDiriClosure p q p'qp
  where
    step :: forall a b p1. (AllA CanExp p1, IsP q) => Prf p1 -> (a -> PolyF (DClo q p1) b) -> Poly (ParMono a b p1) q
    step PrfNil _ = PNil
    step abp@(PrfCons p1) ab1 =
      let monoB = monoBPrf abp
          qOfB = compPrf q monoB
          qOfBPowA = expAPrf abp qOfB
          dCloPQ = diriClosurePrf p1 q
      in P
        (\(a1, a2) -> fmap ((\(b2, fb1) -> (fb1 (), b2)) . fromMonoF) $ fromCompF q monoB $ indexF qOfB (proj1F qOfBPowA dCloPQ (ab1 a1)) a2)
        (step p1 (proj2F qOfBPowA dCloPQ . ab1))

toDiriClosure :: forall p q p'. (AllA CanExp p, AllB CanExp q) => Prf p' -> Prf p -> Prf q -> Poly (p' `X` p) q -> Poly p' (DClo q p)
toDiriClosure PrfNil p q PNil = diriClosurePrf p q // PNil
toDiriClosure abp'@(PrfCons p') p q p'qp = P (step p $ compP p'qp $ inj1P (parMonoPrf abp' p) (parProdPrf p' p)) (toDiriClosure p' p q (compP p'qp (inj2P (parMonoPrf abp' p) (parProdPrf p' p))))
  where
    step :: AllA CanExp p1 => Prf p1 -> Poly (ParMono a b p1) q -> a -> PolyF (DClo q p1) b
    step PrfNil _ _ = Y () absurd
    step abp@(PrfCons p1) (P ab pq) a1 = prodF
      (expF (\a2 -> toCompF q (monoBPrf abp) $ fmap (\(b1, b2) -> Y b2 (const b1)) $ ab (a1, a2)))
      (step p1 pq a1)


-- Categories as comonoids in Poly

class All Eq CanExp p => Comonoid p where
  unit :: Poly p Y
  mult :: Poly p (p <| p)

type Obj p = PolyF p ()

instance Show (Obj '[]) where
  show p = case p of {}
instance (Show (Obj p), Show a) => Show (Obj ('(a, b) ': p)) where
  show (Y a _) = "Y " ++ show a
  show (N p) = "N (" ++ show p ++ ")"

-- Each monomial ay^b represents a collection of @a@ objects with @b@ outgoing arrows.
-- So an arrow is a choice of monomial with a choice of source object @a@ and a choice of outgoing arrow @b@.
data Arr p where
  ArrY :: IsP p => a -> b -> Arr ('(a, b) ': p)
  ArrN :: Arr p -> Arr ('(a, b) ': p)

deriving instance Show (Arr '[])
deriving instance (Show (Arr p), Show a, Show b) => Show (Arr ('(a, b) ': p))

arrPrf :: Arr p -> Prf p
arrPrf ArrY{} = prf
arrPrf (ArrN p) = PrfCons (arrPrf p)

polyCatDom :: Arr p -> Obj p
polyCatDom (ArrY a _) = Y a (const ())
polyCatDom (ArrN p) = N (polyCatDom p)

polyCatCod :: Comonoid p => Arr p -> Obj p
polyCatCod arr = step mult arr arr where
  step :: All Eq CanExp q => Poly p (q <| q) -> Arr p -> Arr q -> Obj q
  step PNil a _ = case a of {}
  step (P _ ppp) (ArrN p) q = step ppp p q
  step (P f _) (ArrY a _) q = loop (fromCompF (arrPrf q) (arrPrf q) (f a)) q
    where
      loop :: PolyF p (PolyF q x) -> Arr p -> Obj q
      loop (Y _ by) (ArrY _ b) = () <$ by b
      loop (N pp) (ArrN p) = loop pp p
      loop _ _ = error "Comonoid laws violated"

polyCatId :: Comonoid p => Obj p -> Arr p
polyCatId = step unit
  where
    step :: Poly p Y -> Obj p -> Arr p
    step PNil o = case o of {}
    step (P ab _) (Y a _) = ArrY a (case ab a of Y () b -> b (); N p -> case p of {})
    step (P _ p) (N o) = ArrN (step p o)

polyCatComp :: Comonoid p => Arr p -> Arr p -> Maybe (Arr p)
polyCatComp arr = step mult arr arr where
  step :: All Eq CanExp q => Poly p (q <| q) -> Arr p -> Arr q -> Arr q -> Maybe (Arr p)
  step PNil a _ = case a of {}
  step (P _ ppp) (ArrN p) q = fmap ArrN <$> step ppp p q
  step (P f _) (ArrY a _) q = fmap (ArrY a) <$> loop (fromCompF (arrPrf q) (arrPrf q) (f a)) q
    where
      loop :: All Eq CanExp q => PolyF p (PolyF q b) -> Arr p -> Arr q -> Maybe b
      loop (Y _ by) (ArrY _ b) = loop2 (by b)
      loop (N pp) (ArrN p) = loop pp p
      loop _ _ = error "Comonoid laws violated"
      loop2 :: All Eq CanExp q => PolyF q b -> Arr q -> Maybe b
      loop2 (Y a' b'y) (ArrY a'' b') = if a' == a'' then Just (b'y b') else Nothing
      loop2 (N pp) (ArrN p) = loop2 pp p
      loop2 _ _ = Nothing


-- A ==> C
-- |     ^
-- +> B -+
data A = A deriving (Eq, Show)
data B = B deriving (Eq, Show)
data C = C deriving (Eq, Show)
type Example = '[ '(A, Fin Nat4), '(B, Fin Nat2), '(C, Fin Nat1) ]
objA, objB, objC :: Obj Example
objA = Y A (const ())
objB = N $ Y B (const ())
objC = N $ N $ Y C (const ())
idA, idB, idC, arrAB, arrBC, arrAC1, arrAC2 :: Arr Example
idA    = ArrY A Fin0
arrAB  = ArrY A Fin1
arrAC1 = ArrY A Fin2
arrAC2 = ArrY A Fin3
idB    = ArrN $ ArrY B Fin0
arrBC  = ArrN $ ArrY B Fin1
idC    = ArrN $ ArrN $ ArrY C Fin0
expand :: PolyF Example (PolyF Example y) -> PolyF (Example <| Example) y
expand = toCompF (prf @Example) (prf @Example)
instance Comonoid Example where
  unit = P (\A -> mkY Fin0) $ P (\B -> mkY Fin0) $ P (\C -> mkY Fin0) $ PNil
  mult = P (\A -> expand $ Y A $ \case
                    Fin0 -> Y A id
                    Fin1 -> N $ Y B $ \case { Fin0 -> Fin1; Fin1 -> Fin2; _ -> err }
                    Fin2 -> N $ N $ Y C $ \case { Fin0 -> Fin2; _ -> err }
                    Fin3 -> N $ N $ Y C $ \case { Fin0 -> Fin3; _ -> err }
                    _ -> err)
         $ P (\B -> expand $ N $ Y B $ \case
                    Fin0 -> N $ Y B id
                    Fin1 -> N $ N $ Y C $ \case { Fin0 -> Fin1; _ -> err }
                    _ -> err)
         $ P (\C -> expand $ N $ N $ Y C $ \case
                    Fin0 -> N $ N $ Y C id
                    _ -> err)
         $ PNil

pattern Fin0 :: Fin ('S n)
pattern Fin0 = FZ
pattern Fin1 :: Fin ('S ('S n))
pattern Fin1 = FS Fin0
pattern Fin2 :: Fin ('S ('S ('S n)))
pattern Fin2 = FS Fin1
pattern Fin3 :: Fin ('S ('S ('S ('S n))))
pattern Fin3 = FS Fin2

err :: a
err = error "Impossible"

-- Lenses

type Lens s t a b = s -> (a, b -> t)

fromLens :: Lens s t a b -> Poly (Mono s t) (Mono a b)
fromLens f = P (\s -> case f s of (a, bt) -> Y a bt) PNil

toLens :: Poly (Mono s t) (Mono a b) -> Lens s t a b
toLens (P qs PNil) s = case qs s of
  Y a bt -> (a, bt)
  N n -> case n of {}

toLens2 :: Poly (Mono s t) (Duo a0 b0 a1 b1) -> s -> Either (a0, b0 -> t) (a1, b1 -> t)
toLens2 (P qs PNil) s = case qs s of
  Y a0 b0t -> Left (a0, b0t)
  N (Y a1 b1t) -> Right (a1, b1t)
  N (N n) -> case n of {}

eitherLens2 :: Poly (Mono (Either a0 a1) (Either b0 b1)) (Duo a0 b0 a1 b1)
eitherLens2 = P (either (\a0 -> Y a0 Left) (\a1 -> N (Y a1 Right))) PNil


-- Systems

type System s p = Poly (Mono s s) p

fromSystem :: System s p -> s -> PolyF p s
fromSystem (P sps PNil) = sps

data Stream a = Str a (Stream a)

fromSystem1 :: System s (Mono a b) -> s -> Stream b -> Stream a
fromSystem1 sys st = case fromSystem sys st of
  Y a bst -> \(Str b bs) -> Str a (fromSystem1 sys (bst b) bs)
  N n -> case n of {}

fromSystem2 :: System s (Duo a0 b0 a1 b1) -> s -> Stream (b0, b1) -> Stream (Either a0 a1)
fromSystem2 sys st = case fromSystem sys st of
  Y a0 bst -> \(Str (b0, _) bs) -> Str (Left a0) (fromSystem2 sys (bst b0) bs)
  N (Y a1 bst) -> \(Str (_, b1) bs) -> Str (Right a1) (fromSystem2 sys (bst b1) bs)
  N (N n) -> case n of {}

flipflop :: System Bool (Mono Bool ())
flipflop = P (\s -> Y s (\() -> not s)) PNil

counter :: System Int (Mono Bool ())
counter = P (\n -> Y (odd n) (\() -> n + 1)) PNil