module Main where
  import AdjacencyList
  import PropertyMap
  import BFS
  import HUnit
  import Queue

  data TestVisitor1 = TestVis1
  instance BFSVisitor TestVisitor1 q [Int] AdjacencyList Edge Vertex where
    discover_vertex vis v g q a = ((idx v):a,q)

  n = 7::Int
  g = adj_list n [(0,1),(1,2),(1,3),(3,4),(0,4),(4,5),(3,6)]
  s = vertex 0

  res = breadth_first_search g s TestVis1 ([]::[Int])

  test1 = TestCase (assertEqual "[0,1,4,2,3,5,6] (reverse res)" [0,1,4,2,3,5,6] (reverse res))

  -- Get the tree edges and bfs numbers
  data TestVisitor2 = TestVis2

  instance (Graph g e v, Ord v) => BFSVisitor TestVisitor2 q (Int,(FMap v Int),[e]) g e v where
    tree_edge vis e g q (n, nmap, tree_edges) = ((n,nmap,(e : tree_edges)), q)
    discover_vertex vis v g q (n, nmap::(FMap v Int), te) = (((n+1), put nmap v n, te), q)

  (num, bfs_num, tree_edges) = breadth_first_search g s TestVis2 (0::Int, init_map (vertices g) (0::Int), ([]::[Edge]))

  -- Construct a graph that is the BFS tree
  bfs_tree = adj_list n [ (idx (src e g), idx (tgt e g)) | e <- tree_edges ]

  -- Convert the graph to a tree data structure labelled by BFS number
  data Tree a = Tr a [Tree a] deriving (Eq, Read, Show)
  bfs_num_tree bfs_num s t = 
    Tr (get bfs_num s) (subtrees (out_edges s t))
    where
      subtrees [] = []
      subtrees (e:es) = 
        (bfs_num_tree bfs_num (tgt e t) t) : (subtrees es)

  good_bfs_tree = Tr (0::Int) [Tr 2 [Tr 5 []],Tr 1 [Tr 4 [Tr 6 []],Tr 3 []]]
  res2 = bfs_num_tree bfs_num s bfs_tree

  test2 = TestCase (assertEqual "good_bfs_tree res2" good_bfs_tree res2)

  tests = TestList [TestLabel "test1" test1, TestLabel "test2" test2]
  main = do runTestTT tests