Day.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | Day convolution of profunctors: @'Day' p q@ convolves @p@ and @q@ along the tensors of the source and
-- target categories, with unit 'DayUnit' and internal hom 'DayExp'. Monoidal profunctors are closed under
-- it, and it preserves 'Procomonad's.
module Proarrow.Profunctor.Instance.Day where

import Proarrow.Category.Instance.Nat (Nat (..))
import Proarrow.Category.Instance.Prof (Prof (..))
import Proarrow.Category.Monoidal
  ( Monoidal (..)
  , MonoidalProfunctor (..)
  , SymMonoidal (..)
  , swap
  , swapInner'
  , unitObj
  )
import Proarrow.Category.Monoidal.Closed (Closed (..))
import Proarrow.Category.Monoidal.CopyDiscard (CopyDiscard (..))
import Proarrow.Category.Monoidal.Distributive (Distributive (..))
import Proarrow.Category.Monoidal.Strength (MonStrong, first', second')
import Proarrow.Core
  ( CAT
  , CategoryOf (..)
  , Profunctor (..)
  , Promonad (..)
  , lmap
  , obj
  , rmap
  , src
  , tgt
  , (//)
  , type (+->)
  )
import Proarrow.Functor (Functor (..))
import Proarrow.Monoid (CocommutativeComonoid, CommutativeMonoid, Comonoid (..), Monoid (..), Supplies)
import Proarrow.Object (pattern Objs)
import Proarrow.Profunctor.Instance.Composition ((:.:) (..))
import Proarrow.Profunctor.Instance.Coproduct ((:+:) (..))
import Proarrow.Promonad (Procomonad (..))

-- | The unit of 'Day' convolution: a pair of arrows through the monoidal units.
data DayUnit a b where
  DayUnit :: a ~> Unit -> Unit ~> b -> DayUnit a b

instance (CategoryOf j, CategoryOf k) => Profunctor (DayUnit :: j +-> k) where
  dimap l r (DayUnit f g) = DayUnit (f . l) (r . g)
  r \\ DayUnit f g = r \\ f \\ g

type Day :: (j +-> k) -> (j +-> k) -> j +-> k

-- | The Day convolution on profunctors.
data Day p q a b where
  Day :: forall c d e f p q a b. a ~> c ** e -> p c d -> q e f -> d ** f ~> b -> Day p q a b

day :: (Monoidal j, Monoidal k, Profunctor (p :: j +-> k), Profunctor q) => p c d -> q e f -> Day p q (c ** e) (d ** f)
day p q = Day (src p ** src q) p q (tgt p ** tgt q)

instance (Profunctor p, Profunctor q) => Profunctor (Day p q) where
  dimap l r (Day f p q g) = Day (lmap l f) p q (rmap r g)
  r \\ Day f _ _ g = r \\ f \\ g

instance (Procomonad p, Procomonad q, Monoidal k) => Procomonad (Day p q :: k +-> k) where
  proextract (Day f p q g) = g . (proextract p ** proextract q) . f
  produplicate (Day f p q g) = case (produplicate p, produplicate q) of
    (p' :.: p'', q' :.: q'') -> let bb = tgt p' ** tgt q' in Day f p' q' bb :.: Day bb p'' q'' g

instance (Profunctor p) => Functor (Day p) where
  map (Prof n) = Prof \(Day f p q g) -> Day f p (n q) g

instance Functor Day where
  map (Prof n) = Nat (Prof \(Day f p q g) -> Day f (n p) q g)

instance (SymMonoidal j, SymMonoidal k, MonoidalProfunctor p, MonoidalProfunctor q) => MonoidalProfunctor (Day p q :: j +-> k) where
  one = Day leftUnitorInv one one leftUnitor \\ unitObj @j \\ unitObj @k
  Day f1 p1 q1 g1 ** Day f2 p2 q2 g2 =
    let f = swapInner' (src p1) (src q1) (src p2) (src q2) . (f1 ** f2)
        g = (g1 ** g2) . swapInner' (tgt p1) (tgt p2) (tgt q1) (tgt q2)
    in Day f (p1 ** p2) (q1 ** q2) g

instance (Monoidal j, Monoidal k) => MonoidalProfunctor (Prof :: CAT (j +-> k)) where
  one = id
  Prof m ** Prof n = Prof \(Day f p q g) -> Day f (m p) (n q) g

instance (Monoidal j, Monoidal k) => Monoidal (j +-> k) where
  type Unit = DayUnit
  type p ** q = Day p q
  withOb2 r = r
  leftUnitor = Prof \(Day f (DayUnit h i) q g) -> dimap (leftUnitor . (h ** src q) . f) (g . (i ** tgt q) . leftUnitorInv) q \\ q
  leftUnitorInv = Prof \q -> Day leftUnitorInv (DayUnit one one) q leftUnitor \\ q
  rightUnitor = Prof \(Day f p (DayUnit h i) g) -> dimap (rightUnitor . (src p ** h) . f) (g . (tgt p ** i) . rightUnitorInv) p \\ p
  rightUnitorInv = Prof \p -> Day rightUnitorInv p (DayUnit one one) rightUnitor \\ p
  associator = Prof \(Day @_ @_ @e1 @f1 f1 (Day @c2 @d2 @e2 @f2 f2 p2@Objs q2@Objs g2) q1@Objs g1) ->
    Day
      (associator @_ @c2 @e2 @e1 . (f2 ** src q1) . f1)
      p2
      (day q2 q1)
      (g1 . (g2 ** tgt q1) . associatorInv @_ @d2 @f2 @f1)
  associatorInv = Prof \(Day @c1 @d1 f1 p1@Objs (Day @c2 @d2 @e2 @f2 f2 p2@Objs q2@Objs g2) g1) ->
    Day
      (associatorInv @_ @c1 @c2 @e2 . (src p1 ** f2) . f1)
      (day p1 p2)
      q2
      (g1 . (tgt p1 ** g2) . associator @_ @d1 @d2 @f2)

instance (SymMonoidal j, SymMonoidal k) => SymMonoidal (j +-> k) where
  swap = Prof \(Day @c @d @e @f f p q g) -> Day (swap @_ @c @e . f) q p (g . swap @_ @f @d) \\ p \\ q

instance (Profunctor p, MonoidalProfunctor p) => Monoid p where
  mempty = Prof \(DayUnit f g) -> dimap f g one
  mappend = Prof \(Day f p q g) -> dimap f g (p ** q)

instance (SymMonoidal j, CopyDiscard k, Supplies CommutativeMonoid j, Profunctor p) => Comonoid (p :: j +-> k) where
  counit = Prof \p -> p // DayUnit discard mempty
  comult = Prof \p -> p // Day copy p p mappend
instance (SymMonoidal j, CopyDiscard k, Supplies CommutativeMonoid j, Profunctor p) => CocommutativeComonoid (p :: j +-> k)

instance (SymMonoidal j, CopyDiscard k, Supplies CommutativeMonoid j) => CopyDiscard (j +-> k)

instance (Monoidal j, Monoidal k) => Distributive (j +-> k) where
  distL = Prof \(Day l a bc r) -> case bc of
    InjL b -> InjL (Day l a b r)
    InjR c -> InjR (Day l a c r)
  distR = Prof \(Day l ab c r) -> case ab of
    InjL a -> InjL (Day l a c r)
    InjR b -> InjR (Day l b c r)
  absorbL = Prof \case {}
  absorbR = Prof \case {}

duoidal
  :: (Monoidal j, Profunctor (p :: j +-> k), Profunctor p', Profunctor q, Profunctor q')
  => (p :.: p') `Day` (q :.: q') ~> (p `Day` q) :.: (p' `Day` q')
duoidal = Prof \(Day f (p :.: p') (q :.: q') g) -> let b = tgt p ** tgt q in Day f p q b :.: Day b p' q' g

-- | The internal hom of 'Day' convolution, making the category of profunctors 'Closed'.
data DayExp p q a b where
  DayExp
    :: forall p q a b. (Ob a, Ob b) => (forall c d e f. e ~> a ** c -> b ** d ~> f -> p c d -> q e f) -> DayExp p q a b

instance (Monoidal j, Monoidal k, Profunctor (p :: j +-> k), Profunctor q) => Profunctor (DayExp p q) where
  dimap l r (DayExp n) = l // r // DayExp \f g p -> n ((l ** src p) . f) (g . (r ** tgt p)) p
  r \\ DayExp n = r \\ n

instance (Monoidal j, Monoidal k, Profunctor (p :: j +-> k)) => Functor (DayExp p) where
  map (Prof n) = Prof \(DayExp m) -> DayExp \f g p -> n (m f g p)

instance (Monoidal j, Monoidal k) => Closed (j +-> k) where
  type p ~~> q = DayExp p q
  withObExp r = r
  curry (Prof n) = Prof \p -> p // DayExp \f g q -> n (Day f p q g)
  apply = Prof \(Day f (DayExp h) q g) -> h f g q
  (^^^) (Prof n) (Prof m) = Prof \(DayExp k) -> DayExp \f g p -> n (k f g (m p))

multDayExp
  :: (SymMonoidal j, SymMonoidal k, Profunctor (p :: j +-> k), Profunctor q, Profunctor p', Profunctor q')
  => (p ~~> q) `Day` (p' ~~> q') ~> (p `Day` p') ~~> (q `Day` q')
multDayExp = Prof \(Day @c @d @e @f g (DayExp pq) (DayExp pq') h) ->
  g //
    h //
      DayExp
        ( \l r (Day i p p' j) ->
            p //
              p' //
                let c = obj @c; d = obj @d; e = obj @e; f = obj @f; c' = src p; d' = tgt p; e' = src p'; f' = tgt p'
                in Day
                     (swapInner' c e c' e' . (g ** i) . l)
                     (pq (c ** c') (d ** d') p)
                     (pq' (e ** e') (f ** f') p')
                     (r . (h ** j) . swapInner' d d' f f')
        )

-- Day p q :: j +-> k can be Strong in multiple ways:
-- 1. Either p or q is strong, and we have: Act a (b ** c) ~ (Act a b) ** c.
-- 2. Or p and q are both strong, and we have: Act a (b ** c) ~ (Act a b) ** (Act a c)

day2comp :: (MonStrong (p :: k +-> k), MonStrong q, Monoidal k) => Day p q ~> p :.: q
day2comp = Prof \(Day @_ @d @e f p q g) -> lmap f (first' @e p) :.: rmap g (second' @d q) \\ p \\ q