-- | Random maze generation.
module MazeGen where

import Control.Monad.Random
import qualified Data.Set as S
import System.Random.Shuffle
import qualified Data.Foldable as F

-- | A cell is currently represented by 2D coordinates (a pair of
--   @Int@ values) but could easily be abstracted over to allow for
--   more interesting topologies.
data Cell = Cell Int Int  deriving (Show, Eq, Ord)

-- | A wall is between a pair of cells.
data Wall = Wall Cell Cell  deriving (Show, Eq, Ord)

-- | A maze has a height, width, a set of cells, and a set of walls.
--   We maintain the invariant that the walls are only between
--   cells that are also present.
data Maze = Maze Int Int
                 (S.Set Cell)
                 (S.Set Wall)
  deriving (Show)

-- | Generate a pretty textual representation of a maze.
showMaze :: Maze -> String
showMaze m@(Maze w h _ _) = unlines $ replicate (2*w+1) '_' : map (showMazeLine m) [0..h-1]

showMazeLine :: Maze -> Int -> String
showMazeLine (Maze w h cells walls) l = "|" ++ concatMap showCell [0..w-1]
  where
    showCell x = (if (Wall (Cell x l) (Cell x (l+1)) `S.member` walls)
                     || l == h - 1
                    then "_" else " ")
                 ++
                 (if (Wall (Cell x l) (Cell (x+1) l) `S.member` walls)
                     || x == w - 1
                     then "|" else "_")

-- | Generate an initial maze of the given width and height with walls
--   between every pair of adjacent cells.
initialMaze :: Int -> Int -> Maze
initialMaze width height = Maze width height cells (fullWalls cells)
  where cells = fullCells width height

-- | Generate the set of cells for a maze with the given width and height.
fullCells :: Int -> Int -> S.Set Cell
fullCells width height = S.fromList [Cell x y | x <- [0..width-1], y <- [0..height-1]]

-- | Generate the set of walls for the given set of cells.
fullWalls :: S.Set Cell -> S.Set Wall
fullWalls cells = S.fromList [Wall c1 c2 | c1 <- S.toList cells, c2 <- S.toList cells, nearby c1 c2]

-- | Determine whether two cells are adjacent.
nearby :: Cell -> Cell -> Bool
nearby (Cell x1 y1) (Cell x2 y2) = (x1 == x2 && y1 + 1 == y2)
                                || (x1 + 1 == x2 && y1 == y2)

-- | Remove the specified wall from a maze.
removeWall :: Wall -> Maze -> Maze
removeWall w (Maze width height cells walls)
  = Maze width height cells (S.delete w walls)

-- | Put the given set of walls into a random order.
randomizeWalls :: RandomGen g => S.Set Wall -> g -> [Wall]
randomizeWalls walls gen = shuffle' ws len gen
  where ws  = S.toList walls
        len = length ws

-- | Perform the random maze generation algorithm: given a list of
--   walls, look at the walls one by one in order through the list.
--   For each wall, if the cells on either side of it belong to
--   different connected components, remove the wall.  Otherwise,
--   leave it intact.  Makes use of 'removeWall'; use 'randomizeWalls'
--   to produce a list suitable for passing as an argument.
knockDownWalls :: Maze -> [Wall] -> Maze
knockDownWalls m@(Maze _ _ cells _) ws = knockDownWalls' m ws initialCellBlocks
  where initialCellBlocks = S.map S.singleton cells

knockDownWalls' m [] _ = m
knockDownWalls' m@(Maze wd ht cells walls) (w@(Wall c1 c2):ws) blocks
  | areConnected c1 c2 blocks = knockDownWalls' m ws blocks
  | otherwise                 = knockDownWalls' (removeWall w m) ws (connect c1 c2 blocks)

areConnected :: Cell -> Cell -> CellBlocks -> Bool
areConnected c1 c2 blocks = F.any (\b -> S.member c1 b && S.member c2 b) blocks

connect :: Cell -> Cell -> CellBlocks -> CellBlocks
connect c1 c2 blocks = S.insert (F.fold s1) s2
  where (s1, s2) = S.partition (\b -> S.member c1 b || S.member c2 b) blocks

-- | The 'CellBlocks' type keeps track of connected components of
--   cells during the generation routine.  We start with one component
--   per cell, and stop when we have just a single component.
type CellBlocks = S.Set (S.Set Cell)

randomMaze :: RandomGen g => g -> Int -> Int -> Maze
randomMaze gen w h = knockDownWalls initMaze walls'
  where initMaze@(Maze _ _ _ walls) = initialMaze w h
        walls' = randomizeWalls walls gen