{-# LANGUAGE AllowAmbiguousTypes #-}

-- | Internal module of "Proarrow.Tools.SMC": functions and traces. It exports everything, also what
-- the public module keeps hidden.
module Proarrow.Tools.SMC.Internal.Closed where

import Data.Kind (Constraint, Type)
import Prelude (type (~))

import Proarrow.Category.Monoidal qualified as M
import Proarrow.Category.Monoidal.Closed (Closed (..))
import Proarrow.Category.Monoidal.Strength (Costrong, trace)
import Proarrow.Core (CAT, CategoryOf (..), Profunctor (..))

import Proarrow.Tools.SMC.Internal.Context
import Proarrow.Tools.SMC.Internal.Pattern
import Proarrow.Tools.SMC.Internal.Syntax
import Proarrow.Tools.SMC.Internal.Term

infixl 8 !

-- | A function: the body receives the argument through a pattern (see /Patterns/). This needs the
-- category to be closed.
{-# INLINE lam #-}
lam
  :: forall {k} d r (a :: SYN k) b t cont
   . (Closed k, Binds (~>) d r a b t cont)
  => (t -> cont)
  -> Term d r (a :-> b)
lam :: forall {k} (d :: Nat) (r :: Ctx k) (a :: SYN k) (b :: SYN k) t
       cont.
(Closed k, Binds (~>) d r a b t cont) =>
(t -> cont) -> Term d r (a :-> b)
lam t -> cont
k = forall (g :: Ctx k) r.
(KnownCtx g, Monoidal k) =>
(Ob (Interp (Mul g)) => r) -> r
forall {k} (g :: Ctx k) r.
(KnownCtx g, Monoidal k) =>
(Ob (Interp (Mul g)) => r) -> r
withCtxOb @r (forall {k} (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
forall (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
withSynOb @a ((Interp (Mul r) ~> Interp (a :-> b)) -> TermP (~>) d r (a :-> b)
forall {k} (p :: CAT k) (g :: Ctx k) (a :: SYN k) (d :: Nat).
p (Interp (Mul g)) (Interp a) -> TermP p d g a
MkTerm (forall k (a :: k) (b :: k) (c :: k).
(Closed k, Ob a, Ob b) =>
((a ** b) ~> c) -> a ~> (b ~~> c)
curry @k @(Interp (Mul r)) @(Interp a) (forall {k} (p :: CAT k) (n :: Nat) (g :: Ctx k) (a :: SYN k)
       (b :: SYN k) t cont.
Binds p n g a b t cont =>
(t -> cont) -> p (Interp (Mul g) ** Interp a) (Interp b)
forall (p :: CAT k) (n :: Nat) (g :: Ctx k) (a :: SYN k)
       (b :: SYN k) t cont.
Binds p n g a b t cont =>
(t -> cont) -> p (Interp (Mul g) ** Interp a) (Interp b)
bound @(~>) @d @r @a @b t -> cont
k))))

-- | Trace: the body receives the value fed back through a pattern (see /Patterns/), and returns
-- it again next to the result. This needs the category to be traced.
{-# INLINE loop #-}
loop
  :: forall {k} {p} (u :: SYN k) b d r t cont
   . (Costrong M.Tensor p, M.SymMonoidal k, KnownObj b, Binds p d r u (b :** u) t cont)
  => (t -> cont)
  -> TermP p d r b
loop :: forall {k} {p :: k +-> k} (u :: SYN k) (b :: SYN k) (d :: Nat)
       (r :: Ctx k) t cont.
(Costrong Tensor p, SymMonoidal k, KnownObj b,
 Binds p d r u (b :** u) t cont) =>
(t -> cont) -> TermP p d r b
loop t -> cont
k =
  forall (g :: Ctx k) r.
(KnownCtx g, Monoidal k) =>
(Ob (Interp (Mul g)) => r) -> r
forall {k} (g :: Ctx k) r.
(KnownCtx g, Monoidal k) =>
(Ob (Interp (Mul g)) => r) -> r
withCtxOb @r
    ( forall {k} (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
forall (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
withSynOb @u
        (forall {k} (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
forall (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
withSynOb @b (p (Interp (Mul r)) (Interp b) -> TermP p d r b
forall {k} (p :: CAT k) (g :: Ctx k) (a :: SYN k) (d :: Nat).
p (Interp (Mul g)) (Interp a) -> TermP p d g a
MkTerm (forall {k} (p :: k +-> k) (u :: k) (x :: k) (y :: k).
(Costrong Tensor p, Ob x, Ob y, Ob u, SymMonoidal k) =>
p (x ** u) (y ** u) -> p x y
forall (p :: k +-> k) (u :: k) (x :: k) (y :: k).
(Costrong Tensor p, Ob x, Ob y, Ob u, SymMonoidal k) =>
p (x ** u) (y ** u) -> p x y
trace @p @(Interp u) @(Interp (Mul r)) @(Interp b) (forall {k} (p :: CAT k) (n :: Nat) (g :: Ctx k) (a :: SYN k)
       (b :: SYN k) t cont.
Binds p n g a b t cont =>
(t -> cont) -> p (Interp (Mul g) ** Interp a) (Interp b)
forall (p :: k +-> k) (n :: Nat) (g :: Ctx k) (a :: SYN k)
       (b :: SYN k) t cont.
Binds p n g a b t cont =>
(t -> cont) -> p (Interp (Mul g) ** Interp a) (Interp b)
bound @p @d @r @u @(b :** u) t -> cont
k))))
    )

-- | Function application. A variable both the function and its argument use is copied.
-- Explicit exponentials determine their result syntax. Embedded function objects, such as
-- results of '(-<)', take their result syntax from the surrounding term.
type Apply :: forall {k}. CAT k -> Type -> Type -> Type -> Constraint
class Apply (p :: CAT k) f x r | f -> p where
  (!)
    :: (Closed k, M.MonoidalProfunctor p, KnownObj (TyOf @k x), KnownObj (TyOf @k r), Merge (CtxOf @k f) (CtxOf @k x))
    => f -> x -> r

instance (x ~ TermP p d g2 a, r ~ TermP p d (Union g1 g2) b) => Apply p (TermP p d g1 (a :-> b :: SYN k)) x r where
  {-# INLINE (!) #-}
  ! :: (Closed k, MonoidalProfunctor p, KnownObj (TyOf x),
 KnownObj (TyOf r),
 Merge (CtxOf (TermP p d g1 (a :-> b))) (CtxOf x)) =>
TermP p d g1 (a :-> b) -> x -> r
(!) = forall {k} (a :: SYN k) (b :: SYN k) (p :: k +-> k) (d :: Nat)
       (g1 :: Ctx k) (g2 :: Ctx k) (f :: SYN k).
(Closed k, MonoidalProfunctor p, KnownObj a, KnownObj b,
 Merge g1 g2, Interp f ~ (Interp a ~~> Interp b)) =>
TermP p d g1 f -> TermP p d g2 a -> TermP p d (Union g1 g2) b
forall (a :: SYN k) (b :: SYN k) (p :: CAT k) (d :: Nat)
       (g1 :: Ctx k) (g2 :: Ctx k) (f :: SYN k).
(Closed k, MonoidalProfunctor p, KnownObj a, KnownObj b,
 Merge g1 g2, Interp f ~ (Interp a ~~> Interp b)) =>
TermP p d g1 f -> TermP p d g2 a -> TermP p d (Union g1 g2) b
applyTerm @a @b

instance
  (f ~ (Interp a ~~> Interp b), x ~ TermP p d g2 a, r ~ TermP p d (Union g1 g2) b)
  => Apply p (TermP p d g1 (F f :: SYN k)) x r
  where
  {-# INLINE (!) #-}
  ! :: (Closed k, MonoidalProfunctor p, KnownObj (TyOf x),
 KnownObj (TyOf r), Merge (CtxOf (TermP p d g1 (F f))) (CtxOf x)) =>
TermP p d g1 (F f) -> x -> r
(!) = forall {k} (a :: SYN k) (b :: SYN k) (p :: k +-> k) (d :: Nat)
       (g1 :: Ctx k) (g2 :: Ctx k) (f :: SYN k).
(Closed k, MonoidalProfunctor p, KnownObj a, KnownObj b,
 Merge g1 g2, Interp f ~ (Interp a ~~> Interp b)) =>
TermP p d g1 f -> TermP p d g2 a -> TermP p d (Union g1 g2) b
forall (a :: SYN k) (b :: SYN k) (p :: CAT k) (d :: Nat)
       (g1 :: Ctx k) (g2 :: Ctx k) (f :: SYN k).
(Closed k, MonoidalProfunctor p, KnownObj a, KnownObj b,
 Merge g1 g2, Interp f ~ (Interp a ~~> Interp b)) =>
TermP p d g1 f -> TermP p d g2 a -> TermP p d (Union g1 g2) b
applyTerm @a @b

{-# INLINE applyTerm #-}
applyTerm
  :: forall {k} (a :: SYN k) b p d g1 g2 f
   . (Closed k, M.MonoidalProfunctor p, KnownObj a, KnownObj b, Merge g1 g2, Interp f ~ (Interp a ~~> Interp b))
  => TermP p d g1 f
  -> TermP p d g2 a
  -> TermP p d (Union g1 g2) b
applyTerm :: forall {k} (a :: SYN k) (b :: SYN k) (p :: k +-> k) (d :: Nat)
       (g1 :: Ctx k) (g2 :: Ctx k) (f :: SYN k).
(Closed k, MonoidalProfunctor p, KnownObj a, KnownObj b,
 Merge g1 g2, Interp f ~ (Interp a ~~> Interp b)) =>
TermP p d g1 f -> TermP p d g2 a -> TermP p d (Union g1 g2) b
applyTerm (MkTerm p (Interp (Mul g1)) (Interp f)
f) (MkTerm p (Interp (Mul g2)) (Interp a)
x) =
  forall {k} (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
forall (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
withSynOb @a
    ( forall {k} (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
forall (s :: SYN k) r. KnownObj s => (Ob (Interp s) => r) -> r
withSynOb @b
        (p (Interp (Mul (Union g1 g2))) (Interp b)
-> TermP p d (Union g1 g2) b
forall {k} (p :: CAT k) (g :: Ctx k) (a :: SYN k) (d :: Nat).
p (Interp (Mul g)) (Interp a) -> TermP p d g a
MkTerm ((Interp (Mul (Union g1 g2))
 ~> (Interp (Mul g1) ** Interp (Mul g2)))
-> ((Interp f ** Interp a) ~> Interp b)
-> p (Interp (Mul g1) ** Interp (Mul g2)) (Interp f ** Interp a)
-> p (Interp (Mul (Union g1 g2))) (Interp b)
forall (c :: k) (a :: k) (b :: k) (d :: k).
(c ~> a) -> (b ~> d) -> p a b -> p c d
forall {j} {k} (p :: j +-> k) (c :: k) (a :: k) (b :: j) (d :: j).
Profunctor p =>
(c ~> a) -> (b ~> d) -> p a b -> p c d
dimap (forall (g1 :: Ctx k) (g2 :: Ctx k).
Merge g1 g2 =>
Interp (Mul (Union g1 g2)) ~> (Interp (Mul g1) ** Interp (Mul g2))
forall {k} (g1 :: Ctx k) (g2 :: Ctx k).
Merge g1 g2 =>
Interp (Mul (Union g1 g2)) ~> (Interp (Mul g1) ** Interp (Mul g2))
merge @g1 @g2) (forall k (a :: k) (b :: k).
(Closed k, Ob a, Ob b) =>
((a ~~> b) ** a) ~> b
apply @k @(Interp a) @(Interp b)) (p (Interp (Mul g1)) (Interp f)
f p (Interp (Mul g1)) (Interp f)
-> p (Interp (Mul g2)) (Interp a)
-> p (Interp (Mul g1) ** Interp (Mul g2)) (Interp f ** Interp a)
forall (x1 :: k) (x2 :: k) (y1 :: k) (y2 :: k).
p x1 x2 -> p y1 y2 -> p (x1 ** y1) (x2 ** y2)
forall {j} {k} (p :: j +-> k) (x1 :: k) (x2 :: j) (y1 :: k)
       (y2 :: j).
MonoidalProfunctor p =>
p x1 x2 -> p y1 y2 -> p (x1 ** y1) (x2 ** y2)
M.** p (Interp (Mul g2)) (Interp a)
x)))
    )