module Main where

import qualified Data.Array as A
import Data.List
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)

import Input
import Util


step :: Ord a => Map (a, a) a -> [a] -> [a]
step rules (a:b:s) = case Map.lookup (a, b) rules of
                       Nothing -> a : step rules (b:s)
                       Just c -> a : c : step rules (b:s)
step _ l = l

maphistogram :: Ord a => [a] -> Map a Int
maphistogram = foldr (\x -> Map.insertWith (+) x 1) mempty

-- Dynamic programming: table
--   (all letter pairs) x (number of steps) x (all letter pairs)
-- where
--   table[p][n][q] means the number of times that q appears in the string
--   produced when stepping n times starting from the string p.
-- Defining relation:
--   table[p][0][q] = if p == q then 1 else 0
--   table[p][n][q] = sum over pairs r:
--                      table[p][n-1][r] * (number of occurrences of q in step(r))
dyn :: Map (Char, Char) Char
    -> Int
    -> (Map Char Int, Int, A.Array ((Int, Int), Int, (Int, Int)) Int)
dyn rules nsteps = (charmap, maxid, table)
  where
    chars = uniq . sort $ concatMap (\((a,b),c) -> [a,b,c]) (Map.assocs rules)
    charlist = zip chars [0::Int ..]
    charmap = Map.fromList charlist
    maxid = snd (fst (Map.deleteFindMax charmap))
    table = A.array (((0, 0), 0, (0, 0)), ((maxid, maxid), nsteps, (maxid, maxid)))
        [(((p1, p2), n, (q1, q2))
         ,if n == 0 then if p == q then 1 else 0
                    else sum [table A.! (p, n - 1, r) * count (q1c, q2c) (step rules [r1c, r2c])
                             | (r1c, r1) <- charlist, (r2c, r2) <- charlist
                             , let r = (r1, r2)])
        | p1 <- [0..maxid], p2 <- [0..maxid]
        , let p = (p1, p2)
        , n <- [0..nsteps]
        , (q1c, q1) <- charlist, (q2c, q2) <- charlist
        , let q = (q1, q2)]
    count (a, b) (x:y:s) | a == x, b == y = 1 + count (a, b) (y:s)
                         | otherwise = count (a, b) (y:s)
    count _ _ = 0

solve :: Map (Char, Char) Char -> String -> Int -> Int
solve rules template nsteps =
    let (charmap, maxid, table) = dyn rules nsteps
        inttemplate = map (charmap Map.!) template
        hist = Map.fromListWith (+) $ concat
            [let n = table A.! ((a, b), nsteps, (c, d))
             in [(c, n), (d, n)]
            | (a, b) <- zip inttemplate (tail inttemplate)
            , c <- [0..maxid]
            , d <- [0..maxid]]
        hist' = Map.map (`div` 2) . Map.unionWith (+) (Map.fromList [(charmap Map.! head template, 1), (charmap Map.! last template, 1)]) $ hist
        -- freqs = Map.assocs hist
    in -- trace (show table ++ "\n" ++ show freqs) $
       maximum (Map.elems hist') - minimum (Map.elems hist')

main :: IO ()
main = do
    template : _ : rules' <- getInput 14
    let rules = Map.fromList (map (\[a,b,' ','-','>',' ',c] -> ((a, b), c)) rules')
    -- let string1 = iterate (step rules) template !! 10
    --     hist1 = Map.elems (maphistogram string1)
    -- print (maximum hist1 - minimum hist1)
    print (solve rules template 10)
    print (solve rules template 40)