content: module Algorithm.SearchSpec (
    main,
    spec
    ) where
  
  import Test.Hspec
  import Algorithm.Search
  import qualified Data.Map as Map
  import Data.Maybe (fromJust)
  
  main :: IO ()
  main = hspec spec
  
  -- | Example cyclic directed unweighted graph
  cyclicUnweightedGraph :: Map.Map Int [Int]
  cyclicUnweightedGraph = Map.fromList [
    (0, [1, 2, 3]),
    (1, [4, 6]),
    (2, [0, 1, 6, 8]),
    (3, [1, 2]),
    (4, [0]),
    (5, [4]),
    (6, [4]),
    (8, [0, 5])
    ]
  
  -- | Example acyclic directed unweighted graph
  acyclicUnweightedGraph :: Map.Map Int [Int]
  acyclicUnweightedGraph = Map.fromList [
    (0, [1, 2, 3]),
    (1, [4]),
    (2, [5]),
    (3, [2]),
    (4, []),
    (5, [])
    ]
  
  -- | Example cyclic directed weighted graph
  cyclicWeightedGraph :: Map.Map Char [(Char, Int)]
  cyclicWeightedGraph = Map.fromList [
    ('a', [('b', 1), ('c', 2)]),
    ('b', [('a', 1), ('c', 2), ('d', 5)]),
    ('c', [('a', 1), ('d', 2)]),
    ('d', [])
    ]
  
  -- | Example for taxicab path finding
  taxicabNeighbors :: (Int, Int) -> [(Int, Int)]
  -- the ordering here is important--for dfs, last state will be visited first
  taxicabNeighbors (x, y) = [(x, y + 1), (x - 1, y), (x, y - 1), (x + 1, y)]
  
  isWall :: (Int, Int) -> Bool
  isWall (x, y) = x == 1 && ((-2) <= y && y <= 1)
  
  taxicabDistance :: (Int, Int) -> (Int, Int) -> Int
  taxicabDistance (x1, y1) (x2, y2) = abs (x2 - x1) + abs (y2 - y1)
  
  taxicabNeighborsBounded :: (Int, Int) -> Maybe [(Int, Int)]
  taxicabNeighborsBounded (x, y)
    | outOfBounds (x, y) = Nothing
    | otherwise = Just $ taxicabNeighbors (x, y)
  
  taxicabDistanceBounded :: (Int, Int) -> (Int, Int) -> Maybe Int
  taxicabDistanceBounded (x1, y1) (x2, y2)
    | outOfBounds (x1, y1) || outOfBounds (x2, y2) = Nothing
    | otherwise = Just $ taxicabDistance (x1, y1) (x2, y2)
  
  outOfBounds :: (Int, Int) -> Bool
  outOfBounds (x, y) = abs x + abs y > 10
  
  isBigWall :: (Int, Int) -> Bool
  isBigWall (x, y) = x == 1 && ((-10) <= y && y <= 10)
  
  spec :: Spec
  spec = do
    describe "bfs" $ do
      let next = (cyclicUnweightedGraph Map.!)
      it "performs breadth-first search" $
        bfs next (== 4) 0 `shouldBe` Just [1, 4]
      it "handles pruning" $
        bfs (next `pruning` odd) (== 4) 0 `shouldBe` Just [2, 6, 4]
      it "returns Nothing when no path is possible" $
        bfs (next `pruning` odd `pruning` (== 6)) (== 4) 0 `shouldBe` Nothing
    describe "dfs" $ do
      let next = (cyclicUnweightedGraph Map.!)
      it "performs depth-first search" $
        dfs next (== 4) 0 `shouldBe` Just [3, 2, 8, 5, 4]
      it "handles pruning" $
        dfs (next `pruning` odd) (== 4) 0 `shouldBe` Just [2, 6, 4]
      it "returns Nothing when no path is possible" $
        dfs (next `pruning` odd `pruning` (== 6)) (== 4) 0 `shouldBe` Nothing
      it "handles doubly-inserted nodes" $
        dfs (acyclicUnweightedGraph Map.!) (==4) 0 `shouldBe` Just [1, 4]
    describe "dijkstra" $ do
      let next = map fst . (cyclicWeightedGraph Map.!)
          cost a b = fromJust . lookup b $ cyclicWeightedGraph Map.! a
      it "performs dijkstra's algorithm" $
        dijkstra next cost (== 'd') 'a'
          `shouldBe` Just (4, ['c', 'd'])
      it "handles pruning" $
        dijkstra (next `pruning` (== 'c')) cost (== 'd') 'a'
          `shouldBe` Just (6, ['b', 'd'])
      it "returns Nothing when no path is possible" $
        dijkstra (next `pruning` (== 'b') `pruning` (== 'c')) cost (== 'd') 'a'
          `shouldBe` Nothing
      it "handles zero-length solutions" $
        dijkstra next cost (== 'd') 'd' `shouldBe` Just (0, [])
    describe "dijkstra with cost" $ do
      let new_cost old_cost (char, step_cost) = (char, old_cost * 10 + step_cost)
          next (char, cost) = map (new_cost cost) (cyclicWeightedGraph Map.! char)
          isChar char1 (char2, _cost) = char1 == char2
      it "performs dijkstra's algorithm" $
        dijkstraAssocCost next (== 'd') 'a'
          `shouldBe` Just (15, ['b', 'd'])
      it "handles pruning" $
        dijkstraAssocCost (next `pruning` isChar 'b') (== 'd') 'a'
          `shouldBe` Just (22, ['c', 'd'])
      it "returns Nothing when no path is possible" $
        dijkstraAssocCost (next `pruning` isChar 'b' `pruning` isChar 'c') (== 'd') 'a'
          `shouldBe` Nothing
      it "handles zero-length solutions" $
        dijkstraAssocCost next (== 'd') 'd' `shouldBe` Just (0, [])
    describe "aStar" $ do
      let start = (0, 0)
          end = (2, 0)
      it "performs the A* algorithm" $
        aStar taxicabNeighbors taxicabDistance (taxicabDistance end) (== end)
          start
          `shouldBe` Just (2, [(1, 0), (2, 0)])
      it "handles pruning" $
        aStar (taxicabNeighbors `pruning` isWall) taxicabDistance
          (taxicabDistance end) (== end) start
          `shouldBe` Just (6, [(0, 1), (0, 2), (1, 2), (2, 2), (2, 1), (2, 0)])
      it "returns Nothing when no path is possible" $
        aStar
          (taxicabNeighbors
            `pruning` isWall
            `pruning` (\ p -> taxicabDistance p (0,0) > 1)
          )
          taxicabDistance
          (taxicabDistance end)
          (== end)
          start
          `shouldBe` Nothing
      it "handles zero-length solutions" $
        aStar taxicabNeighbors taxicabDistance (taxicabDistance end) (== start)
          start
          `shouldBe` Just (0, [])
    describe "bfsM" $ do
      let start = (0, 0)
          end = (2, 0)
      it "performs monadic breadth-first search" $ do
        bfsM taxicabNeighborsBounded (return . (== end)) start
          `shouldBe` Just (Just [(1, 0), (2, 0)])
        bfsM
          (taxicabNeighborsBounded `pruningM` (return . isWall))
          (return . (== end))
          start
          `shouldBe` Just (Just [(0,1),(0,2),(1,2),(2,2),(2,1),(2,0)])
      it "handles cyclic graphs" $ do
        let nextM = return . map fst . (cyclicWeightedGraph Map.!)
        bfsM nextM (return . (== 'd')) 'a'
          `shouldBe` Just (Just ['b', 'd'])
      it "correctly handles monadic behavior" $ do
        bfsM
          (taxicabNeighborsBounded `pruningM` (return . isBigWall))
          (return . (== end))
          start
          `shouldBe` Nothing
        bfsM taxicabNeighborsBounded (const Nothing) start
          `shouldBe` Nothing
    describe "dfsM" $ do
      let start = (0, 0)
          end = (2, 0)
      it "performs monadic depth-first search" $
        dfsM taxicabNeighborsBounded (return . (== end)) start
          `shouldBe` Just (Just [(1, 0), (2, 0)])
      it "handles doubly-inserted nodes" $ do
        let nextM = return . (acyclicUnweightedGraph Map.!)
        dfsM nextM (return . (== 4)) 0 `shouldBe` Just (Just [1, 4])
      it "correctly handles monadic behavior" $ do
        dfsM
          (taxicabNeighborsBounded `pruningM` (return . isBigWall))
          (return . (== end))
          start
          `shouldBe` Nothing
        dfsM taxicabNeighborsBounded (const Nothing) start
          `shouldBe` Nothing
    describe "dijkstraM" $ do
      let start = (0, 0)
          end = (2, 0)
      it "performs monadic dijkstra's algorithm" $
        dijkstraM
          taxicabNeighborsBounded
          taxicabDistanceBounded
          (return . (== end))
          start
          `shouldBe` Just (Just (2, [(1, 0), (2, 0)]))
      it "handles cyclic graphs" $ do
        let nextM = return . map fst . (cyclicWeightedGraph Map.!)
            costM a b = lookup b $ cyclicWeightedGraph Map.! a
        dijkstraM nextM costM (return . (== 'd')) 'a'
          `shouldBe` Just (Just (4, ['c', 'd']))
        dijkstraM (nextM `pruningM` (return . (== 'c'))) costM
          (return . (== 'd')) 'a'
          `shouldBe` Just (Just (6, ['b', 'd']))
      it "handles zero-length solutions" $ do
        let nextM = return . map fst . (cyclicWeightedGraph Map.!)
            costM a b = lookup b $ cyclicWeightedGraph Map.! a
        dijkstraM nextM costM (return . (== 'd')) 'd'
          `shouldBe` Just (Just (0, []))
      it "correctly handles monadic behavior" $ do
        dijkstraM
          (taxicabNeighborsBounded `pruningM` (return . isBigWall))
          taxicabDistanceBounded
          (return . (== end))
          start
          `shouldBe` Nothing
        dijkstraM
          taxicabNeighborsBounded
          ((const . const) Nothing :: (Int, Int) -> (Int, Int) -> Maybe Int)
          (return . (== end))
          start
          `shouldBe` Nothing
        dijkstraM
          (taxicabNeighborsBounded `pruningM` (return . isBigWall))
          taxicabDistanceBounded
          (const Nothing)
          start
          `shouldBe` Nothing
    describe "dijkstraAssocCostM" $ do
      let start = (0, 0)
          end = (2, 0)
          taxicabNeighborsBounded' (from, cost1) = do
            tos <- taxicabNeighborsBounded from
            cost2s <- mapM (taxicabDistanceBounded from) tos
            pure $ zip tos $ (cost1 +) <$> cost2s
      it "performs monadic dijkstra's algorithm" $
        dijkstraAssocCostM
          taxicabNeighborsBounded'
          (return . (== end))
          start
          `shouldBe` Just (Just (2, [(1, 0), (2, 0)]))
      it "handles cyclic graphs" $ do
        let nextM (char, cost) = return $ fmap (cost +) <$> (cyclicWeightedGraph Map.! char)
            isChar char1 (char2, _cost) = char1 == char2
        dijkstraAssocCostM nextM (return . (== 'd')) 'a'
          `shouldBe` Just (Just (4, ['c', 'd']))
        dijkstraAssocCostM (nextM `pruningM` (return . isChar 'c'))
          (return . (== 'd')) 'a'
          `shouldBe` Just (Just (6, ['b', 'd']))
      it "handles zero-length solutions" $ do
        let nextM (char, cost) = return $ fmap (cost +) <$> (cyclicWeightedGraph Map.! char)
        dijkstraAssocCostM nextM  (return . (== 'd')) 'd'
          `shouldBe` Just (Just (0, []))
      it "correctly handles monadic behavior" $ do
        dijkstraAssocCostM
          (taxicabNeighborsBounded' `pruningM` (return . isBigWall . fst))
          (return . (== end))
          start
          `shouldBe` Nothing
        dijkstraAssocCostM
          (taxicabNeighborsBounded' `pruningM` (return . isBigWall . fst))
          (const Nothing)
          start
          `shouldBe` Nothing
    describe "aStarM" $ do
      let start = (0, 0)
          end = (2, 0)
      it "performs a monadic A* algorithm" $
        aStarM
          (taxicabNeighborsBounded `pruningM` (return . isWall))
          taxicabDistanceBounded
          (taxicabDistanceBounded end)
          (return . (== end))
          start
          `shouldBe`
          Just (Just (6, [(0, 1), (0, 2), (1, 2), (2, 2), (2, 1), (2, 0)]))
      it "handles zero-length solutions" $
        aStarM taxicabNeighborsBounded taxicabDistanceBounded
          (taxicabDistanceBounded end) (return . (== start)) start
          `shouldBe` Just (Just (0, []))
      it "correctly handles monadic behavior" $ do
        aStarM
          (taxicabNeighborsBounded `pruningM` (return . isBigWall))
          taxicabDistanceBounded
          (taxicabDistanceBounded end)
          (return . (== end))
          start
          `shouldBe` Nothing
        aStarM
          taxicabNeighborsBounded
          ((const . const) Nothing :: (Int, Int) -> (Int, Int) -> Maybe Int)
          (taxicabDistanceBounded end)
          (return . (== end))
          start
          `shouldBe` Nothing
        aStarM
          taxicabNeighborsBounded
          taxicabDistanceBounded
          (const Nothing)
          (return . (== end))
          start
          `shouldBe` Nothing
        aStarM
          taxicabNeighborsBounded
          taxicabDistanceBounded
          (taxicabDistanceBounded end)
          (const Nothing)
          start
          `shouldBe` Nothing
    describe "incrementalCosts" $ do
      let cost a b = fromJust . lookup b $ cyclicWeightedGraph Map.! a
      it "gives the incremental costs along a path" $
        incrementalCosts cost ['a', 'b', 'd'] `shouldBe` [1, 5]
      it "handles zero-length paths" $ do
        incrementalCosts cost [] `shouldBe` []
        incrementalCosts cost ['a'] `shouldBe` []
    describe "incrementalCostsM" $ do
      let costM a b = lookup b $ cyclicWeightedGraph Map.! a
      it "gives monadic incremental costs along a path" $
        incrementalCostsM costM ['a', 'b', 'd'] `shouldBe` Just [1, 5]
      it "handles zero-length paths" $ do
        incrementalCostsM costM [] `shouldBe` Just []
        incrementalCostsM costM ['a'] `shouldBe` Just []
      it "correctly handles monadic behavior" $
        incrementalCostsM costM ['a', 'd'] `shouldBe` Nothing
    describe "pruning" $ do
      it "filters according to its predicate" $
        ((\x -> [x, 2*x, 3*x]) `pruning` even) 3 `shouldBe` [3, 9]
    describe "pruningAssoc" $ do
      it "filters according to its predicate" $
        ((\x -> [(x, 1), (2*x, 2)]) `pruningAssoc` (\(_, cost) -> cost > 1)) 3 `shouldBe` [(3, 1)]
    describe "pruningM" $ do
      it "filters according to its predicate" $
        ((\x -> (Just [x, 2*x, 3*x])) `pruningM` (\x -> Just (even x))) 3 `shouldBe` Just [3, 9]
    describe "pruningAssocM" $ do
       it "filters according to its predicate" $
         ((\x -> Just [(x, 1), (2*x, 2)]) `pruningAssocM` (\(_, cost) -> Just (cost > 1))) 3 `shouldBe` Just [(3, 1)]
  