import Control.Monad
import Data.List
import Data.Map (Map)
import Data.Map qualified as Map

type Pos = (Int, Int)

data Keypad = Keypad (Map Char Pos) Pos

makeKeypad :: [[Char]] -> Keypad
makeKeypad rows =
  (\m -> Keypad m (m Map.! 'A')) $
    Map.fromList [(c, (i, j)) | (i, r) <- zip [0 ..] rows, (j, c) <- zip [0 ..] r, c /= '_']

numericKeypad = makeKeypad ["789", "456", "123", "_0A"]

directionalKeypad = makeKeypad ["_^A", "<v>"]

movesToButton :: Keypad -> Pos -> Pos -> [[Char]]
movesToButton (Keypad btns _) (i1, j1) (i2, j2) =
  let di = i2 - i1
      dj = j2 - j1
      v = replicate (abs di) $ if di > 0 then 'v' else '^'
      h = replicate (abs dj) $ if dj > 0 then '>' else '<'
      hv = guard ((i1, j2) `elem` btns) >> return (h ++ v)
      vh = guard ((i2, j1) `elem` btns) >> return (v ++ h)
   in (++ ['A']) <$> nub (hv ++ vh)

sequenceMoves :: Keypad -> [Char] -> [[Char]]
sequenceMoves keypad@(Keypad btns start) =
  (fmap concat . sequence) . (zipWith (movesToButton keypad) =<< (start :)) . map (btns Map.!)

indirectLength :: Int -> [Char] -> Int
indirectLength levels = go levels
  where
    Keypad btns start = directionalKeypad
    go l = sum . (zipWith (\p1 p2 -> lengths Map.! (l, p1, p2)) =<< (start :)) . map (btns Map.!)
    lengths =
      let ps = Map.elems btns
       in Map.fromList [((l, p1, p2), bestLength l p1 p2) | l <- [1 .. levels], p1 <- ps, p2 <- ps]
    bestLength l p1 p2 =
      minimum . map (if l == 1 then length else go (l - 1)) $ movesToButton directionalKeypad p1 p2

complexity :: Int -> String -> Int
complexity bots code = bestLength code * read (init code)
  where
    bestLength = minimum . map (indirectLength bots) . sequenceMoves numericKeypad

main = do
  input <- lines <$> readFile "input21"
  mapM_ (print . sum . flip map input . complexity) [2, 25]