Day5.hs

module Day5 where

import           Advent.AoC
import           Control.Applicative        ((<|>))
import           Control.Monad              (replicateM)
import           Data.Foldable              (foldl')
import           Data.IntMap                (IntMap)
import qualified Data.IntMap.Strict         as IM
import           Data.Maybe                 (mapMaybe)
import           Text.Megaparsec            (anySingle, endBy, manyTill, sepBy)
import           Text.Megaparsec.Char       (space, upperChar)
import qualified Text.Megaparsec.Char.Lexer as L

type Stacks = IntMap [Char]

data Instruction = Instruction Int Int Int
  deriving (Eq, Show)

getInput :: FilePath -> IO (Stacks, [Instruction])
getInput = parseFile ((,) <$> parseStacks <*> parseInstruction `endBy` "\n")
  where
    parseStacks = (IM.unionsWith (<>) <$> parseStackLine `sepBy` "\n") <*
                    lexeme (manyTill anySingle "\n") -- ignore the position numbers

    parseStackLine = IM.fromList . mapMaybe sequence . zip [1..] <$> maybeEntry `sepBy` " "

    maybeEntry = (Just <$> anEntry) <|> (Nothing <$ "   ")
    anEntry = "[" *> replicateM 1 upperChar <* "]"

    parseInstruction = Instruction <$> ("move " *> lexeme num) <*> ("from " *> lexeme num) <*> ("to " *> num)
    lexeme = L.lexeme space
    num = L.decimal

go :: ([Char] -> [Char]) -> Stacks -> [Instruction] -> String
go shuffle stacks = foldMap (take 1) . foldl' run stacks
  where
    run st (Instruction n f t) = IM.adjust (shuffle movef <>) t $ IM.insert f newf st
      where
        (movef, newf) = splitAt n (st IM.! f)

part1 :: IO String
part1 = uncurry (go reverse) <$> getInput "input/day5"

part2 :: IO String
part2 = uncurry (go id) <$> getInput "input/day5"