{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
import qualified Data.IntMap.Strict as IntMap
import qualified Data.IntSet as IntSet
import qualified Data.Sequence as Seq
import qualified Data.Foldable as F
import Data.IntMap.Strict(IntMap)
import Data.IntSet(IntSet)
import Data.Sequence(Seq)
import Control.Monad.State
import Data.Monoid
import Data.Maybe

type Node = Int
type LNode a = (Node, a)
type Edge = (Int, Int)
type LEdge a = (Edge, a)

newtype Graph a = Graph {
  unGraph :: IntMap (Seq (LNode a))
  }

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

empty = Graph IntMap.empty

addUEdge s t w = Graph . IntMap.insertWith (\old new -> old <> new) t (Seq.singleton (s, w)) . IntMap.insertWith (\old new -> old <> new) s (Seq.singleton (t, w)) . unGraph
addEdge s t w = Graph . IntMap.insertWith (\old new -> old <> new) s (Seq.singleton (t, w)) . unGraph

mkGraph :: [LEdge a] -> Graph a
mkGraph = foldl acc empty
  where acc g (!(!s, !t), !w) = addEdge s t w g

mkUGraph :: [LEdge a] -> Graph a
mkUGraph = foldl acc empty
  where acc g (!(!s, !t), !w) = addUEdge s t w g

data Policy = BFS | DFS

enqueue BFS s q = q <> s
enqueue DFS s q = s <> q

fsg :: (MonadState IntSet m, Monoid n, Num a) => (Node -> Node -> Int -> a -> n) -> Graph a -> n -> Policy -> Seq (Node, (Node, Int, a)) -> m n
fsg f g m p (Seq.viewl -> Seq.EmptyL) = return m
fsg f g m p (Seq.viewl -> ((!t, (!s, !l, !z))) Seq.:< q) = do
  visited <- get
  if t `IntSet.member` visited then fsg f g m p q else do
    put (IntSet.insert t visited)
    let adjs = maybe Seq.empty id (IntMap.lookup t (unGraph g))
        notVisited = fmap (\(x, y) -> (x, (t, l + 1, y + z))) . Seq.filter (\(x, y) -> x `IntSet.notMember` visited) $ adjs
    fsg f g (f t s l z <> m) p (enqueue p notVisited q)

start z = Seq.singleton (z, (z, 1, 0) )
nodes s p l w = Seq.singleton s

bfsWith f z g = evalState (fsg f g mempty BFS s) IntSet.empty
  where s = start z
dfsWith f z g = evalState (fsg f g mempty BFS s) IntSet.empty
  where s = start z

bfs z g = F.toList . evalState (fsg nodes g mempty BFS s) $ IntSet.empty
  where s = start z

dfs z g = F.toList . evalState (fsg nodes g mempty DFS s) $ IntSet.empty
  where s = start z