Learning functional maze

maze.hs

{-

Funcitonal Maze
Run: runhaskell <SOURCE> <WIDTH::Int> <HEIGHT::Int> <RANDOM_SEED::Int>
Control:
    - [W] [A] [S] [D]: move up, left, down, right
    - [Q] [E|SPACE]: rotate counter-clockwise, clockwise
    - [ESC]: quit

-}

import Control.Monad.Trans.State
import System.Random
import Control.Applicative
import Control.Monad
import Data.List
import System.Environment
import System.IO
import Data.Maybe

data Dir = North | South | West | East deriving(Show, Eq, Ord)
data TurnDir = CW | CCW deriving(Show, Eq)

newtype Cell = Cell (Point, [Dir]) deriving(Show, Eq)

type Point = (Int, Int)
type Board = [Cell]

--          size     reserved   stack      final
generate :: Point -> [Point] -> Board -> State StdGen Board
generate _ _ [] = return []
generate size reserved (Cell (x, xdirs):xs) = do
    randGen <- get
    let (nextDirs, randGen') = runState (potentialNext size x reserved) randGen
    let nextPoints = pointOfDir x <$> nextDirs
    let nextCells = cellOfDir x <$> nextDirs
    let (rest, randGen'') = runState (generate size (reserved ++ nextPoints) (xs ++ nextCells)) randGen'
    put randGen''
    return (Cell(x, xdirs ++ nextDirs) : rest)

potentialNext :: Point -> Point -> [Point] -> State StdGen [Dir]
potentialNext size (x, y) avail = do
    randGen0 <- get
    let canPointBeAdded point = (point `notElem` avail) && validPont size point
    let (randR, randGen1) = randomR (1, 5) randGen0 :: (Int, StdGen)
    let (randT, randGen2) = randomR (1, 5) randGen1 :: (Int, StdGen)
    let (randB, randGen3) = randomR (1, 5) randGen2 :: (Int, StdGen)
    let (randL, randGen4) = randomR (1, 5) randGen3 :: (Int, StdGen)
    let r = [East  | randR <= 3 && canPointBeAdded (x + 1, y)]
    let t = [North | randT <= 3 && canPointBeAdded (x, y - 1)]
    let b = [South | randB <= 3 && canPointBeAdded (x, y + 1)]
    let l = [West  | randL <= 3 && canPointBeAdded (x - 1, y)]
    let (potentials, randGen5) = runState (shuffle $ concat [l, r, t, b]) randGen4
    put randGen5
    return potentials

validPont :: Point -> Point -> Bool
validPont (w, h) (x, y)
    | x < 0  = False
    | y < 0  = False
    | x >= w = False
    | y >= h = False
    | otherwise = True

shuffle :: [a] -> State StdGen [a]
shuffle [] = return []
shuffle (x:xs) = do
    randGen <- get
    let (randNum, randGen') = randomR (1, 2) randGen :: (Int, StdGen)
    let (rest, randGen'') = runState (shuffle xs) randGen'
    put randGen''
    return (if randNum == 1 then x : rest else rest ++ [x])

revDir :: Dir -> Dir
revDir North = South
revDir South = North
revDir East  = West
revDir West  = East

pointOfDir :: Point -> Dir -> Point
pointOfDir (x, y) West  = (x - 1, y)
pointOfDir (x, y) East  = (x + 1, y)
pointOfDir (x, y) North = (x, y - 1)
pointOfDir (x, y) South = (x, y + 1)

cellOfDir :: Point -> Dir -> Cell
cellOfDir (x, y) West  = Cell ((x - 1, y), [East])
cellOfDir (x, y) East  = Cell ((x + 1, y), [West])
cellOfDir (x, y) North = Cell ((x, y - 1), [South])
cellOfDir (x, y) South = Cell ((x, y + 1), [North])

dirsToChar :: [Dir] -> Char
dirsToChar [] = ' '
dirsToChar [North] = '╨'
dirsToChar [South] = '╥'
dirsToChar [West] = '╡'
dirsToChar [East] = '╞'
dirsToChar [North, South] = '║'
dirsToChar [North, East] = '╚'
dirsToChar [North, West] = '╝'
dirsToChar [West, East] = '═'
dirsToChar [South, East] = '╔'
dirsToChar [South, West] = '╗'
dirsToChar [North, West, East] = '╩'
dirsToChar [North, South, East] = '╠'
dirsToChar [North, South, West] = '╣'
dirsToChar [South, West, East] = '╦'
dirsToChar [North, South, West, East] = '╬'
dirsToChar _ = '?'

goto :: Point -> IO ()
goto (x, y) = putStr ("\ESC[" ++ show (y + 2) ++ ";" ++ show (x + 2) ++ "H")

--           original connected
printCell :: Board -> [Point] -> Cell -> IO ()
printCell b conns c@(Cell (p, dirs)) = do
    goto p
    if p `elem` conns then putStr "\ESC[93m" else putStr "\ESC[39m"
    putStr [dirsToChar $ sort dirs]

getIntFromArg :: Int -> IO Int
getIntFromArg pos = read . (!!pos) <$> getArgs

clearScreen :: IO ()
clearScreen = putStrLn "\ESC[2J"

applyN :: Int -> (a -> a) -> a -> a
applyN 0 _ a = a
applyN n f a = applyN (n - 1) f $ f a

turnCell :: TurnDir -> [Dir] -> [Dir]
turnCell CCW dirs = applyN 3 (turnCell CW) dirs
turnCell CW dirs = concat [n, s, w, e]
    where
        n = [East  | North `elem` dirs]
        s = [West  | South `elem` dirs]
        w = [North | West  `elem` dirs]
        e = [South | East  `elem` dirs]

turnBoard :: Point -> TurnDir -> Board -> Board
turnBoard _ _ [] = []
turnBoard p@(px, py) dir (c@(Cell ((cx, cy), dirs)):cs)
    | px == cx && py == cy = Cell ((px, py), turnCell dir dirs) : cs
    | otherwise = c : turnBoard p dir cs

getCell :: Point -> Board -> Maybe Cell
getCell _ [] = Nothing
getCell p@(px, py) (c@(Cell ((cx, cy), _)):cs)
    | px == cx && py == cy = Just c
    | otherwise = getCell p cs

hasConnection :: Cell -> Dir -> Bool
hasConnection (Cell (_, dirs)) dir = dir `elem` dirs

connectionFor :: Board -> Point -> Dir -> Maybe Cell
connectionFor board p dir = do
    newCell <- getCell (pointOfDir p dir) board
    if hasConnection newCell (revDir dir) then return newCell else Nothing

--       stack    visited    original
reach :: Board -> [Point] -> Board -> [Point]
reach [] _ _ = []
reach (c@(Cell (p, dirs)):cs) visited b = p : reach (cs ++ newReach) (p : visited) b
    where
        newReach = filter (\(Cell (p, _)) -> p `notElem` visited) $ map fromJust $ filter (/=Nothing) $ map (connectionFor b p) dirs

shuffleBoard :: Board -> State StdGen Board
shuffleBoard [] = return []
shuffleBoard (c:cs) = do
    randGen <- get
    let (spin, randGen') = randomR (0, 3) randGen :: (Int, StdGen)
    let (Cell (p, dirs)) = c
    let newDirs = applyN spin (turnCell CW) dirs
    let (rest, randGen'') = runState (shuffleBoard cs) randGen'
    put randGen''
    return (Cell (p, newDirs):rest)

play :: Point -> Point -> Board -> IO ()
play size@(w, h) (x, y) board = do
    let connected = reach [head board] [] board
    if length connected == length board then return () else playTurn size (x, y) board connected

playTurn :: Point -> Point -> Board -> [Point] -> IO ()
playTurn size@(w, h) (x, y) board connected = do
    clearScreen
    mapM_ (printCell board connected) board
    goto (x, y)
    keyStroke <- getChar
    case keyStroke of
        'w' -> play size (x, y - 1) board
        's' -> play size (x, y + 1) board
        'a' -> play size (x - 1, y) board
        'd' -> play size (x + 1, y) board
        'q' -> play size (x, y) (turnBoard (x, y) CCW board)
        'e' -> play size (x, y) (turnBoard (x, y) CW board)
        ' ' -> play size (x, y) (turnBoard (x, y) CW board)
        '\ESC' -> return ()
        _ -> play size (x, y) board

main = do
    hSetBuffering stdin NoBuffering
    width <- getIntFromArg 0
    height <- getIntFromArg 1
    randGen <- getIntFromArg 2
    let (board, randGen') = runState (generate (width, height) [(0, 0)] [Cell((0, 0), [North])]) (mkStdGen randGen)
    play (width, height) (0, 0) $ evalState (shuffleBoard board) randGen'
    goto (0, height)
    putStr "\n"