import Data.Char

main :: IO ()
main = do
  input <- readFile "input"
--  print $ part1 (init input)
  print $ part2 (init input)

----

part2 input =
  let blocks = expandBlocks 0 . map digitToInt $ input
      lastId = (length blocks) `div` 2
      defrag = expandForChk $ moveFiles lastId blocks
  in dot defrag [0..]
  where moveFiles 0 blocks = blocks
        moveFiles id blocks = moveFiles (id-1) (moveFile id blocks)

data Blocks = Free Int | File Int Int
instance Show Blocks where
  show (Free n) = "Free " ++ show n
  show (File id n) = "(File " ++ show id ++ ") " ++ show n

expandBlocks id [] = []
expandBlocks id [n] = [(File id n)]
expandBlocks id (n1:n2:ns) =
  (File id n1):(Free n2):(expandBlocks (id+1) ns)

expandForChk [] = []
expandForChk ((Free n):blocks) 
 | n == 0 = expandForChk blocks
 | otherwise = 0:expandForChk ((Free (n-1)):blocks)
expandForChk ((File id n):blocks)
 | n == 0 = expandForChk blocks
 | otherwise = id:expandForChk ((File id (n-1)):blocks)

moveFile id blocks =
  let f = findFile id blocks
      sp = findSpaceBefore f [] blocks
      (File _ n) = f
  in case sp of
    Nothing -> blocks
    Just (pref, ((Free m):rest)) -> pref ++ (f:removeFile id ((Free (m-n)):rest))

findFile id ((File fid n):blocks)
  | fid == id = (File fid n)
  | otherwise = findFile id blocks
findFile id (_:blocks) = findFile id blocks

findSpaceBefore (File id n) pref (File fid m:blocks)
  | id == fid = Nothing
  | otherwise = findSpaceBefore (File id n) ((File fid m):pref) blocks
findSpaceBefore (File id n) pref (Free m:blocks)
  | n <= m = Just (reverse pref, Free m:blocks)
  | otherwise = findSpaceBefore (File id n) ((Free m):pref) blocks

removeFile id blocks = mergeFree $ justRemoveFile id blocks

justRemoveFile id [] = []
justRemoveFile id (Free n:blocks) = (Free n:justRemoveFile id blocks)
justRemoveFile id (File fid n:blocks)
  | fid == id = (Free n:blocks)
  | otherwise = (File fid n:justRemoveFile id blocks)

mergeFree [] = []
mergeFree [Free _] = []
mergeFree (Free 0:blocks) = mergeFree blocks
mergeFree (Free n1:Free n2:blocks) = mergeFree (Free (n1+n2):blocks)
mergeFree (b:blocks) = (b:mergeFree blocks)

----

part1 input =
  let blocks = map digitToInt input
      files = expandIDs 0 (everyOther blocks)
      filesBak = reverse files
      n = length files
      defrag = take n $ mergeAlternate blocks files filesBak
  in dot defrag [0..]

mergeAlternate [n] l1 l2 = take n l1
mergeAlternate (n1:n2:ns) l1 l2 =
  let (l1a,l1b) = splitAt n1 l1
      (l2a,l2b) = splitAt n2 l2
  in l1a ++ l2a ++ (mergeAlternate ns l1b l2b)

expandIDs id [] = []
expandIDs id (0:ns) = expandIDs (id+1) ns
expandIDs id (n:ns) = (id:expandIDs id (n-1:ns))

everyOther [] = []
everyOther [x] = [x]
everyOther (x:_:xs) = (x:everyOther xs)

dot l1 l2 = sum $ zipWith (*) l1 l2