{-# LANGUAGE AllowAmbiguousTypes #-}

-- | Relations and weighted graphs on a bare set of points, given as a table: a list of edges with
-- their weights in the enriching category @v@. 'Edges' is an enriched profunctor on the discrete
-- category of an 'Indexed' kind (over a discrete base there is nothing to be compatible with), so
-- it composes, and its Kleene closure 'Proarrow.Category.Enriched.Thin.Composition.Closure' is
-- reachability for 'BOOL' weights and shortest paths for 'COST' weights.
module Proarrow.Profunctor.Instance.Edges where

import Data.Kind (Constraint)
import Data.Type.Equality qualified as Eq
import Prelude (type (~))

import Proarrow.Category.Enriched (EnrichedProfunctor (..))
import Proarrow.Category.Enriched.Quantale (Quantale (..))
import Proarrow.Category.Enriched.Thin
  ( DecidableProfunctor (..)
  , Decision (..)
  , Equal
  , Indexed
  , KnownIndex
  , ThinProfunctor
  , decideEq
  )
import Proarrow.Category.Instance.Bool (BOOL (..), Booleans (..), If)
import Proarrow.Category.Instance.Cost (COST)
import Proarrow.Category.Instance.Discrete (DISCRETE (..), Discrete (..), deltaAct)
import Proarrow.Category.Monoidal (Monoidal (..))
import Proarrow.Colimit.Initial (HasInitialObject (..))
import Proarrow.Core (CategoryOf (..), Kind, Profunctor (..), obj, type (+->))
import Proarrow.Limit.BinaryProduct (type (&&))

-- | A weighted graph on the bare set of points of an 'Indexed' kind, given as a list of edges with
-- their weights in @v@: an enriched profunctor on the discrete category, since over a discrete base
-- there is nothing to be compatible with. Unlisted pairs are at 'InitialObject', a pair listed twice
-- takes its first weight, and an element is a pair at 'Unit' weight.
type Edges :: forall {k} {v}. [(k, k, v)] -> DISCRETE k +-> DISCRETE k
data Edges es a b where
  Edge :: (Ob a, Ob b, WeightOf es a b ~ Unit) => Edges es a b

type WeightOf :: forall {k} {v}. [(k, k, v)] -> DISCRETE k -> DISCRETE k -> v
type family WeightOf es a b where
  WeightOf '[] a b = InitialObject
  WeightOf ('(x, y, w) ': es) a b = If (Equal a (D x) && Equal b (D y)) w (WeightOf es a b)

instance (Indexed k) => Profunctor (Edges (es :: [(k, k, v)])) where
  dimap :: forall (c :: DISCRETE k) (a :: DISCRETE k) (b :: DISCRETE k)
       (d :: DISCRETE k).
(c ~> a) -> (b ~> d) -> Edges es a b -> Edges es c d
dimap c ~> a
Discrete c a
Refl b ~> d
Discrete b d
Refl Edges es a b
e = Edges es c d
Edges es a b
e
  (Ob a, Ob b) => r
r \\ :: forall (a :: DISCRETE k) (b :: DISCRETE k) r.
((Ob a, Ob b) => r) -> Edges es a b -> r
\\ Edges es a b
Edge = r
(Ob a, Ob b) => r
r

-- | The edge list, reflected to the value level.
type EdgeList :: forall {k} {v}. [(k, k, v)] -> Kind
data EdgeList es where
  ENil :: EdgeList '[]
  ECons :: forall x y w es. (KnownIndex x, KnownIndex y, Ob w) => EdgeList es -> EdgeList ('(x, y, w) ': es)

type KnownEdges :: forall {k} {v}. [(k, k, v)] -> Constraint
class KnownEdges es where
  edges :: EdgeList es
instance KnownEdges '[] where
  edges :: EdgeList '[]
edges = EdgeList '[]
forall k v. EdgeList '[]
ENil
instance (KnownIndex x, KnownIndex y, Ob w, KnownEdges es) => KnownEdges ('(x, y, w) ': es) where
  edges :: EdgeList ('(x, y, w) : es)
edges = EdgeList es -> EdgeList ('(x, y, w) : es)
forall {k} {v} (x :: k) (y :: k) (w :: v) (es :: [(k, k, v)]).
(KnownIndex x, KnownIndex y, Ob w) =>
EdgeList es -> EdgeList ('(x, y, w) : es)
ECons EdgeList es
forall {k} {v} (es :: [(k, k, v)]). KnownEdges es => EdgeList es
edges

-- | A graph with 'BOOL' weights is a relation on the points: decided by walking the edge list.
instance (Indexed k, KnownEdges es) => ThinProfunctor (Edges (es :: [(k, k, BOOL)]))

instance (Indexed k, KnownEdges es) => DecidableProfunctor (Edges (es :: [(k, k, BOOL)])) where
  type Holds (Edges es) a b = WeightOf es a b
  decide :: forall (a :: DISCRETE k) (b :: DISCRETE k).
(Ob a, Ob b) =>
Decision (Edges es) a b (Holds (Edges es) a b)
decide @a @b = EdgeList es -> Decision (Edges es) a b (WeightOf es a b)
forall (es' :: [(k, k, BOOL)]).
(WeightOf es' a b ~ WeightOf es a b) =>
EdgeList es' -> Decision (Edges es) a b (WeightOf es' a b)
go (forall (es :: [(k, k, BOOL)]). KnownEdges es => EdgeList es
forall {k} {v} (es :: [(k, k, v)]). KnownEdges es => EdgeList es
edges @es)
    where
      go
        :: forall (es' :: [(k, k, BOOL)])
         . (WeightOf es' a b ~ WeightOf es a b)
        => EdgeList es' -> Decision (Edges es) a b (WeightOf es' a b)
      go :: forall (es' :: [(k, k, BOOL)]).
(WeightOf es' a b ~ WeightOf es a b) =>
EdgeList es' -> Decision (Edges es) a b (WeightOf es' a b)
go EdgeList es'
ENil = Decision (Edges es) a b 'FLS
Decision (Edges es) a b (WeightOf es' a b)
forall {j} {k} (p :: j +-> k) (a :: k) (b :: j).
Decision p a b 'FLS
No
      go (ECons @x @y @w EdgeList es
es') = case (forall {k} (a :: k) (b :: k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
forall (a :: DISCRETE k) (b :: DISCRETE k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
decideEq @a @(D x), forall {k} (a :: k) (b :: k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
forall (a :: DISCRETE k) (b :: DISCRETE k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
decideEq @b @(D y)) of
        (Yes a :~: D x
Eq.Refl, Yes b :~: D y
Eq.Refl) -> case forall {k} (a :: k). (CategoryOf k, Ob a) => Obj a
forall (a :: BOOL). (CategoryOf BOOL, Ob a) => Obj a
obj @w of
          Obj w
Booleans w w
Tru -> Edges es a b -> Decision (Edges es) a b 'TRU
forall {j} {k} (p :: j +-> k) (a :: k) (b :: j).
p a b -> Decision p a b 'TRU
Yes Edges es a b
forall {k} {v} (a :: DISCRETE k) (b :: DISCRETE k)
       (es :: [(k, k, v)]).
(Ob a, Ob b, WeightOf es a b ~ Unit) =>
Edges es a b
Edge
          Obj w
Booleans w w
Fls -> Decision (Edges es) a b 'FLS
Decision (Edges es) a b (WeightOf es' a b)
forall {j} {k} (p :: j +-> k) (a :: k) (b :: j).
Decision p a b 'FLS
No
        (Decision (:~:) a (D x) (NatEq (Index (UN D a)) (Index x))
No, Decision (:~:) b (D y) (NatEq (Index (UN D b)) (Index y))
_) -> EdgeList es -> Decision (Edges es) a b (WeightOf es a b)
forall (es' :: [(k, k, BOOL)]).
(WeightOf es' a b ~ WeightOf es a b) =>
EdgeList es' -> Decision (Edges es) a b (WeightOf es' a b)
go EdgeList es
es'
        (Yes a :~: D x
_, Decision (:~:) b (D y) (NatEq (Index (UN D b)) (Index y))
No) -> EdgeList es -> Decision (Edges es) a b (WeightOf es a b)
forall (es' :: [(k, k, BOOL)]).
(WeightOf es' a b ~ WeightOf es a b) =>
EdgeList es' -> Decision (Edges es) a b (WeightOf es' a b)
go EdgeList es
es'
  toHolds :: forall (a :: DISCRETE k) (b :: DISCRETE k) r.
Edges es a b
-> ((Holds (Edges es) a b ~ 'TRU, Ob a, Ob b) => r) -> r
toHolds Edges es a b
Edge (Holds (Edges es) a b ~ 'TRU, Ob a, Ob b) => r
r = r
(Holds (Edges es) a b ~ 'TRU, Ob a, Ob b) => r
r

-- | The weight of a pair, reflected to the value level by walking the edge list.
withObWeight
  :: forall {k} {v} (es :: [(k, k, v)]) a b r
   . (Quantale v, Indexed k, KnownEdges es, KnownIndex a, KnownIndex b)
  => ((Ob (WeightOf es a b)) => r) -> r
withObWeight :: forall {k} {v} (es :: [(k, k, v)]) (a :: DISCRETE k)
       (b :: DISCRETE k) r.
(Quantale v, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
withObWeight Ob (WeightOf es a b) => r
r = EdgeList es -> (Ob (WeightOf es a b) => r) -> r
forall (es' :: [(k, k, v)]).
EdgeList es' -> (Ob (WeightOf es' a b) => r) -> r
go (forall (es :: [(k, k, v)]). KnownEdges es => EdgeList es
forall {k} {v} (es :: [(k, k, v)]). KnownEdges es => EdgeList es
edges @es) r
Ob (WeightOf es a b) => r
r
  where
    go :: forall (es' :: [(k, k, v)]). EdgeList es' -> ((Ob (WeightOf es' a b)) => r) -> r
    go :: forall (es' :: [(k, k, v)]).
EdgeList es' -> (Ob (WeightOf es' a b) => r) -> r
go EdgeList es'
ENil Ob (WeightOf es' a b) => r
r' = r
Ob (WeightOf es' a b) => r
r'
    go (ECons @x @y EdgeList es
es') Ob (WeightOf es' a b) => r
r' = case (forall {k} (a :: k) (b :: k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
forall (a :: DISCRETE k) (b :: DISCRETE k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
decideEq @a @(D x), forall {k} (a :: k) (b :: k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
forall (a :: DISCRETE k) (b :: DISCRETE k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
decideEq @b @(D y)) of
      (Yes a :~: D x
Eq.Refl, Yes b :~: D y
Eq.Refl) -> r
Ob (WeightOf es' a b) => r
r'
      (Decision (:~:) a (D x) (NatEq (Index (UN D a)) (Index x))
No, Decision (:~:) b (D y) (NatEq (Index (UN D b)) (Index y))
_) -> EdgeList es -> (Ob (WeightOf es a b) => r) -> r
forall (es' :: [(k, k, v)]).
EdgeList es' -> (Ob (WeightOf es' a b) => r) -> r
go EdgeList es
es' r
Ob (WeightOf es' a b) => r
Ob (WeightOf es a b) => r
r'
      (Yes a :~: D x
_, Decision (:~:) b (D y) (NatEq (Index (UN D b)) (Index y))
No) -> EdgeList es -> (Ob (WeightOf es a b) => r) -> r
forall (es' :: [(k, k, v)]).
EdgeList es' -> (Ob (WeightOf es' a b) => r) -> r
go EdgeList es
es' r
Ob (WeightOf es' a b) => r
Ob (WeightOf es a b) => r
r'

-- | A unit into a weight is an edge at the unit, since an object above the unit is the unit.
enrichedEdge
  :: forall {k} {v} (es :: [(k, k, v)]) a b
   . (Quantale v, Indexed k, KnownEdges es, KnownIndex a, KnownIndex b)
  => Unit ~> WeightOf es a b -> Edges es a b
enrichedEdge :: forall {k} {v} (es :: [(k, k, v)]) (a :: DISCRETE k)
       (b :: DISCRETE k).
(Quantale v, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Unit ~> WeightOf es a b) -> Edges es a b
enrichedEdge Unit ~> WeightOf es a b
f = EdgeList es -> (Unit ~> WeightOf es a b) -> Edges es a b
forall (es' :: [(k, k, v)]).
(WeightOf es' a b ~ WeightOf es a b) =>
EdgeList es' -> (Unit ~> WeightOf es' a b) -> Edges es a b
go (forall (es :: [(k, k, v)]). KnownEdges es => EdgeList es
forall {k} {v} (es :: [(k, k, v)]). KnownEdges es => EdgeList es
edges @es) Unit ~> WeightOf es a b
f
  where
    go
      :: forall (es' :: [(k, k, v)])
       . (WeightOf es' a b ~ WeightOf es a b)
      => EdgeList es' -> Unit ~> WeightOf es' a b -> Edges es a b
    go :: forall (es' :: [(k, k, v)]).
(WeightOf es' a b ~ WeightOf es a b) =>
EdgeList es' -> (Unit ~> WeightOf es' a b) -> Edges es a b
go EdgeList es'
ENil Unit ~> WeightOf es' a b
g = forall v r. Quantale v => (Unit ~> InitialObject) -> r
unitIsNotBottom @v Unit ~> InitialObject
Unit ~> WeightOf es' a b
g
    go (ECons @x @y @w EdgeList es
es') Unit ~> WeightOf es' a b
g = case (forall {k} (a :: k) (b :: k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
forall (a :: DISCRETE k) (b :: DISCRETE k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
decideEq @a @(D x), forall {k} (a :: k) (b :: k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
forall (a :: DISCRETE k) (b :: DISCRETE k).
(KnownIndex a, KnownIndex b) =>
Decision (:~:) a b (Equal a b)
decideEq @b @(D y)) of
      (Yes a :~: D x
Eq.Refl, Yes b :~: D y
Eq.Refl) -> forall v (w :: v) r.
(Quantale v, Ob w) =>
(Unit ~> w) -> ((w ~ Unit) => r) -> r
unitIsTop @v @w Unit ~> w
Unit ~> WeightOf es' a b
g Edges es a b
(w ~ Unit) => Edges es a b
forall {k} {v} (a :: DISCRETE k) (b :: DISCRETE k)
       (es :: [(k, k, v)]).
(Ob a, Ob b, WeightOf es a b ~ Unit) =>
Edges es a b
Edge
      (Decision (:~:) a (D x) (NatEq (Index (UN D a)) (Index x))
No, Decision (:~:) b (D y) (NatEq (Index (UN D b)) (Index y))
_) -> EdgeList es -> (Unit ~> WeightOf es a b) -> Edges es a b
forall (es' :: [(k, k, v)]).
(WeightOf es' a b ~ WeightOf es a b) =>
EdgeList es' -> (Unit ~> WeightOf es' a b) -> Edges es a b
go EdgeList es
es' Unit ~> WeightOf es' a b
Unit ~> WeightOf es a b
g
      (Yes a :~: D x
_, Decision (:~:) b (D y) (NatEq (Index (UN D b)) (Index y))
No) -> EdgeList es -> (Unit ~> WeightOf es a b) -> Edges es a b
forall (es' :: [(k, k, v)]).
(WeightOf es' a b ~ WeightOf es a b) =>
EdgeList es' -> (Unit ~> WeightOf es' a b) -> Edges es a b
go EdgeList es
es' Unit ~> WeightOf es' a b
Unit ~> WeightOf es a b
g

-- | A graph with 'COST' weights: a weighted graph, whose closure is shortest paths.
instance (Indexed k, KnownEdges es) => EnrichedProfunctor COST (Edges (es :: [(k, k, COST)])) where
  type ProObj COST (Edges es) a b = WeightOf es a b
  withProObj :: forall (a :: DISCRETE k) (b :: DISCRETE k) r.
(Ob a, Ob b) =>
(Ob (ProObj COST (Edges es) a b) => r) -> r
withProObj @a @b = forall (es :: [(k, k, COST)]) (a :: DISCRETE k) (b :: DISCRETE k)
       r.
(Quantale COST, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
forall {k} {v} (es :: [(k, k, v)]) (a :: DISCRETE k)
       (b :: DISCRETE k) r.
(Quantale v, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
withObWeight @es @a @b
  underlying :: forall (a :: DISCRETE k) (b :: DISCRETE k).
Edges es a b -> Unit ~> ProObj COST (Edges es) a b
underlying Edges es a b
Edge = forall {k} (a :: k). (CategoryOf k, Ob a) => Obj a
forall (a :: COST). (CategoryOf COST, Ob a) => Obj a
obj @(Unit :: COST)
  enriched :: forall (a :: DISCRETE k) (b :: DISCRETE k).
(Ob a, Ob b) =>
(Unit ~> ProObj COST (Edges es) a b) -> Edges es a b
enriched @a @b = forall (es :: [(k, k, COST)]) (a :: DISCRETE k) (b :: DISCRETE k).
(Quantale COST, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Unit ~> WeightOf es a b) -> Edges es a b
forall {k} {v} (es :: [(k, k, v)]) (a :: DISCRETE k)
       (b :: DISCRETE k).
(Quantale v, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Unit ~> WeightOf es a b) -> Edges es a b
enrichedEdge @es @a @b
  rmap :: forall (a :: DISCRETE k) (b :: DISCRETE k) (c :: DISCRETE k).
(Ob a, Ob b, Ob c) =>
(HomObj COST b c ** ProObj COST (Edges es) a b)
~> ProObj COST (Edges es) a c
rmap @a @b @c =
    forall (es :: [(k, k, COST)]) (a :: DISCRETE k) (b :: DISCRETE k)
       r.
(Quantale COST, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
forall {k} {v} (es :: [(k, k, v)]) (a :: DISCRETE k)
       (b :: DISCRETE k) r.
(Quantale v, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
withObWeight @es @a @b (forall (es :: [(k, k, COST)]) (a :: DISCRETE k) (b :: DISCRETE k)
       r.
(Quantale COST, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
forall {k} {v} (es :: [(k, k, v)]) (a :: DISCRETE k)
       (b :: DISCRETE k) r.
(Quantale v, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
withObWeight @es @a @c (forall {k} {v} (x :: k) (y :: k) (w :: v) (w' :: v).
(Quantale v, KnownIndex x, KnownIndex y, Ob w, Ob w') =>
((x ~ y) => w :~: w') -> (Delta v (Equal x y) ** w) ~> w'
forall (x :: DISCRETE k) (y :: DISCRETE k) (w :: COST)
       (w' :: COST).
(Quantale COST, KnownIndex x, KnownIndex y, Ob w, Ob w') =>
((x ~ y) => w :~: w') -> (Delta COST (Equal x y) ** w) ~> w'
deltaAct @b @c @(WeightOf es a b) @(WeightOf es a c) WeightOf es a b :~: WeightOf es a c
WeightOf es a c :~: WeightOf es a c
(b ~ c) => WeightOf es a b :~: WeightOf es a c
forall {k} (a :: k). a :~: a
Eq.Refl))
  lmap :: forall (a :: DISCRETE k) (b :: DISCRETE k) (c :: DISCRETE k).
(Ob a, Ob b, Ob c) =>
(HomObj COST c a ** ProObj COST (Edges es) a b)
~> ProObj COST (Edges es) c b
lmap @a @b @c =
    forall (es :: [(k, k, COST)]) (a :: DISCRETE k) (b :: DISCRETE k)
       r.
(Quantale COST, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
forall {k} {v} (es :: [(k, k, v)]) (a :: DISCRETE k)
       (b :: DISCRETE k) r.
(Quantale v, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
withObWeight @es @a @b (forall (es :: [(k, k, COST)]) (a :: DISCRETE k) (b :: DISCRETE k)
       r.
(Quantale COST, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
forall {k} {v} (es :: [(k, k, v)]) (a :: DISCRETE k)
       (b :: DISCRETE k) r.
(Quantale v, Indexed k, KnownEdges es, KnownIndex a,
 KnownIndex b) =>
(Ob (WeightOf es a b) => r) -> r
withObWeight @es @c @b (forall {k} {v} (x :: k) (y :: k) (w :: v) (w' :: v).
(Quantale v, KnownIndex x, KnownIndex y, Ob w, Ob w') =>
((x ~ y) => w :~: w') -> (Delta v (Equal x y) ** w) ~> w'
forall (x :: DISCRETE k) (y :: DISCRETE k) (w :: COST)
       (w' :: COST).
(Quantale COST, KnownIndex x, KnownIndex y, Ob w, Ob w') =>
((x ~ y) => w :~: w') -> (Delta COST (Equal x y) ** w) ~> w'
deltaAct @c @a @(WeightOf es a b) @(WeightOf es c b) WeightOf es a b :~: WeightOf es a b
WeightOf es a b :~: WeightOf es c b
(c ~ a) => WeightOf es a b :~: WeightOf es c b
forall {k} (a :: k). a :~: a
Eq.Refl))