module Day21monad

import Prelude
import Node
import Aoc
import SortedMap
import Parser

record State st a where
  constructor ST
  runState : st → (st × a)

instance ∀ st. Applicative (State st) where
  return x = ST (\st => (st, x))
  (ST sf) <*> (ST sx) = ST $ \st =>
    let (st', f) = sf st in
    let (st'', x) = sx st' in
    (st'', f x)

instance ∀ st. Monad (State st) where
  pure x = ST (\st => (st, x))
  bind (ST run) f = ST $ \st =>
    let (st', a) = run st in
    let (ST run') = f a in
    run' st'

get : ∀ st. State st st
get = ST (\st => (st, st))

put : ∀ st. st → State st Unit
put st = ST (\_ => (st, MkUnit))


min : Int → Int → Int
min a b = if a < b then a else b

gridPoints : String → List (Char × Int × Int)
gridPoints text = go 0 0 (unpack text) Nil
  where
    -- might as well be tail recursive
    go : Int → Int → List Char → List (Char × Int × Int) → List (Char × Int × Int)
    go row col Nil points = points
    go row col ('\n' :: cs) points = go (row + 1) 0 cs points
    go row col (c :: cs) points = go row (col + 1) cs ((c,row,col) :: points)

data Dir : U where North East South West : Dir

dirs : List Dir
dirs = (North :: South :: East :: West :: Nil)

move : Point → Dir → Point
move (r, c) North  = (r - 1, c)
move (r, c) East   = (r, c + 1)
move (r, c) South  = (r + 1, c)
move (r, c) West   = (r, c - 1)

Costs : U
Costs = SortedMap (Point × Point) Int

-- linked list of keypads
record Keypad where
  constructor KP
  name : String
  start : Point
  interdit : Point
  costs : Costs -- cache of costs
  next : Maybe Keypad

getPaths : Point → Point → Point → List (List Dir)
getPaths interdit pt@(a,b) to@(c,d) =
  if pt == to then Nil :: Nil else
  if pt == interdit then Nil else
  join $ map check dirs
  where
    check : Dir → List (List Dir)
    check North = if c < a then map (_::_ North) $ getPaths interdit (move pt North) (c,d) else Nil
    check South = if a < c then map (_::_ South) $ getPaths interdit (move pt South) (c,d) else Nil
    check East = if b < d then map (_::_ East) $ getPaths interdit (move pt East) (c,d) else Nil
    check West = if d < b then map (_::_ West) $ getPaths interdit (move pt West) (c,d) else Nil

updateCost : Point × Point → Int → State Keypad Unit
updateCost path cost = do
  kp <- get
  put (KP (.name kp) (.start kp) (.interdit kp) (updateMap path cost (.costs kp)) (.next kp))

keyPos : Dir → Point
keyPos North = (0,1)
keyPos South = (1,1)
keyPos East  = (1,2)
keyPos West  = (1,0)

-- cost to run a path in a keypad
pathCost : Point → Point → State Keypad Int

-- cost of sequence of points (run in parent keypad)
-- for numpad, we pick points from the map, for the rest map keyPos ...
seqCost : Point → List Point → Int → State Keypad Int
seqCost cur Nil cost = pure cost
seqCost cur (pt :: pts) cost = do
  cost' <- pathCost cur pt
  seqCost pt pts (cost + cost')

-- cost of best path from -> to in kp
pathCost from to = do
  kp <- get {Keypad}
  case lookupMap (from, to) (.costs kp) of
    Just (_, cost) => pure {State Keypad} cost
    Nothing =>
      let (path :: paths) = getPaths (.interdit kp) from to | _ => fatalError "empty paths" in
      case kp of
        (KP n s i c Nothing) => pure 1
        (KP n s i c (Just kp')) => do
          put kp'
          cost <- mincost path paths
          kp' <- get
          put $ KP n s i c (Just kp')
          -- need to ensure cost goes into correct kp
          updateCost (from,to) cost
          pure cost
  where
    xlate : List Dir → Point -> List Point
    xlate Nil a = a :: Nil
    xlate (d :: ds) a = keyPos d :: xlate ds a

    mincost : List Dir → List (List Dir) → State Keypad Int
    mincost path paths = do
      kp <- get
      cost <- seqCost (0,2) (xlate path $ kp.start) 0
      case paths of
        Nil => pure cost
        (path :: paths) => do
          cost' <- mincost path paths
          pure $ min cost cost'

fromList : ∀ k v. {{Ord k}} {{Eq k}} → List (k × v) → SortedMap k v
fromList xs = foldMap (\ a b => b) EmptyMap xs

getNum : String → Int
getNum str = case number (unpack str) of
  Right (n, _) => n
  _ => 0

runOne : SortedMap Char Point → String → State Keypad (Int × Int)
runOne numpad str = do
  let pts = map snd $ mapMaybe (flip lookupMap numpad) $ unpack str
  res <- seqCost (3,2) pts 0
  pure (getNum str, res)

makeKeypad : Int → Keypad -> Keypad
makeKeypad 0 kp = kp
makeKeypad n kp = makeKeypad (n - 1) $ KP (show n) (0,2) (0,0) EmptyMap (Just kp)

run : String -> IO Unit
run fn = do
  putStrLn fn
  text <- readFile fn
  let codes = split (trim text) "\n"

  -- the space is illegal spot
  let numpad = fromList $ filter (not ∘ _==_ ' ' ∘ fst) $ gridPoints "789\n456\n123\n 0A"
  printLn $ toList numpad

  let rob1 = KP "r1" (0,2) (0,0) EmptyMap Nothing
  let robn = makeKeypad 2 rob1
  let kp = KP "kp" (3,2) (3,0) EmptyMap (Just robn)
  let p1 = foldl _+_ 0 $ map (uncurry _*_) $ snd $ .runState (traverse (runOne numpad) codes) kp
  putStrLn $ "part1 " ++ show p1

  let rob1 = KP "r1" (0,2) (0,0) EmptyMap Nothing
  let robn = makeKeypad 25 rob1
  let kp = KP "kp" (3,2) (3,0) EmptyMap (Just robn)
  let p2 = foldl _+_ 0 $ map (uncurry _*_) $ snd $ .runState (traverse (runOne numpad) codes) kp
  putStrLn $ "part2 " ++ show p2

main : IO Unit
main = do
  run "aoc2024/day21/eg.txt"
  run "aoc2024/day21/input.txt"