{-# LANGUAGE TupleSections #-}
module Day14 where

import qualified Data.Map as M
import Data.List (sort, sortOn)

type Pair = (Char,Char)
type PairInsertion = M.Map Pair Char

parseRule :: String -> (Pair,Char)
parseRule (a:b:s) = ((a,b), last s)

convertPair :: PairInsertion -> Pair -> [Pair]
convertPair m p@(a,b) = case M.lookup p m of
    Just c -> [(a,c), (c,b)]
    Nothing -> [(a,b)]

stepPairs :: PairInsertion -> [(Pair,Int)] -> [(Pair, Int)]
stepPairs m = M.assocs . M.fromListWith (+) . concatMap (\(p,i) -> map (\np -> (np,i)) (convertPair m p))

pairsToCount :: [(Pair, Int)] -> [(Char, Int)]
pairsToCount = map (\(c,i) -> (c, (i+1)`div`2)) . M.assocs . M.fromListWith (+) . concatMap (\((a,b), freq) -> [(a,freq), (b,freq)])

stringToPairs :: String -> [Pair]
stringToPairs s = zip s (tail s)

main :: IO ()
main = do
    template <- map (\p -> (p,1)) . stringToPairs <$> getLine
    getLine
    transforms <- M.fromList . map parseRule . lines <$> getContents
    print $ pairsToCount template

    putStr "part 1: "
    let step10 = sortOn snd $ pairsToCount $ iterate (stepPairs transforms) template !! 10
    print $ snd (last step10) - snd (head step10)

    putStr "part 2: "
    let step10 = sortOn snd $ pairsToCount $ iterate (stepPairs transforms) template !! 40
    print $ snd (last step10) - snd (head step10)


    print "Bye"