Untitled

parser :: (MonadParsec e s m, IsString (Tokens s), Token s ~ Char, Integral a) => m [(Tokens s, (a, [Tokens s]))]
parser = parseLine `sepEndBy` eol where
    parseLine = do
        string "Valve "
        src <- name
        string " has flow rate="
        w <- L.decimal
        string "; tunnel leads to valve " <|> string "; tunnels lead to valves "
        dsts <- name `sepEndBy1` string ", "
        pure (src, (w, dsts))
    name = takeWhile1P Nothing isAlphaNum

search :: (Ord a, Ord b) => (a -> (Maybe b, [(b, a)])) -> (b, a) -> [a]
search next = search' Set.empty Nothing . Heap.singleton @Heap.FstMaxPolicy where
    search' seen bestEstimate (Heap.view -> Just ((b, a), heap))
      | Set.member a seen = search' seen bestEstimate heap
      | fromMaybe False $ (<) <$> potential <*> bestEstimate = search' seen bestEstimate heap
      | otherwise = a : search' seen' (Just $ maybe b (max b) bestEstimate) heap' where
            (potential, nexts) = next a
            seen' = Set.insert a seen
            heap' = foldl' (flip Heap.insert) heap $ filter (flip Set.notMember seen' . snd) nexts
    search' _ _ _ = []

shortestPaths :: (Ord a, Monoid b, Ord b) => Map (a, a) b -> Map (a, a) b
shortestPaths es =
    Set.foldl' (\d a -> Set.foldl' (\d b -> Set.foldl' (flip $ update a b) d vs) d vs) d0 vs
  where
    vs = Set.fromList $ concat [[a, b] | (a, b) <- Map.keys es]
    d0 = Map.union es $ Map.fromSet (const mempty) $ Set.mapMonotonic (join (,)) vs
    update c b a d
      | Just x <- d Map.!? (a, c)
      , Just y <- d Map.!? (c, b)
      = Map.insertWith min (a, b) (x <> y) d
      | otherwise = d

(//) :: [a] -> [(Int, a)] -> [a]
as // ias = update (zip [0..] as) ias where
    update as@((i, a):as') bs@((j, b):bs') = case compare i j of
        LT -> a:update as' bs
        EQ -> b:update as' bs
        GT -> update as bs'
    update as _ = snd <$> as
infixl 9 //

day16 :: Int -> Int -> Text -> Either (ParseErrorBundle Text Void) Int
day16 n m input = do
    gr <- Map.fromList <$> parse parser "day16.txt" input
    let distances = fmap getSum . shortestPaths $ Map.fromList
            [((a, b), Sum 1) | (a, (_, bs)) <- Map.assocs gr, b <- bs]
        next (_, _, _, _, 0) = (Nothing, [])
        next (rooms, valves, flow, total, time)
          | null options = (Just estimate, [(estimate, (rooms, valves, flow, estimate, 0))])
          | otherwise = (Just potential, options)
          where
            estimate = total + flow * time
            potential = estimate + sum
              [ maximum $ 0 :
                  [ rate * (time - d - 1)
                  | (room, age) <- nubOrd rooms
                  , d <- maybeToList $ subtract age <$> distances Map.!? (room, room')
                  , 0 <= d && d < time
                  ]
              | (room', rate) <- Map.assocs valves
              ]
            moves = IntMap.fromListWith (IntMap.unionWith (<>))
              [ (d, IntMap.singleton i [(room', rate)])
              | (i, (room, age)) <- zip [0..] rooms
              , (room', rate) <- Map.assocs valves
              , d <- maybeToList $ subtract age <$> distances Map.!? (room, room')
              , 0 <= d && d < time
              ]
            options =
              [ ( estimate + rate * (time - d - 1)
                , ( sort $ (second (d + 1 +) <$> rooms) // zip is ((, 0) <$> rooms')
                  , Map.withoutKeys valves $ Set.fromList rooms'
                  , flow + rate
                  , total + flow * (d + 1)
                  , time - d - 1
                  )
                )
              | (d, moves') <- IntMap.assocs moves
              , (is, moves'') <- fmap unzip . filterM (const [False, True]) $ IntMap.assocs moves' 
              , moves'''@(_:_) <- sequence moves''
              , let (rooms', sum -> rate) = unzip moves'''
              , and . zipWith Set.notMember rooms' $ scanl' (flip Set.insert) Set.empty rooms'
              ]
        max' total (rooms, valves, flows, total', time)
          | total' > total = traceShow (rooms, Map.keys valves, flows, total', time) total'
          | otherwise = total
    pure . foldl' max' 0 $ search next
        (0, (replicate n ("AA", 0), Map.filter (> 0) $ fst <$> gr, 0, 0, m))