Ran.hs

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

-- | The right Kan extension of a profunctor @p@ along @j@, written @j '|>' p@: the universal @g@ with
-- @g ':.:' j ~> p@. 'Ran' and 'Proarrow.Profunctor.Instance.Rift.Rift' are swapped compared to the
-- @profunctors@ package.
module Proarrow.Profunctor.Instance.Ran where

import Prelude (type (~))

import Proarrow.Category.Instance.Nat (Nat (..))
import Proarrow.Category.Instance.Opposite (OPPOSITE (..), Op (..))
import Proarrow.Category.Instance.Prof (Prof (..))
import Proarrow.Core (CategoryOf (..), Profunctor (..), Promonad (..), lmap, rmap, (//), type (+->))
import Proarrow.Functor (Functor (..), FunctorForRep)
import Proarrow.Limit (HasLimits (..))
import Proarrow.Profunctor.Corepresentable (Corep (..), Corepresentable (..), corepUniv, withObCorep)
import Proarrow.Profunctor.Instance.Composition ((:.:) (..))
import Proarrow.Profunctor.Instance.Star (Star, pattern Star)
import Proarrow.Profunctor.Representable (CorepStar, Rep (..), Representable (..), repUniv, withObRep)
import Proarrow.Promonad (Procomonad (..), RelativeMonad (..))

type j |> p = Ran (OP j) p

-- | The data type behind @j '|>' p@. Its universal property is 'ranUniv' and 'runRanProf'.
type Ran :: OPPOSITE (i +-> j) -> i +-> k -> j +-> k
data Ran j p a b where
  Ran :: (Ob a, Ob b) => {unRan :: forall x. j b x -> p a x} -> Ran (OP j) p a b

runRan :: (Profunctor j) => j b x -> Ran (OP j) p a b -> p a x
runRan j (Ran k) = k j \\ j

runRanProf :: (Profunctor j, Profunctor p) => (j |> p) :.: j ~> p
runRanProf = Prof \(r :.: j) -> runRan j r

ranUniv :: (Profunctor j, Profunctor g) => (g :.: j) ~> p -> g ~> j |> p
ranUniv (Prof n) = Prof \g -> g // Ran \j -> n (g :.: j)

flipRan :: (FunctorForRep j, Profunctor p) => Corep j |> p ~> p :.: Rep j
flipRan = Prof \(Ran k) -> k corepUniv :.: repUniv

flipRanInv :: (FunctorForRep j, Profunctor p) => p :.: Rep j ~> Corep j |> p
flipRanInv = Prof \(p :.: f) -> p // f // Ran \g -> rmap (coindex g . index f) p

instance (Profunctor p, Profunctor j) => Profunctor (Ran (OP j) p) where
  dimap l r (Ran k) = l // r // Ran (lmap l . k . lmap r)
  r \\ Ran{} = r

instance (Profunctor j) => Functor (Ran (OP j)) where
  map (Prof n) = Prof \(Ran k) -> Ran (n . k)

instance Functor Ran where
  map (Op (Prof n)) = Nat (Prof \(Ran k) -> Ran (k . n))

instance (p ~ j, Profunctor p) => Promonad (Ran (OP j) p) where
  id = Ran id
  Ran l . Ran r = Ran (r . l)

instance (HasLimits j k, Representable d) => Representable (Ran (OP j) (d :: i +-> k)) where
  type Ran (OP j) d % a = Limit j d % a
  index = index @(Limit j d) . limitUniv @j @k @d (\(Ran k' :.: j) -> k' j)
  repUniv @a = withObRep @(Limit j d) @a (Ran \j -> limit (repUniv :.: j))

type PWRan j p a = (j |> p) % a

-- a ~> PWRan j p b = forall x. (b ~> j x) -> a ~> p x
-- a ~> PWRan j p b = forall x. a ~> (p x ^ (b ~> j x))
-- PWRan j p b = forall x. p x ^ (b ~> j x)
class (Representable p, Representable j, Representable (j |> p)) => PointwiseRightKanExtension j p
instance (Representable p, Representable j, Representable (j |> p)) => PointwiseRightKanExtension j p

type PWLift j p a = (j |> p) %% a

-- PWLift j p a ~> b = forall x. (j b ~> x) -> p a ~> x
-- PWLift j p a ~> b = p a ~> j b
class (Corepresentable j, Corepresentable p, Corepresentable (j |> p)) => PointwiseLeftKanLift j p
instance (Corepresentable j, Corepresentable p, Corepresentable (j |> p)) => PointwiseLeftKanLift j p

instance (Corepresentable g, Corepresentable f, Profunctor j, f ~ g |> j) => RelativeMonad j (CorepStar g :.: CorepStar f) where
  relReturn @a = let f = corepUniv @f @a in runRan (corepUniv @g) f \\ f
  relBind @b @a j = withObCorep @f @b (corepMap @g @(f %% a) @(f %% b) (coindex @f @a (Ran (\g -> rmap (coindex g) j)))) \\ j

-- | The right Kan extension is the right adjoint of the precomposition functor.
instance (Profunctor j) => Corepresentable (Star (Ran (OP j))) where
  type Star (Ran (OP j)) %% p = p :.: j
  coindex (Star (Prof n)) = Prof \(p :.: j) -> runRan j (n p)
  cotabulate (Prof n) = Star (Prof \p -> p // Ran \q -> n (p :.: q))
  corepMap f = unNat (map f)

ranCompose :: (Profunctor i, Profunctor j, Profunctor p) => i |> (j |> p) ~> (i :.: j) |> p
ranCompose = Prof \k -> k // Ran \(i :.: j) -> runRan j (runRan i k)

ranComposeInv :: (Profunctor i, Profunctor j, Profunctor p) => (i :.: j) |> p ~> i |> (j |> p)
ranComposeInv = Prof \k -> k // Ran \i -> i // Ran \j -> runRan (i :.: j) k

ranHom :: (Profunctor p) => p ~> (~>) |> p
ranHom = Prof \p -> p // Ran (`rmap` p)

ranHomInv :: (Profunctor p) => (~>) |> p ~> p
ranHomInv = Prof \(Ran k) -> k id

compAsRan :: (Promonad p, Ob a) => (p |> (p |> p)) a a
compAsRan = Ran \p -> p // Ran (. p)

instance (Procomonad j) => Promonad (Star (Ran (OP j))) where
  id = Star (unNat (map (Op (Prof proextract))) . ranHom)
  Star l . Star r = Star (unNat (map (Op (Prof produplicate))) . ranCompose . map l . r)