{-# LANGUAGE MultiWayIf #-}

import Control.Arrow

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

import qualified Data.List as List
import qualified Data.Maybe as Maybe
import qualified Data.Array.MArray as MArray

toNumber '0' = 0
toNumber '1' = 1
toNumber '2' = 2
toNumber '3' = 3
toNumber '4' = 4
toNumber '5' = 5
toNumber '6' = 6
toNumber '7' = 7
toNumber '8' = 8
toNumber '9' = 9

parse :: String -> [Int]
parse s = filter (/= '\n')
        >>> map toNumber
        >>> zip [0..]
        >>> List.concatMap (\ (index, n) -> if index `mod` 2 == 0 then replicate n (index `div` 2) else replicate n (-1))
        $ s

calculateChecksum :: [Int] -> Int
calculateChecksum = zip [0..]
        >>> filter (snd >>> (/= -1))
        >>> map (uncurry (*))
        >>> sum

moveFiles :: [Int] -> ST s Int
moveFiles bs = do
        let bLength = length bs
        marray <- MArray.newListArray (1, bLength) bs
        moveFiles' marray 1 bLength
        elems <- MArray.getElems marray
        return $ calculateChecksum elems


moveFiles' :: STUArray s Int Int -> Int -> Int -> ST s ()
moveFiles' a start stop
        | start == stop = return ()
        | otherwise = do
                stopBlock <- MArray.readArray a stop

                if stopBlock == -1
                then
                        moveFiles' a start (pred stop)
                else
                        do
                                startBlock <- MArray.readArray a start
                                if startBlock == -1
                                then
                                        do
                                                MArray.writeArray a start stopBlock
                                                MArray.writeArray a stop (-1)
                                                moveFiles' a (succ start) (pred stop) 
                                else
                                        moveFiles' a (succ start) stop

countConsecutive :: STUArray s Int Int -> Int -> Int -> ST s Int
countConsecutive a i step = do
        block <- MArray.readArray a i
        let nextI = i + step
        bounds <- MArray.getBounds a
        if      | MArray.inRange bounds nextI ->
                        do
                                nextBlock <- MArray.readArray a nextI
                                if nextBlock == block
                                then
                                        do
                                                steps <- countConsecutive a nextI step
                                                return $ 1 + steps
                                else
                                        return 1
                | otherwise -> return 1

findEmpty :: STUArray s Int Int -> Int -> Int -> Int -> ST s (Maybe Int)
findEmpty a i l s = do
        block <- MArray.readArray a i
        blockLength <- countConsecutive a i 1
        let nextI = i + blockLength
        bounds <- MArray.getBounds a
        let nextInBounds = MArray.inRange bounds nextI

        if      | i >= s                           -> return $! Nothing
                | block == -1 && blockLength >= l  -> return $ Just i
                | block /= -1 && nextInBounds      -> findEmpty a nextI l s
                | blockLength <= l && nextInBounds -> findEmpty a nextI l s
                | not nextInBounds                 -> return $! Nothing

moveDefragmenting :: [Int] -> ST s Int
moveDefragmenting bs = do
        let bLength = length bs
        marray <- MArray.newListArray (1, bLength) bs
        moveDefragmenting' marray bLength
        elems <- MArray.getElems marray
        return $ calculateChecksum elems

moveDefragmenting' :: STUArray s Int Int -> Int -> ST s ()
moveDefragmenting' a 1    = return ()
moveDefragmenting' a stop
        | otherwise = do
                stopBlock  <- MArray.readArray a stop
                stopLength <- countConsecutive a stop (-1)
                targetBlock <- findEmpty a 1 stopLength stop

                elems <- MArray.getElems a

                let nextStop = stop - stopLength
                bounds <- MArray.getBounds a
                let nextStopInRange = MArray.inRange bounds nextStop
                
                if      | stopBlock == -1
                                -> moveDefragmenting' a nextStop
                        | Maybe.isJust targetBlock 
                                -> do
                                        let target = Maybe.fromJust targetBlock
                                        mapM_ (\ o -> MArray.writeArray a (stop - o) (-1)) [0..stopLength - 1]
                                        mapM_ (\ o -> MArray.writeArray a (target + o) stopBlock) [0..stopLength - 1]
                                        if nextStopInRange then moveDefragmenting' a nextStop else return ()
                        | nextStopInRange -> moveDefragmenting' a nextStop
                        | otherwise -> return ()
                                

part1 bs = runST $ moveFiles bs
part2 bs = runST $ moveDefragmenting bs

main = getContents
        >>= print
        . (part1 &&& part2)
        . parse