import           Data.List       (maximumBy, minimumBy)
import           Data.Map.Strict (Map, (!))
import qualified Data.Map.Strict as Map
import           Data.Ord        (comparing)
import           Itertools       (flatmap, windows)
import           Parsing         (splitByString)
import           Reducers        (leastMostOcc)

main :: IO ()
main = do
  input <- getContents
  let
    polymer = parsePolymer input
  (putStrLn . show . solution1) polymer
  (putStrLn . show . solution2) polymer

parsePolymer :: String -> ([Char], Map [Char] Char)
parsePolymer input = (template input, rules input)
  where
    template :: String -> String
    template = head . splitByString "\n\n"
    rules :: String -> Map [Char] Char
    rules = Map.fromList . map tupelize . map (splitByString " -> ") . lines . last . splitByString "\n\n"
    tupelize :: [String] -> ([Char], Char)
    tupelize untupled = (head untupled, (head . last) untupled)

solution1 :: ([Char], Map [Char] Char) -> Int
solution1 (template, rules) = uncurry (-) (leastMostOcc (polymerizeN rules template 10))

solution2 :: ([Char], Map [Char] Char) -> Int
solution2 (template, rules) = ((snd . most) polyCounts) - ((snd . least) polyCounts)
  where
    polyCounts :: Map Char Int
    polyCounts = polyCounter rules template 40
    least = minimumBy (comparing snd) . Map.assocs
    most = maximumBy (comparing snd) . Map.assocs

polymerizeN :: Map [Char] Char -> [Char] -> Int -> [Char]
polymerizeN rules template i = iterate (polymerize rules) template !! i

polymerize :: Map [Char] Char -> [Char] -> [Char]
polymerize rules template = step template
  where
    step :: [Char] -> [Char]
    step []      = []
    step [final] = [final]
    step (x:xs)  = x:(rules ! [x, head xs]):(step xs)

-- Since we don't actually need to synthesize the polymer, but only
-- need to know how many of each elements occur, we can get around
-- having to create a ridiculously large list by operating on the rule
-- pairs instead.
--
-- E.g. NNCB -> NCNBCHB (see testInput1), in this case we turn:
--
--   - NN -> NC, CN
--   - NC -> NB, BC
--   - CB -> CH, HB
--
-- All we need to do is keep track of the counts of these
-- sub-sequences, not of the full polymer.
--
-- To get back to the individual elements, we then just need to
-- de-window it all again.
polyCounter :: Map [Char] Char -> [Char] -> Int -> Map Char Int
polyCounter rules template iterations = elCounts
  where
    initial :: Map [Char] Int
    initial = (Map.fromListWith (+) . map (\e -> (e, 1)) . windows 2) template
    step :: Map [Char] Int -> Map [Char] Int
    step = Map.fromListWith (+) . flatmap resultingPolies . Map.assocs
    resultingPolies :: ([Char], Int) -> [([Char], Int)]
    resultingPolies (p, i) = [([head p, c], i), ([c, last p], i)]
      where c = rules ! p
    polyCounts :: Map [Char] Int
    polyCounts = iterate step initial !! iterations

    -- Because the windows are size 2, every element of a sub-binding is
    -- counted twice (e.g. NBCBC -> NB, BC, CB, BC), except the first and
    -- last element. We know the first and last elements (because we
    -- operated on windows, they'll match the first and last of the
    -- template), so we can un-window the whole batch by dividing their
    -- number of occurrences by two, except for the first and last, where
    -- we subtract the number by 1 first and add 1 back later.
    elCounts :: Map Char Int
    elCounts = Map.mapWithKey deWindow naiveSum
      where
        els :: ([Char], Int) -> [(Char, Int)]
        els (p, i) = [(head p, i), (last p, i)]
        naiveSum :: Map Char Int
        naiveSum = (Map.fromListWith (+) . flatmap els . Map.assocs) polyCounts
        deWindow :: Char -> Int -> Int
        deWindow c i
          | c == head template || c == last template = (i - 1) `quot` 2 + 1
          | otherwise = i `quot` 2

-- Tests

testInput1 = unlines [
  "NNCB",
  "",
  "CH -> B",
  "HH -> N",
  "CB -> H",
  "NH -> C",
  "HB -> C",
  "HC -> B",
  "HN -> C",
  "NN -> C",
  "BH -> H",
  "NC -> B",
  "NB -> B",
  "BN -> B",
  "BB -> N",
  "BC -> B",
  "CC -> N",
  "CN -> C"]

parsedTestInput1 = parsePolymer testInput1
test1 = solution1 parsedTestInput1 -- 1588
test2 = solution2 parsedTestInput1 -- 2188189693529