import Data.List
import Data.Maybe

data Cell = Clr | Obs deriving (Show, Eq)
type Grid = [[Cell]]
type Coo = (Int, Int)
data Dir = U | R | D | L deriving (Show, Eq)
type Guard = (Coo, Dir)

(!?) :: Grid -> Coo -> Maybe Cell
grid !? (row, col)
  | row < 0                     = Nothing
  | row >= length grid          = Nothing
  | col < 0                     = Nothing
  | col >= (length $ head grid) = Nothing
  | otherwise                   = Just (grid !! row !! col)

findGuard :: Int -> [[Char]] -> Guard
findGuard r (h:t)
  | c == Nothing = findGuard (r + 1) t
  | otherwise    = ((r, (fromJust c)), U)
  where c = elemIndex '^' h

consume :: [Char] -> (Grid, Guard)
consume file =
  let toCell = \ln -> case ln of
                        '#' -> Obs
                        otherwise -> Clr
      grid = map (map toCell) (lines file)
      guard = findGuard 0 (lines file) in
        (grid, guard)

step :: Grid -> Guard -> Maybe Guard
step grid (coo@(r,c), U)
  | cell == Nothing = Nothing
  | cell == Just Obs = Just (coo, R)
  | cell == Just Clr = Just (ncoo, U)
  where
    ncoo = (r - 1, c)
    cell = grid !? ncoo
step grid (coo@(r,c), R)
  | cell == Nothing = Nothing
  | cell == Just Obs = Just (coo, D)
  | cell == Just Clr = Just (ncoo, R)
  where
    ncoo = (r, c + 1)
    cell = grid !? ncoo
step grid (coo@(r,c), D)
  | cell == Nothing = Nothing
  | cell == Just Obs = Just (coo, L)
  | cell == Just Clr = Just (ncoo, D)
  where
    ncoo = (r + 1, c)
    cell = grid !? ncoo
step grid (coo@(r,c), L)
  | cell == Nothing = Nothing
  | cell == Just Obs = Just (coo, U)
  | cell == Just Clr = Just (ncoo, L)
  where
    ncoo = (r, c - 1)
    cell = grid !? ncoo

(?:) :: Eq a => a -> [a] -> [a]
el ?: arr
  | el `elem` arr = arr
  | otherwise = el:arr

route :: Grid -> [Coo] -> Guard -> [Coo]
route grid seen guard@(coo, dir)
  | nguard == Nothing = coo ?: seen
  | otherwise         = route grid (coo ?: seen) (fromJust nguard)
  where
    nguard = step grid guard

route2 :: Grid -> [Guard] -> Bool
route2 grid seen@(guard:_)
  | nguard == Nothing             = False
  | (fromJust nguard) `elem` seen = True
  | otherwise                     = route2 grid ((fromJust nguard):seen)
  where
    nguard = step grid guard

putObs :: Grid -> Coo -> Grid
putObs grid (r, c) =
  let (br,tr:ar) = splitAt r grid
      (bc,_:ac)  = splitAt c tr in
        br ++
        [bc ++ [Obs] ++ ac] ++
        ar

part1 :: Grid -> Guard -> Int
part1 grid guard =
  length $ route grid [] guard

part2 :: Grid -> Guard -> Int
part2 grid guard =
  length $
    filter (route2A . putObsA) $
    route grid [] guard
  where
    route2A = (flip route2) [guard]
    putObsA = putObs grid

main = do
  file <- readFile "day06.input"
  let (grid, guard) = consume file
  putStrLn ("Part 1: " ++ (show $ part1 grid guard))
  putStrLn ("Part 2: " ++ (show $ part2 grid guard))