{-# LANGUAGE GADTs #-}

module HGraph.Undirected.AdjacencyMap
       ( Graph
       , emptyGraph
       , module HGraph.Undirected
       )
where

import HGraph.Undirected
import qualified Data.Map as M
import qualified Data.Set as S

data Graph a where
  Graph :: Ord a => M.Map a (S.Set a) -> Int -> Graph a

emptyGraph :: Ord a => Graph a
emptyGraph :: Graph a
emptyGraph = Map a (Set a) -> Int -> Graph a
forall a. Ord a => Map a (Set a) -> Int -> Graph a
Graph Map a (Set a)
forall k a. Map k a
M.empty Int
0

instance UndirectedGraph Graph where
  empty :: Graph a -> Graph a
empty (Graph Map a (Set a)
_ Int
_) = Map a (Set a) -> Int -> Graph a
forall a. Ord a => Map a (Set a) -> Int -> Graph a
Graph Map a (Set a)
forall k a. Map k a
M.empty Int
0
  vertices :: Graph a -> [a]
vertices (Graph Map a (Set a)
adj Int
_) = Map a (Set a) -> [a]
forall k a. Map k a -> [k]
M.keys Map a (Set a)
adj
  numVertices :: Graph a -> b
numVertices (Graph Map a (Set a)
adj Int
_) = Int -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> b) -> Int -> b
forall a b. (a -> b) -> a -> b
$ Map a (Set a) -> Int
forall k a. Map k a -> Int
M.size Map a (Set a)
adj
  edges :: Graph a -> [(a, a)]
edges (Graph Map a (Set a)
adj Int
_) = [(a
v,a
u) | (a
v, Set a
nv) <- Map a (Set a) -> [(a, Set a)]
forall k a. Map k a -> [(k, a)]
M.assocs Map a (Set a)
adj, a
u <- Set a -> [a]
forall a. Set a -> [a]
S.toList Set a
nv, a
u a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
v]
  numEdges :: Graph a -> b
numEdges (Graph Map a (Set a)
_ Int
numE) = Int -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> b) -> Int -> b
forall a b. (a -> b) -> a -> b
$ Int
numE
  linearizeVertices :: Graph a -> (Graph Int, [(Int, a)])
linearizeVertices g :: Graph a
g@(Graph Map a (Set a)
adj Int
_) = (Graph Int
g', [(Int, a)]
assocs)
    where
      assocs :: [(Int, a)]
assocs = [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] (Map a (Set a) -> [a]
forall k a. Map k a -> [k]
M.keys Map a (Set a)
adj)
      ltoi :: Map a Int
ltoi = [(a, Int)] -> Map a Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(a, Int)] -> Map a Int) -> [(a, Int)] -> Map a Int
forall a b. (a -> b) -> a -> b
$ [a] -> [Int] -> [(a, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Map a (Set a) -> [a]
forall k a. Map k a -> [k]
M.keys Map a (Set a)
adj) [Int
0..]
      g' :: Graph Int
g' = ((Int, Int) -> Graph Int -> Graph Int)
-> Graph Int -> [(Int, Int)] -> Graph Int
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Graph Int -> (Int, Int) -> Graph Int)
-> (Int, Int) -> Graph Int -> Graph Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip Graph Int -> (Int, Int) -> Graph Int
forall (t :: * -> *) a. Mutable t => t a -> (a, a) -> t a
addEdge) ((Int -> Graph Int -> Graph Int) -> Graph Int -> [Int] -> Graph Int
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Graph Int -> Int -> Graph Int) -> Int -> Graph Int -> Graph Int
forall a b c. (a -> b -> c) -> b -> a -> c
flip Graph Int -> Int -> Graph Int
forall (t :: * -> *) a. Mutable t => t a -> a -> t a
addVertex) Graph Int
forall a. Ord a => Graph a
emptyGraph (((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst [(Int, a)]
assocs)) ([(Int, Int)] -> Graph Int) -> [(Int, Int)] -> Graph Int
forall a b. (a -> b) -> a -> b
$
            [ (Map a Int
ltoi Map a Int -> a -> Int
forall k a. Ord k => Map k a -> k -> a
M.! a
u, Map a Int
ltoi Map a Int -> a -> Int
forall k a. Ord k => Map k a -> k -> a
M.! a
v) | (a
u,a
v) <- Graph a -> [(a, a)]
forall (t :: * -> *) a. UndirectedGraph t => t a -> [(a, a)]
edges Graph a
g ]


instance Adjacency Graph where
  neighbors :: Graph a -> a -> [a]
neighbors (Graph Map a (Set a)
adj Int
_) a
v = Set a -> [a]
forall a. Set a -> [a]
S.toList (Set a -> [a]) -> Set a -> [a]
forall a b. (a -> b) -> a -> b
$ Map a (Set a)
adj Map a (Set a) -> a -> Set a
forall k a. Ord k => Map k a -> k -> a
M.! a
v
  degree :: Graph a -> a -> b
degree (Graph Map a (Set a)
adj Int
_) a
v = Int -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> b) -> Int -> b
forall a b. (a -> b) -> a -> b
$ Set a -> Int
forall a. Set a -> Int
S.size (Set a -> Int) -> Set a -> Int
forall a b. (a -> b) -> a -> b
$ Map a (Set a)
adj Map a (Set a) -> a -> Set a
forall k a. Ord k => Map k a -> k -> a
M.! a
v
  edgeExists :: Graph a -> (a, a) -> Bool
edgeExists (Graph Map a (Set a)
adj Int
_) (a
v,a
u) = a
u a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` (Map a (Set a)
adj Map a (Set a) -> a -> Set a
forall k a. Ord k => Map k a -> k -> a
M.! a
v)
  inducedSubgraph :: Graph a -> [a] -> Graph a
inducedSubgraph (Graph Map a (Set a)
adj Int
numE) [a]
vs = Map a (Set a) -> Int -> Graph a
forall a. Ord a => Map a (Set a) -> Int -> Graph a
Graph Map a (Set a)
adj' (Int -> Graph a) -> Int -> Graph a
forall a b. (a -> b) -> a -> b
$ ((Int -> Set a -> Int) -> Int -> Map a (Set a) -> Int
forall a b k. (a -> b -> a) -> a -> Map k b -> a
M.foldl' (\Int
s Set a
n -> Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Set a -> Int
forall a. Set a -> Int
S.size Set a
n) Int
0 Map a (Set a)
adj') Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
    where
      adj' :: Map a (Set a)
adj' = (Set a -> Set a) -> Map a (Set a) -> Map a (Set a)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\Set a
n -> Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.intersection Set a
n Set a
svs) (Map a (Set a) -> Map a (Set a)) -> Map a (Set a) -> Map a (Set a)
forall a b. (a -> b) -> a -> b
$ Map a (Set a) -> Set a -> Map a (Set a)
forall k a. Ord k => Map k a -> Set k -> Map k a
M.restrictKeys Map a (Set a)
adj Set a
svs
      svs :: Set a
svs = [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList [a]
vs

instance Mutable Graph where
  addVertex :: Graph a -> a -> Graph a
addVertex (Graph Map a (Set a)
adj Int
nE) a
v = Map a (Set a) -> Int -> Graph a
forall a. Ord a => Map a (Set a) -> Int -> Graph a
Graph (a -> Set a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
v Set a
forall a. Set a
S.empty Map a (Set a)
adj) Int
nE
  removeVertex :: Graph a -> a -> Graph a
removeVertex g :: Graph a
g@(Graph Map a (Set a)
adj Int
nE) a
v = 
    Map a (Set a) -> Int -> Graph a
forall a. Ord a => Map a (Set a) -> Int -> Graph a
Graph (a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
v (Map a (Set a) -> Map a (Set a)) -> Map a (Set a) -> Map a (Set a)
forall a b. (a -> b) -> a -> b
$ (a -> Map a (Set a) -> Map a (Set a))
-> Map a (Set a) -> [a] -> Map a (Set a)
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Set a -> Set a) -> a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.delete a
v)) Map a (Set a)
adj [a]
nv) (Int
nE Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Graph a -> a -> Int
forall (t :: * -> *) b a.
(Adjacency t, Integral b) =>
t a -> a -> b
degree Graph a
g a
v))
    where
      nv :: [a]
nv = Graph a -> a -> [a]
forall (t :: * -> *) a. Adjacency t => t a -> a -> [a]
neighbors Graph a
g a
v
  addEdge :: Graph a -> (a, a) -> Graph a
addEdge g :: Graph a
g@(Graph Map a (Set a)
adj Int
nE) (a
v,a
u)
    | Graph a -> (a, a) -> Bool
forall (t :: * -> *) a. Adjacency t => t a -> (a, a) -> Bool
edgeExists Graph a
g (a
v,a
u) = Graph a
g
    | Bool
otherwise = Map a (Set a) -> Int -> Graph a
forall a. Ord a => Map a (Set a) -> Int -> Graph a
Graph Map a (Set a)
adj' (Int
nE Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      where
        adj' :: Map a (Set a)
adj' = (Set a -> Set a -> Set a)
-> a -> Set a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.union a
v (a -> Set a
forall a. a -> Set a
S.singleton a
u) (Map a (Set a) -> Map a (Set a)) -> Map a (Set a) -> Map a (Set a)
forall a b. (a -> b) -> a -> b
$ (Set a -> Set a -> Set a)
-> a -> Set a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.union a
u (a -> Set a
forall a. a -> Set a
S.singleton a
v) Map a (Set a)
adj
  removeEdge :: Graph a -> (a, a) -> Graph a
removeEdge g :: Graph a
g@(Graph Map a (Set a)
adj Int
nE) (a
v,a
u)
    | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Graph a -> (a, a) -> Bool
forall (t :: * -> *) a. Adjacency t => t a -> (a, a) -> Bool
edgeExists Graph a
g (a
v,a
u) = Graph a
g
    | Bool
otherwise = Map a (Set a) -> Int -> Graph a
forall a. Ord a => Map a (Set a) -> Int -> Graph a
Graph Map a (Set a)
adj' (Int
nE Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
      where
        adj' :: Map a (Set a)
adj' = (Set a -> Set a) -> a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.delete a
u) a
v (Map a (Set a) -> Map a (Set a)) -> Map a (Set a) -> Map a (Set a)
forall a b. (a -> b) -> a -> b
$ (Set a -> Set a) -> a -> Map a (Set a) -> Map a (Set a)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.delete a
v) a
u Map a (Set a)
adj