{-# 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 M.empty 0

instance UndirectedGraph Graph where
  empty (Graph _ _) = Graph M.empty 0
  vertices (Graph adj _) = M.keys adj
  numVertices (Graph adj _) = fromIntegral $ M.size adj
  edges (Graph adj _) = [(v,u) | (v, nv) <- M.assocs adj, u <- S.toList nv, u >= v]
  numEdges (Graph _ numE) = fromIntegral $ numE
  linearizeVertices g@(Graph adj _) = (g', assocs)
    where
      assocs = zip [0..] (M.keys adj)
      ltoi = M.fromList $ zip (M.keys adj) [0..]
      g' = foldr (flip addEdge) (foldr (flip addVertex) emptyGraph (map fst assocs)) $
            [ (ltoi M.! u, ltoi M.! v) | (u,v) <- edges g ]


instance Adjacency Graph where
  neighbors (Graph adj _) v = S.toList $ adj M.! v
  degree (Graph adj _) v = fromIntegral $ S.size $ adj M.! v
  edgeExists (Graph adj _) (v,u) = u `S.member` (adj M.! v)
  inducedSubgraph (Graph adj numE) vs = Graph adj' $ (M.foldl' (\s n -> s + S.size n) 0 adj') `div` 2
    where
      adj' = M.map (\n -> S.intersection n svs) $ M.restrictKeys adj svs
      svs = S.fromList vs

instance Mutable Graph where
  addVertex (Graph adj nE) v = Graph (M.insert v S.empty adj) nE
  removeVertex g@(Graph adj nE) v = 
    Graph (M.delete v $ foldr (M.adjust (S.delete v)) adj nv) (nE - (degree g v))
    where
      nv = neighbors g v
  addEdge g@(Graph adj nE) (v,u)
    | edgeExists g (v,u) = g
    | otherwise = Graph adj' (nE + 1)
      where
        adj' = M.insertWith S.union v (S.singleton u) $ M.insertWith S.union u (S.singleton v) adj
  removeEdge g@(Graph adj nE) (v,u)
    | not $ edgeExists g (v,u) = g
    | otherwise = Graph adj' (nE - 1)
      where
        adj' = M.adjust (S.delete u) v $ M.adjust (S.delete v) u adj