Haskell Quiz/Amazing Mazes/Solution Burton

module Main
    where
 
import Random
import System
import System.IO
import System.IO.Unsafe
import Maybe
import List
 
{--
Algorithm is Dijkstra depth-first search 
The solver is augmented by picking neighbour which is closest to the goal
rather than a random one
 
*Main> doSol 10 10 Nothing
+---+---+---+---+---+---+---+---+---+---+
| x   x | x   x   x |       |           |
+---+   +   +---+   +   +---+   +   +   +
| x   x | x |   | x |           |   |   |
+   +---+   +   +   +---+---+---+   +---+
| x   x   x |   | x   x |       |       |
+---+---+---+   +---+   +   +   +---+   +
|               | x   x |   |       |   |
+   +   +---+   +   +---+   +---+   +   +
|   |   |       | x |       |   |   |   |
+---+   +   +---+   +---+   +   +   +   +
|       |   | x   x |       |   |   |   |
+   +---+---+   +---+   +---+   +   +   +
|   |       | x   x |       |   |   |   |
+   +   +   +---+   +---+   +   +   +   +
|       |   | x   x |       |           |
+---+---+   +   +---+   +---+---+---+---+
|           | x   x |                   |
+   +---+---+---+   +---+---+---+---+   +
|                 x   x   x   x   x   x |
+---+---+---+---+---+---+---+---+---+---+
 
--}
 
data Direction = N | E | S | W
               deriving (Show, Eq, Enum)
data Cell = GenCell [Direction] [Direction] -- Cell borders walls
          deriving Show
type Maze = [Cell]
type Neighbour = (Maybe Cell, Int, Direction) --(cell, position, direction in which it lies)
data CLIFlag = Help | Version | Mode | Input String | Output String | Start Int | End Int
               deriving Show
type Solution = (Maze, [Int], [Int]) --(m, path, backtracks)
 
borders, walls :: Cell -> [Direction]
borders (GenCell bs ws) = bs
walls (GenCell bs ws) = ws
smaze :: Solution -> Maze
smaze (m, s, b) = m 
spath :: Solution -> [Int]
spath (m, s, b) = s
sback :: Solution -> [Int]
sback (m, s, b) = b  
ncell :: Neighbour -> Maybe Cell
ncell (c, p, d) = c                 
npos :: Neighbour -> Int
npos (c, p, d) = p
ndir :: Neighbour -> Direction
ndir (c, p, d) = d 
 
--get a maze with all walls intact
genBlockedMaze :: Int -> Int -> Maze
genBlockedMaze w h = topRow ++ (genMaze' (h-2) []) ++ bottomRow
                     where genMaze' 0 m = m
                           genMaze' (h+1) m   = midRow ++ genMaze' h m
                           genCells 0 bs ws   = []
                           genCells x bs ws   = (GenCell bs ws) : (genCells (x-1) bs ws)
                           topRow      = (GenCell [W,N] [E,S]) : (genCells (w-2) [N] [E .. W]) ++ [(GenCell [N, E] [S,W])]
                           bottomRow   = (GenCell [W,S] [N, E]) : (genCells (w-2) [S] [N, E, W]) ++ [(GenCell [S, E] [N, W])]
                           midRow      = (GenCell [W] [N .. S]) : (genCells (w-2) [] [N .. W]) ++ [(GenCell [E] [N, S, W])]
 
 
getRandNum :: Int -> Int
getRandNum n = unsafePerformIO $ getStdRandom $ randomR (0,n)
 
getDir :: Int -> Direction
getDir n = if n < 4 
           then [N .. W]!!n 
           else error ("can't give a direction for " ++ (show n)) 
 
randElem :: [a] -> a
randElem xs = xs!!(getRandNum ((length xs)-1))
 
--get the dimensions of a maze
dims :: Maze -> (Int, Int)
dims m = (w, (h+1))
         where w = findW m 1
               h = findH m 1
               findW [] x     = -1
               findW (c:cs) x = if E `elem` (borders c) then x else findW cs (x+1)
               findH [] x     = -1
               findH (c:cs) x = if S `elem` (borders c) then x `div` w else findH cs (x+1)
 
--get the neighbours of a position in the maze 
neighbs :: Maze -> Int -> [Neighbour]
neighbs m i = filter (isJust . ncell) [neighb (-1) 0 W, neighb 1 0 E, neighb 0 (-1) N, neighb 0 1 S]
              where len = length m
                    (w,h) = dims m
                    neighb x y d | newPos >= len || newPos < 0 = (Nothing, -1, d) --past beginning or end
                                 | newX >= w || newX < 0       = (Nothing, -1, d) --last row
                                 | newY >= h || newY < 0       = (Nothing, -1, d) --last col
                                 | otherwise                   = (Just (m!!newPos), newPos, d)
                                 where newPos = i+x + (y*w)
                                       newX   = x + (i `mod` h)
                                       newY   = (i `div` w) + y
 
--generate a perfect maze given the dimensions
genMaze :: Int -> Int -> Maze
genMaze w h = genMaze' m [0] 0 1 
              where m = genBlockedMaze w h
                    t = length m
                    genMaze' m' stk cur vis | vis == t = m'
                                            | null newNeighbs = genMaze' m' (tail stk) (head $ tail stk) vis
                                            | otherwise = genMaze' m'' (neighbPos:stk) neighbPos (vis+1)
                                            where newNeighbs = filter (\x -> (npos x `notElem` stk) && allwalls (fromJust $ ncell x)) (neighbs m' cur)
                                                  neighb = randElem newNeighbs
                                                  neighbWall = opp curWall
                                                  neighbPos = npos neighb
                                                  Just neighbCell = ncell neighb
                                                  curWall = ndir neighb
                                                  m'' = splice (splice m' (newNeighb, neighbPos)) (newCur, cur)
                                                  newCur = GenCell (borders curcell) (delete curWall (walls curcell))
                                                  newNeighb = GenCell (borders neighbCell) (delete neighbWall $ walls neighbCell)
                                                  curcell = m'!!cur
 
 
opp :: Direction -> Direction
opp N = S
opp E = W
opp S = N
opp W = E
 
allwalls :: Cell -> Bool
allwalls c = all (\d -> d `elem` walls c || d `elem` borders c) ds
           where ds = [N .. W]
 
splice :: Maze -> (Cell, Int) -> Maze
splice [] _ = []
splice m (c, 0) = c : (tail m)
splice m (c, i) = if i >= length m then error "cannot splice " else take i m ++ [c] ++ drop (i+1) m  
 
topBorder :: Int -> String
topBorder i = (topBorder' i "") ++ "+\n"
              where topBorder' 0 acc = acc
                    topBorder' (n+1) acc = topBorder' n ("+---" ++ acc)
 
asciiMaze :: Solution -> String
asciiMaze (m, s, b) = topBorder w ++ asciiRows m 0
    where (w,h) = dims m
          asciiRows [] c = "\n"
          asciiRows m c = asciiRow (take w m) [] 0 c s b ++ asciiRows (drop w m) (c+w)
 
asciiRow :: [Cell] -> [Cell] -> Int -> Int -> [Int] -> [Int] -> String
asciiRow [] acc i pos sol bk = ""
asciiRow [c] acc 0 pos sol bk = showCell c i pos sol bk ++ "|\n" ++ asciiRow (reverse (c:acc)) [] 1 (pos-length (c:acc)) sol bk
asciiRow [c] acc i pos sol bk = showCell c i pos sol bk ++ "+\n"
asciiRow (c:cs) acc 0 pos sol bk = showCell c i pos sol bk ++ asciiRow cs (c:acc) 0 (pos+1) sol bk
asciiRow (c:cs) acc i pos sol bk = showCell c i pos sol bk ++ asciiRow cs (c:acc) 1 (pos+1) sol bk
 
showCell :: Cell -> Int -> Int -> [Int] -> [Int] -> String
showCell c 0 pos sol bk | pos `elem` sol = showCellMS c
                        | pos `elem` bk  = showCellMB c
                        | otherwise      = showCellM c
showCell c _ _ _ _ = showCellB c
 
showCellB, showCellM, showCellMS, showCellMB :: Cell -> String
showCellB c | S `elem` borders c || S `elem` walls c = "+---"
            | otherwise                              = "+   "
showCellM c | W `elem` borders c || W `elem` walls c = "|   "
            | otherwise                              = "    "
showCellMS c | W `elem` borders c || W `elem` walls c = "| x "
             | otherwise                              = "  x "
showCellMB c | W `elem` borders c || W `elem` walls c = "| - "
             | otherwise                              = "  - "
 
printMaze :: Maze -> IO ()
printMaze m = putStr $ asciiMaze (m ,[], [])
 
solveMaze :: Maze -> Int -> Int -> Solution
solveMaze m start end = solveMaze' [start] [] start
                        where solveMaze' sol bk cur | cur == end = (m, sol, bk)
                                                    | null newNeighbs = solveMaze' (tail sol) ((head sol):bk) $ head $ tail sol
                                                    | otherwise = solveMaze' (neighbPos:sol) bk neighbPos
                                                    where newNeighbs = (filter (\x -> npos x `notElem` sol && npos x `notElem` bk && nowalls curcell x) $ neighbs m cur
                                                          curcell    = m!!cur
                                                          --neighb  = randElem newNeighbs
                                                          neighb = mhat newNeighbs end $ dims m
                                                          neighbPos  = npos neighb
 
coords :: Int -> (Int, Int) -> (Int, Int)
coords i (w, h) = (i `mod` w, i `div` h)
 
xpos, ypos :: (Int, Int) -> Int
xpos = fst
ypos = snd
 
--get neighbour which is closest to the goal
mhat :: [Neighbour] -> Int -> (Int, Int) -> Neighbour
mhat ns g (w, h) = head (qsort ns)
    where mhatdist n = abs (xpos (coords (npos n) (w, h)) - xposg) + abs (ypos (coords (npos n) (w, h)) - yposg)
          xposg = xpos (coords g (w, h))
          yposg = ypos (coords g (w, h))
          qsort [] = []
          qsort (n:ns) = qsort (filter closer ns) ++ [n] ++ qsort (filter further ns)
              where closer n' = mhatdist n' <= mhatdist n
                    further   = not . closer
 
nowalls :: Cell -> Neighbour -> Bool
nowalls c n = case ncell n of
               Nothing -> False
               Just c' ->
              dir `notElem` walls c
              && dir `notElem` borders c
              && opp dir `notElem` borders c'
              && opp dir `notElem` walls c'
            where dir = ndir n
 
printSolution :: Solution -> IO ()
printSolution = putStr . asciiMaze
 
proc :: [String] -> IO ()
proc args | length args /= 3 && length args /= 4 = usage
          | o == "--gen"                         = doGen w h f
          | o == "--sol"                         = doSol w h f
          | otherwise                            = usage
          where w = read $ args!!0
                h = read $ args!!1
                o = args!!2
                f = if length args == 4 then Just $ args!!3 else Nothing
 
doGen :: Int -> Int -> Maybe String -> IO ()
doGen w h Nothing  = printMaze $ genMaze w h
doGen w h (Just s) = writeFile s $ asciiMaze ((genMaze w h), [], [])
 
doSol :: Int -> Int -> Maybe String -> IO ()
doSol w h Nothing  = printSolution $ solveMaze (genMaze w h) 0 ((w*h)-1)
doSol w h (Just s) = writeFile s $ asciiMaze m
                     where m = solveMaze (genMaze w h) 0 ((w*h)-1)
 
usage :: IO ()
usage = do putStrLn "USAGE:"
           putStrLn "maze w h [--gen|--sol] [OUTFILE] [--start=START] [--end=END]"
 
main :: IO ()
main = do x <- getArgs
          proc x