-- | A structure representing unique vertices and their interrelations
module Mini.Data.Graph (
  -- * A Note on Performance
  -- $note

  -- * Type
  Graph,

  -- * Primitive Recursion
  graph,

  -- * Algorithms
  distance,
  layers,
  path,
  reachable,
  sort,

  -- * Construction
  empty,
  fromList,
  singleton,

  -- * Modification
  add,
  remove,
  connect,
  disconnect,
  transpose,

  -- * Query
  assocs,
  edges,
  vertices,
  indegree,
  indegrees,
  outdegree,
  outdegrees,
  lookup,
  lookupGE,
  lookupGT,
  lookupLE,
  lookupLT,
  lookupMax,
  lookupMin,
  member,
  sourceMax,
  sourceMin,
  sources,
  sinkMax,
  sinkMin,
  sinks,
) where

import Data.Bifunctor (
  second,
 )
import Data.Bool (
  bool,
 )
import Data.List (
  unfoldr,
 )
import Mini.Data.Map (
  Map,
 )
import qualified Mini.Data.Map as Map (
  delete,
  empty,
  foldlWithKey,
  foldrWithKey,
  insertWith,
  lookup,
  lookupGE,
  lookupGT,
  lookupLE,
  lookupLT,
  lookupMax,
  lookupMin,
  member,
  singleton,
  toAscList,
  unionWith,
 )
import Mini.Data.Set (
  Set,
 )
import qualified Mini.Data.Set as Set (
  delete,
  difference,
  empty,
  fromList,
  member,
  null,
  singleton,
  size,
  toAscList,
 )
import Prelude (
  Bool,
  Eq,
  Foldable,
  Int,
  Maybe (
    Just,
    Nothing
  ),
  Monoid,
  Ord,
  Semigroup,
  Show,
  any,
  compare,
  concat,
  flip,
  fmap,
  foldMap,
  foldr,
  fst,
  maybe,
  mempty,
  show,
  uncurry,
  ($),
  (+),
  (.),
  (<$>),
  (<>),
  (==),
 )

{-
 - A Note on Performance
 -}

{- $note
In order to provide a friendly user interface, some performance has been
sacrificed. The internal adjacency lists are implemented via maps rather than
arrays, meaning accesses are done in logarithmic time rather than constant time.
-}

{-
 - Type
 -}

-- | A graph with directed edges between vertices of type /a/
data Graph a
  = Graph
      (Map a (Set a))
      (Map a (Set a))

instance (Eq a) => Eq (Graph a) where
  (Graph _ oes1) == (Graph _ oes2) = oes1 == oes2

instance (Ord a) => Ord (Graph a) where
  compare (Graph _ oes1) (Graph _ oes2) = compare oes1 oes2

instance (Show a) => Show (Graph a) where
  show = show . assocs

instance Foldable Graph where
  foldr f b = foldr f b . vertices

instance (Ord a) => Semigroup (Graph a) where
  (Graph ies1 oes1) <> (Graph ies2 oes2) =
    Graph
      (Map.unionWith (<>) ies1 ies2)
      (Map.unionWith (<>) oes1 oes2)

instance (Ord a) => Monoid (Graph a) where
  mempty = empty

{-
 - Primitive Recursion
 -}

-- | Primitive recursion on graphs (internally represented by adjacency lists)
graph
  :: (Map a (Set a) -> Map a (Set a) -> b)
  -- ^ Function applied to the adjacency lists of the graph:
  -- incoming edges, outgoing edges
  -> Graph a
  -- ^ The graph
  -> b
graph f (Graph ies oes) = f ies oes

{-
 - Algorithms
 -}

-- | Get the shortest distance in a graph between a vertex and another
distance :: (Ord a) => Graph a -> a -> a -> Maybe Int
distance g s t =
  foldr
    (\a b -> bool ((+ 1) <$> b) (Just 0) $ t `Set.member` a)
    Nothing
    $ bfs g s

-- | Breadth-first search for the hierarchy in a graph from a starting vertex
layers :: (Ord a) => Graph a -> a -> [[a]]
layers g = fmap Set.toAscList . bfs g

-- | Check whether there is a path in a graph from a vertex to another
path :: (Ord a) => Graph a -> a -> a -> Bool
path g s t = any (t `Set.member`) $ bfs g s

-- | Get the reachable vertices in a graph from a starting vertex
reachable :: (Ord a) => Graph a -> a -> [a]
reachable g = concat . layers g

-- | Topologically sort a graph (assumes acyclicity)
sort :: (Ord a) => Graph a -> [a]
sort = unfoldr $ \g -> (\u -> (u, remove u g)) <$> sourceMin g

{-
 - Construction
 -}

-- | The empty graph
empty :: Graph a
empty = Graph Map.empty Map.empty

-- | Make a graph from a list of vertex associations
fromList :: (Ord a) => [(a, [a])] -> Graph a
fromList = foldr (uncurry connect) empty

-- | Make a graph with an isolated vertex
singleton :: a -> Graph a
singleton u = Graph (Map.singleton u Set.empty) (Map.singleton u Set.empty)

{-
 - Modification
 -}

-- | Add an isolated vertex to a graph unless already present
add :: (Ord a) => a -> Graph a -> Graph a
add u = connect u []

-- | Remove a vertex and its associations from a graph
remove :: (Ord a) => a -> Graph a -> Graph a
remove u (Graph ies oes) =
  Graph
    (Set.delete u <$> Map.delete u ies)
    (Set.delete u <$> Map.delete u oes)

-- | Add edges from a vertex to a list of vertices in a graph
connect :: (Ord a) => a -> [a] -> Graph a -> Graph a
connect u vs (Graph ies oes) =
  uncurry Graph $
    foldr
      ( \v (ies', oes') ->
          ( Map.insertWith (<>) v (Set.singleton u) ies'
          , Map.insertWith (<>) v Set.empty oes'
          )
      )
      ( Map.insertWith (<>) u Set.empty ies
      , Map.insertWith (<>) u (Set.fromList vs) oes
      )
      vs

-- | Remove edges from a vertex to a list of vertices in a graph
disconnect :: (Ord a) => a -> [a] -> Graph a -> Graph a
disconnect u vs (Graph ies oes) =
  Graph
    ( foldr
        (\v -> Map.insertWith (flip Set.difference) v $ Set.singleton u)
        ies
        vs
    )
    (Map.insertWith (flip Set.difference) u (Set.fromList vs) oes)

-- | Reverse the edges of a graph
transpose :: Graph a -> Graph a
transpose (Graph ies oes) = Graph oes ies

{-
 - Query
 -}

-- | Get the vertex associations of a graph
assocs :: Graph a -> [(a, [a])]
assocs (Graph _ oes) = second Set.toAscList <$> Map.toAscList oes

-- | Get the edges of a graph
edges :: Graph a -> [(a, a)]
edges (Graph _ oes) =
  Map.foldrWithKey
    (\u -> flip $ foldr (\v -> (:) (u, v)))
    []
    oes

-- | Get the vertices of a graph
vertices :: Graph a -> [a]
vertices (Graph _ oes) = fst <$> Map.toAscList oes

-- | Get the number of incoming edges to a vertex in a graph
indegree :: (Ord a) => a -> Graph a -> Maybe Int
indegree v (Graph ies _) = Set.size <$> Map.lookup v ies

-- | Get the number of incoming edges of each vertex in a graph
indegrees :: Graph a -> [(a, Int)]
indegrees (Graph ies _) = Map.toAscList $ Set.size <$> ies

-- | Get the number of outgoing edges from a vertex in a graph
outdegree :: (Ord a) => a -> Graph a -> Maybe Int
outdegree u (Graph _ oes) = Set.size <$> Map.lookup u oes

-- | Get the number of outgoing edges of each vertex in a graph
outdegrees :: Graph a -> [(a, Int)]
outdegrees (Graph _ oes) = Map.toAscList $ Set.size <$> oes

-- | Get the associations of a vertex from a graph
lookup :: (Ord a) => a -> Graph a -> Maybe [a]
lookup u (Graph _ oes) = Set.toAscList <$> Map.lookup u oes

-- | Get the associations of the least vertex greater than or equal to a vertex
lookupGE :: (Ord a) => a -> Graph a -> Maybe (a, [a])
lookupGE a (Graph _ oes) = second Set.toAscList <$> Map.lookupGE a oes

-- | Get the associations of the least vertex strictly greater than a vertex
lookupGT :: (Ord a) => a -> Graph a -> Maybe (a, [a])
lookupGT a (Graph _ oes) = second Set.toAscList <$> Map.lookupGT a oes

-- | Get the associations of the greatest vertex less than or equal to a vertex
lookupLE :: (Ord a) => a -> Graph a -> Maybe (a, [a])
lookupLE a (Graph _ oes) = second Set.toAscList <$> Map.lookupLE a oes

-- | Get the associations of the greatest vertex strictly less than a vertex
lookupLT :: (Ord a) => a -> Graph a -> Maybe (a, [a])
lookupLT a (Graph _ oes) = second Set.toAscList <$> Map.lookupLT a oes

-- | Get the associations of the maximum vertex from a graph
lookupMax :: Graph a -> Maybe (a, [a])
lookupMax (Graph _ oes) = second Set.toAscList <$> Map.lookupMax oes

-- | Get the associations of the minimum vertex from a graph
lookupMin :: Graph a -> Maybe (a, [a])
lookupMin (Graph _ oes) = second Set.toAscList <$> Map.lookupMin oes

-- | Check whether a vertex is in a graph
member :: (Ord a) => a -> Graph a -> Bool
member u (Graph _ oes) = u `Map.member` oes

-- | Get the maximum vertex with no incoming edges from a graph
sourceMax :: Graph a -> Maybe a
sourceMax (Graph ies _) =
  Map.foldlWithKey
    ( \b k ->
        bool
          b
          (Just k)
          . Set.null
    )
    Nothing
    ies

-- | Get the minimum vertex with no incoming edges from a graph
sourceMin :: Graph a -> Maybe a
sourceMin (Graph ies _) =
  Map.foldrWithKey
    ( \k a b ->
        bool
          b
          (Just k)
          $ Set.null a
    )
    Nothing
    ies

-- | Get the vertices with no incoming edges from a graph
sources :: Graph a -> [a]
sources (Graph ies _) =
  Map.foldrWithKey
    ( \k a b ->
        bool
          b
          (k : b)
          $ Set.null a
    )
    []
    ies

-- | Get the maximum vertex with no outgoing edges from a graph
sinkMax :: Graph a -> Maybe a
sinkMax (Graph _ oes) =
  Map.foldlWithKey
    ( \b k ->
        bool
          b
          (Just k)
          . Set.null
    )
    Nothing
    oes

-- | Get the minimum vertex with no outgoing edges from a graph
sinkMin :: Graph a -> Maybe a
sinkMin (Graph _ oes) =
  Map.foldrWithKey
    ( \k a b ->
        bool
          b
          (Just k)
          $ Set.null a
    )
    Nothing
    oes

-- | Get the vertices with no outgoing edges from a graph
sinks :: Graph a -> [a]
sinks (Graph _ oes) =
  Map.foldrWithKey
    ( \k a b ->
        bool
          b
          (k : b)
          $ Set.null a
    )
    []
    oes

{-
 - Helpers
 -}

-- | Breadth-first search for the hierarchy in a graph from a starting vertex
bfs :: (Ord a) => Graph a -> a -> [Set a]
bfs (Graph _ oes) s =
  foldMap
    ( \vs ->
        Set.singleton s
          : unfoldr
            ( \((us, es), ds) ->
                bool
                  ( Just
                      ( us
                      ,
                        ( foldr
                            ( \u b@(us', es') ->
                                maybe
                                  b
                                  ( \vs' ->
                                      ( (us' <> vs') `Set.difference` ds
                                      , Map.delete u es'
                                      )
                                  )
                                  $ Map.lookup u es
                            )
                            (Set.empty, es)
                            us
                        , ds <> us
                        )
                      )
                  )
                  Nothing
                  $ Set.null us
            )
            ((vs, oes), Set.singleton s)
    )
    $ Map.lookup s oes