module HGraph.Directed.AdjacencyMap
       ( Digraph
       , emptyDigraph
       , module HGraph.Directed
       )
where

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

type DirectedNeighborhood a = (S.Set a, S.Set a)
data Digraph a where
  Digraph :: Ord a => M.Map a (DirectedNeighborhood a) -> Digraph a

emptyDigraph :: Ord a => Digraph a
emptyDigraph = Digraph M.empty

instance DirectedGraph Digraph where
  empty (Digraph d) = Digraph M.empty
  numVertices (Digraph d) = fromIntegral $ M.size d
  vertices (Digraph d) = M.keys d
  arcs (Digraph d) = concatMap (\(v, (_,o)) -> [(v,u) | u <- S.toList o]) $ M.assocs d
  isVertex (Digraph d) v = v `M.member` d
  linearizeVertices g@(Digraph adj) = (g', assocs)
    where
      assocs = zip [0..] (M.keys adj)
      ltoi = M.fromList $ zip (M.keys adj) [0..]
      g' = foldr addArc (foldr addVertex emptyDigraph (map fst assocs)) $
            [ (ltoi M.! u, ltoi M.! v) | (u,v) <- arcs g ]

instance Adjacency Digraph where
  outneighbors (Digraph d) v = S.toList $ snd $ d M.! v
  inneighbors  (Digraph d) v = S.toList $ fst $ d M.! v
  arcExists (Digraph d) (v,u) = u `S.member` (snd $ d M.! v)

instance Mutable Digraph where
  addVertex v (Digraph d) =
    Digraph (M.insertWith (\_ o -> o) v (S.empty, S.empty) d)
  removeVertex v g@(Digraph d) = 
      let Digraph d' = foldr removeArc g
                     (  (map (\u -> (v,u)) $ outneighbors g v)
                     ++ (map (\u -> (u,v)) $ inneighbors  g v))
      in Digraph $ M.delete v d'
            
  addArc (v,u) (Digraph d) =
    Digraph ( M.adjust (\(i,o) -> (i, S.insert u o)) v
            $ M.adjust (\(i,o) -> (S.insert v i, o)) u d)
  removeArc (v,u) (Digraph d)  =
    Digraph ( M.adjust (\(i,o) -> (i, S.delete u o)) v
            $ M.adjust (\(i,o) -> (S.delete v i, o)) u d)