{-# LANGUAGE MultiWayIf #-}

import Control.Arrow
import Data.Bifunctor hiding (first, second)

import Data.Array.Unboxed (UArray)
import Data.Array.ST (STUArray)
import Control.Monad.ST (ST, runST)
import Control.Monad (join)

import qualified Data.Char as Char
import qualified Data.List as List
import qualified Data.Set as Set
import qualified Data.Array.Unboxed as UArray
import qualified Data.Array.ST as MArray

parse :: String -> (UArray (Int, Int) Char, String)
parse s = (grid, orders)
        where 
                l = lines s
                orderLines = drop 1 . dropWhile (/= "") $ l
                orders     = foldl (++) "" $ orderLines
                gridLines  = takeWhile (/= "") $ l
                gridHeight = length gridLines
                gridWidth  = length . head $ gridLines
                grid       = UArray.listArray ((1, 1), (gridHeight, gridWidth)) . foldl (++) "" $ gridLines

moveRobot :: UArray (Int, Int) Char -> String -> ST s (UArray (Int, Int) Char)
moveRobot g s = do
        let robotPosition = maybe (error "Robot not in grid") fst . List.find ((== '@') . snd) . UArray.assocs $ g
        mg <- MArray.thaw g
        moveRobot' mg robotPosition s

type RobotPosition = (Int, Int)

walkDirection :: (Int, Int) -> (Int, Int) -> [(Int, Int)]
walkDirection p d = iterate (.+. d) p

orderDirection :: Char -> (Int, Int)
orderDirection '>' = ( 0,  1)
orderDirection '<' = ( 0, -1)
orderDirection '^' = (-1,  0)
orderDirection 'v' = ( 1,  0)

(y1, x1) .+. (y2, x2) = (y1 + y2, x1 + x2)

(y1, x1) .*. (y2, x2) = (y1 * y2, x1 * x2)

countBarrels :: STUArray s (Int, Int) Char -> RobotPosition -> (Int, Int) -> ST s Int
countBarrels g p d = do
        currentTile <- MArray.readArray g p
        if currentTile == 'O' then
                do
                        n <- countBarrels g (p .+. d) d
                        return $! n + 1
        else
                return 0

moveRobot' :: STUArray s (Int, Int) Char -> RobotPosition -> String -> ST s (UArray (Int, Int) Char)
moveRobot' g _ [] = MArray.freeze g
moveRobot' g p (o:os) = do
        let direction = orderDirection o
        let nextCoordinate = p .+. direction
        nextTile <- MArray.readArray g nextCoordinate
        case nextTile of
                '#' -> moveRobot' g p os
                '.' -> MArray.writeArray g p '.' 
                        *> MArray.writeArray g nextCoordinate '@'
                        *> moveRobot' g nextCoordinate os
                'O' -> do
                        barrelCount <- countBarrels g nextCoordinate direction
                        let postBarrelPosition = p .+. (direction .*. (1 + barrelCount, 1 + barrelCount))
                        postBarrelTile <- MArray.readArray g postBarrelPosition
                        case postBarrelTile of
                                '#' -> moveRobot' g p os
                                '.' -> MArray.writeArray g p '.'
                                        *> MArray.writeArray g nextCoordinate '@'
                                        *> MArray.writeArray g postBarrelPosition 'O'
                                        *> moveRobot' g nextCoordinate os

part1 (g, o) = UArray.assocs
        >>> filter (snd >>> (== 'O'))
        >>> map (uncurry (+) . ((*100) *** id) . (join bimap pred) . fst)
        >>> sum
        $ g'
        where 
                g' = runST $ (moveRobot g o)

translate :: Char -> String
translate '#'  = "##"
translate '.'  = ".."
translate '@'  = "@."
translate 'O'  = "[]"
translate '\n' = "\n"
translate c    = [c]

moveWideRobot :: UArray (Int, Int) Char -> String -> ST s (UArray (Int, Int) Char)
moveWideRobot g s = do
        let robotPosition = maybe (error "Robot not in grid") fst . List.find ((== '@') . snd) . UArray.assocs $ g
        mg <- MArray.thaw g
        moveWideRobot' mg robotPosition s

moveChestHorizontally g p d = do
        tile <- MArray.readArray g p
        case tile of
                '.' -> return True
                '#' -> return False
                _   -> do
                        let p' = p .+. d
                        canMove <- moveChestHorizontally g p' d
                        if canMove then MArray.writeArray g p' tile else return ()
                        return canMove

boxCounterpart ('[', (y, x)) = (']', (y, x+1))
boxCounterpart (']', (y, x)) = ('[', (y, x-1))

moveChestVertically :: STUArray s (Int, Int) Char -> [(Int, Int)] -> (Int, Int) -> ST s Bool
moveChestVertically g [] d = return True
moveChestVertically g ps d = do
        tiles <- flip zip ps <$> mapM (MArray.readArray g) ps
        let counterParts = List.map boxCounterpart . List.filter (fst >>> flip List.elem "[]") $ tiles
        let tiles' = List.nub $ tiles ++ counterParts
        if | any ((== '#') . fst) tiles' -> return False
           | otherwise -> do
                let boxTiles = List.filter (fst >>> flip List.elem "[]") $ tiles'
                let boxPositions = List.map snd $ boxTiles
                let positionsAhead = List.map (.+. d) $ boxPositions
                success <- moveChestVertically g positionsAhead d
                if success then do
                        mapM_ (second (.+. d) >>> uncurry (flip (MArray.writeArray g))) boxTiles
                        mapM_ (flip (MArray.writeArray g) '.') boxPositions
                else 
                        return ()

                return $ success


moveWideRobot' :: STUArray s (Int, Int) Char -> RobotPosition -> String -> ST s (UArray (Int, Int) Char)
moveWideRobot' g p [] = MArray.freeze g
moveWideRobot' g p (o:os) = do
        let direction = orderDirection o
        let nextCoordinate = p .+. direction
        nextTile <- MArray.readArray g nextCoordinate
        case nextTile of
                '#' -> moveWideRobot' g p os
                '.' -> MArray.writeArray g p '.'
                        *> MArray.writeArray g nextCoordinate '@'
                        *> moveWideRobot' g nextCoordinate os
                '[' -> do
                        success <- if o == '>' 
                        then 
                                moveChestHorizontally g nextCoordinate direction
                        else
                                moveChestVertically g [nextCoordinate, second succ nextCoordinate] direction

                        if success then do
                                MArray.writeArray g p '.'
                                MArray.writeArray g nextCoordinate '@'
                                moveWideRobot' g nextCoordinate os
                        else
                                moveWideRobot' g p os
                ']' -> do
                        success <- if o == '<'
                        then
                                moveChestHorizontally g nextCoordinate direction
                        else
                                moveChestVertically g [nextCoordinate, second pred nextCoordinate] direction

                        if success then do
                                MArray.writeArray g p '.'
                                MArray.writeArray g nextCoordinate '@'
                                moveWideRobot' g nextCoordinate os
                        else
                                moveWideRobot' g p os

part2 (g, o) = UArray.assocs
        >>> List.filter (snd >>> (== '['))
        >>> map (uncurry (+) . ((*100) *** id) . (join bimap pred) . fst)
        >>> sum
        $ g'
        where
                g' = runST $ (moveWideRobot g o)

main = getContents
        >>= print
        . (part1 *** part2)
        . join bimap parse
        . second (List.concatMap translate)
        . join (,)