import Data.Array
import Data.List.Split

main :: IO ()
main = do
  input <- readFile "input"
  let [grid, motion] = splitOn "\n\n" input
  let m = filter (/= '\n') motion
  print $ part2 (makeGrid $ scaleGrid $ lines grid) m 
--  print $ part1 (makeGrid $ lines grid) m 


part2 grid motion =
  let start = fst . head . filter ((== '@') . snd) . assocs $ grid
      final = walkGrid2 grid start motion
  in sum . map (calcCoords . fst) . filter ((== '[') . snd) . assocs $ final

walkGrid2 grid pos [] = grid
walkGrid2 grid pos (m:motion) =
  let (g, p) = moveOnce2 grid pos (getDir m)
  in walkGrid2 g p motion

-- urgh, a bit ugly
moveOnce2 grid pos (dy,0) =
  let next = move pos dir
  in case grid!next of
       '#' -> (grid, pos)
       '.' -> (grid // [(next, grid!pos), (pos, '.')], next)
       ']' -> moveTwo (move pos (0,-1))
       '[' -> moveTwo (move pos (0,1))
  where dir = (dy,0)
        moveTwo pos2 =
          let n1 = move pos dir
              n2 = move pos2 dir
              (g', pos') = moveOnce2 grid n1 dir
              (g'', pos2') = moveOnce2 g' n2 dir
          in if g''!n1 == '.' && g''!n2 == '.'
             then (g'' // [(n1, g''!pos), (pos, '.')], n1)
             else (grid, pos)
moveOnce2 grid pos (0,dx) =
  let next = move pos dir
  in case grid!next of
       '#' -> (grid, pos)
       '.' -> (grid // [(next, grid!pos), (pos, '.')], next)
       _ ->  let (g,p) = moveOnce2 grid next dir
             in if g!next == '.' then  (g // [(next, g!pos), (pos, '.')], next)
                else (grid, pos)
  where dir = (0,dx)

scaleGrid g = map (concatMap subst) g
  where subst '#' = "##"
        subst 'O' = "[]"
        subst '.' = ".."
        subst '@' = "@."

-----

part1 grid motion =
  let start = fst . head . filter ((== '@') . snd) . assocs $ grid
      final = walkGrid grid start motion
  in sum . map (calcCoords . fst) . filter ((== 'O') . snd) . assocs $ final

calcCoords (y,x) = y*100 + x

walkGrid grid pos [] = grid
walkGrid grid pos (m:motion) =
  let (g, p) = moveOnce grid pos (getDir m)
  in walkGrid g p motion

moveOnce grid pos dir =
  let next = move pos dir
      end = findNonBox grid pos dir
  in if (grid!end) == '#' then (grid, pos)
     else (grid // [(pos, '.'), (end, 'O'), (next, '@')], next)

findNonBox grid pos dir = 
  let next = move pos dir
  in if grid!next /= 'O' then next
     else findNonBox grid next dir

move (py,px) (dy,dx) = (py+dy,px+dx)

getDir 'v' = (1,0)
getDir '^' = (-1,0)
getDir '>' = (0,1)
getDir '<' = (0,-1)

makeGrid grid = array bounds elements
  where n = length grid
        m = length (head grid)
        bounds = ((0,0), (n-1,m-1))
        elements = [((i,j), (grid!!i!!j)) |
                    i <- [0..n-1], j <- [0..m-1]]