#!/usr/bin/env nix-shell
#!nix-shell
import Aoc (readAndParseStdin)
import Data.Char (ord)
import qualified Data.Map as Map
import qualified Data.Map.Ordered as OM
import Text.Parsec (char, choice, digit, many, many1, noneOf, parse, sepBy)
import Text.Parsec.String (Parser)
main = do
input <- readAndParseStdin parseInput
print $ part1 input
print $ part2 input
part1 :: [String] -> Int
part1 input = sum $ map hash input
part2 :: [String] -> Int
part2 input = calcScore $ foldl folder (Map.fromList (zip [0 ..] $ replicate 256 OM.empty)) input
where
calcScore :: Map.Map Int (OM.OMap String Int) -> Int
calcScore = Map.foldrWithKey (\i v acc -> acc + ((i + 1) * sum (zipWith (*) [1 ..] (map snd $ OM.assocs v)))) 0
folder :: Map.Map Int (OM.OMap String Int) -> String -> Map.Map Int (OM.OMap String Int)
folder acc input =
let (label, add, focalLength) = parsePart2Input input
hashedLabel = hash label
updateBox = OM.alter (\_ -> if add then Just focalLength else Nothing) label
in Map.adjust updateBox hashedLabel acc
hash :: String -> Int
hash = foldl (\acc c -> ((acc + ord c) * 17) `rem` 256) 0
parsePart2Input :: String -> (String, Bool, Int)
parsePart2Input input = case parse doParse "" input of
Left parseError -> error $ show parseError
Right doc -> doc
where
doParse = do
label <- many1 (noneOf "-=")
operation <-
choice
[ True <$ char '=',
False <$ char '-'
]
focalLength <- read <$> many digit
return (label, operation, focalLength)
parseInput :: Parser [String]
parseInput = many1 (noneOf ",\n") `sepBy` char ','