{-# Language QuasiQuotes, ImportQualifiedPost, BlockArguments #-}
{-|
Module      : Main
Description : Day 16 solution
Copyright   : (c) Eric Mertens, 2022
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2022/day/16>

>>> :{
:main +
    "Valve AA has flow rate=0; tunnels lead to valves DD, II, BB\n\
    \Valve BB has flow rate=13; tunnels lead to valves CC, AA\n\
    \Valve CC has flow rate=2; tunnels lead to valves DD, BB\n\
    \Valve DD has flow rate=20; tunnels lead to valves CC, AA, EE\n\
    \Valve EE has flow rate=3; tunnels lead to valves FF, DD\n\
    \Valve FF has flow rate=0; tunnels lead to valves EE, GG\n\
    \Valve GG has flow rate=0; tunnels lead to valves FF, HH\n\
    \Valve HH has flow rate=22; tunnel leads to valve GG\n\
    \Valve II has flow rate=0; tunnels lead to valves AA, JJ\n\
    \Valve JJ has flow rate=21; tunnel leads to valve II\n"
:}
1651
1707

-}
module Main where

import Data.List (tails, foldl')
import Data.Map qualified as Map
import Data.IntMap (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Maybe (maybeToList)

import Advent (format)
import Advent.SmallSet (SmallSet)
import Advent.SmallSet qualified as SmallSet
import Advent.Tokenize

-- |
-- >>> :main
-- 1820
-- 2602
main :: IO ()
IO ()
main =
 do (Int
aa, [(Int, Int, [Int])]
input) <-
      [(String, Int, [String])] -> (Int, [(Int, Int, [Int])])
renumber ([(String, Int, [String])] -> (Int, [(Int, Int, [Int])]))
-> IO [(String, Int, [String])] -> IO (Int, [(Int, Int, [Int])])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      [format|2022 16
        (Valve %s has flow rate=%u;
         tunnel(|s) lead(|s) to valve(|s) %s&(, )%n)*|]

    let distances1 :: IntMap (IntMap Int)
distances1 = [(Int, IntMap Int)] -> IntMap (IntMap Int)
forall a. [(Int, a)] -> IntMap a
IntMap.fromList [(Int
k, [(Int, Int)] -> IntMap Int
forall a. [(Int, a)] -> IntMap a
IntMap.fromList [(Int
v,Int
1) | Int
v <- [Int]
vs]) | (Int
k, Int
_, [Int]
vs) <- [(Int, Int, [Int])]
input]
    let distances :: IntMap (IntMap Int)
distances  = [Int] -> IntMap (IntMap Int) -> IntMap (IntMap Int)
fw (IntMap (IntMap Int) -> [Int]
forall a. IntMap a -> [Int]
IntMap.keys IntMap (IntMap Int)
distances1) IntMap (IntMap Int)
distances1
    let flows :: IntMap Int
flows      = [(Int, Int)] -> IntMap Int
forall a. [(Int, a)] -> IntMap a
IntMap.fromList [(Int
k, Int
n) | (Int
k, Int
n, [Int]
_) <- [(Int, Int, [Int])]
input, Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0]
    let graph :: Edges
graph      = Int -> IntMap [(Int, Int, Int)] -> Edges
buildEdges Int
aa (IntMap [(Int, Int, Int)] -> Edges)
-> IntMap [(Int, Int, Int)] -> Edges
forall a b. (a -> b) -> a -> b
$
                     ([(Int, Int, Int)] -> [(Int, Int, Int)] -> [(Int, Int, Int)])
-> [(Int, [(Int, Int, Int)])] -> IntMap [(Int, Int, Int)]
forall a. (a -> a -> a) -> [(Int, a)] -> IntMap a
IntMap.fromListWith [(Int, Int, Int)] -> [(Int, Int, Int)] -> [(Int, Int, Int)]
forall a. [a] -> [a] -> [a]
(++)
                        [(Int
src, [(Int
dst,Int
costInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1,Int
flow)])
                            | (Int
src,IntMap Int
m) <- IntMap (IntMap Int) -> [(Int, IntMap Int)]
forall a. IntMap a -> [(Int, a)]
IntMap.assocs IntMap (IntMap Int)
distances
                            , Int
src Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
aa Bool -> Bool -> Bool
|| Int -> IntMap Int -> Bool
forall a. Int -> IntMap a -> Bool
IntMap.member Int
src IntMap Int
flows
                            , (Int
dst,Int
cost) <- IntMap Int -> [(Int, Int)]
forall a. IntMap a -> [(Int, a)]
IntMap.assocs IntMap Int
m
                            , Int
src Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
dst
                            , Int
flow <- Maybe Int -> [Int]
forall a. Maybe a -> [a]
maybeToList (Int -> IntMap Int -> Maybe Int
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
dst IntMap Int
flows)]

    let routeValues1 :: IntMap Int
routeValues1 = Edges -> Int -> IntMap Int
solve Edges
graph Int
30
    Int -> IO ()
forall a. Show a => a -> IO ()
print (IntMap Int -> Int
forall a. Ord a => IntMap a -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum IntMap Int
routeValues1)

    let routeValues2 :: IntMap Int
routeValues2 = Edges -> Int -> IntMap Int
solve Edges
graph Int
26
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Int
v1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
v2
        | (Int
open1,Int
v1) : [(Int, Int)]
elephants <- [(Int, Int)] -> [[(Int, Int)]]
forall a. [a] -> [[a]]
tails (IntMap Int -> [(Int, Int)]
forall a. IntMap a -> [(Int, a)]
IntMap.assocs IntMap Int
routeValues2)
        , (Int
open2,Int
v2) <- [(Int, Int)]
elephants
        , SmallSet -> SmallSet -> Bool
SmallSet.disjoint (Word64 -> SmallSet
SmallSet.SmallSet (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
open1))
                            (Word64 -> SmallSet
SmallSet.SmallSet (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
open2))])

-- | Find the maximum water flow achievable from activating all possible combinations
-- of valves.
solve ::
    Edges      {- graph: source to (dest, distance, flow) -} ->
    Int        {- starting time -} ->
    IntMap Int {- map of opened valves to maximum flow -}
solve :: Edges -> Int -> IntMap Int
solve Edges
start Int
time0 = (Int -> Int -> Int) -> [(Int, Int)] -> IntMap Int
forall a. (a -> a -> a) -> [(Int, a)] -> IntMap a
IntMap.fromListWith Int -> Int -> Int
forall a. Ord a => a -> a -> a
max ([S] -> [(Int, Int)]
forall {a}. Num a => [S] -> [(a, Int)]
go [Int -> Edges -> SmallSet -> Int -> S
S Int
time0 Edges
start SmallSet
SmallSet.empty Int
0])
    where
        go :: [S] -> [(a, Int)]
go [S]
xs = [(Word64 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (SmallSet -> Word64
SmallSet.setRep SmallSet
open),Int
flow) | S Int
_ Edges
_ SmallSet
open Int
flow <- [S]
xs] [(a, Int)] -> [(a, Int)] -> [(a, Int)]
forall a. [a] -> [a] -> [a]
++ (S -> [(a, Int)]) -> [S] -> [(a, Int)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([S] -> [(a, Int)]
go ([S] -> [(a, Int)]) -> (S -> [S]) -> S -> [(a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. S -> [S]
step) [S]
xs
        step :: S -> [S]
step (S Int
t (Node [(Edges, SmallSet, Int, Int)]
graph) SmallSet
open Int
flow) =
            [Int -> Edges -> SmallSet -> Int -> S
S Int
t' Edges
graph' (SmallSet -> SmallSet -> SmallSet
SmallSet.union SmallSet
next SmallSet
open) (Int
flow Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
t' Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
valve)
                | (Edges
graph', SmallSet
next, Int
cost, Int
valve) <- [(Edges, SmallSet, Int, Int)]
graph
                , SmallSet -> SmallSet -> Bool
SmallSet.disjoint SmallSet
next SmallSet
open
                , let t' :: Int
t' = Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
cost
                , Int
t' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0]

data S = S !Int Edges !SmallSet !Int

newtype Edges = Node [(Edges, SmallSet, Int, Int)]

renumber :: [(String, Int, [String])] -> (Int, [(Int, Int, [Int])])
renumber :: [(String, Int, [String])] -> (Int, [(Int, Int, [Int])])
renumber [(String, Int, [String])]
xs = (String, [(String, Int, [String])]) -> (Int, [(Int, Int, [Int])])
forall a b. AutoToken a b => a -> b
autoTokenize (String
"AA", [(String, Int, [String])]
xs)

-- | Replace all the string names with sequentially assigned Int names to
-- speed up comparisons and enable the use of SmallSet
buildEdges :: Int -> IntMap [(Int, Int, Int)] -> Edges
buildEdges :: Int -> IntMap [(Int, Int, Int)] -> Edges
buildEdges Int
aa IntMap [(Int, Int, Int)]
graph = IntMap Edges
m IntMap Edges -> Int -> Edges
forall a. IntMap a -> Int -> a
IntMap.! Int
aa
    where
        m :: IntMap Edges
m = ([(Int, Int, Int)] -> Edges)
-> IntMap [(Int, Int, Int)] -> IntMap Edges
forall a b. (a -> b) -> IntMap a -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([(Edges, SmallSet, Int, Int)] -> Edges
Node ([(Edges, SmallSet, Int, Int)] -> Edges)
-> ([(Int, Int, Int)] -> [(Edges, SmallSet, Int, Int)])
-> [(Int, Int, Int)]
-> Edges
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Int, Int) -> (Edges, SmallSet, Int, Int))
-> [(Int, Int, Int)] -> [(Edges, SmallSet, Int, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int, Int) -> (Edges, SmallSet, Int, Int)
f) IntMap [(Int, Int, Int)]
graph
        f :: (Int, Int, Int) -> (Edges, SmallSet, Int, Int)
f (Int
n,Int
x,Int
y) = (IntMap Edges
m IntMap Edges -> Int -> Edges
forall a. IntMap a -> Int -> a
IntMap.! Int
n, Int -> SmallSet
SmallSet.singleton Int
n, Int
x, Int
y)

-- | Floyd-Warshall shortest paths
fw ::
    [Int]               {- ^ all vertices -} ->
    IntMap (IntMap Int) {- ^ distances between a pair of vertices -} ->
    IntMap (IntMap Int) {- ^ shortest distance between two vertices -}
fw :: [Int] -> IntMap (IntMap Int) -> IntMap (IntMap Int)
fw [Int]
keys = (Int -> IntMap (IntMap Int) -> IntMap (IntMap Int))
-> IntMap (IntMap Int) -> IntMap (IntMap Int)
forall {b}. (Int -> b -> b) -> b -> b
each \Int
k -> (Int -> IntMap (IntMap Int) -> IntMap (IntMap Int))
-> IntMap (IntMap Int) -> IntMap (IntMap Int)
forall {b}. (Int -> b -> b) -> b -> b
each \Int
i -> (Int -> IntMap (IntMap Int) -> IntMap (IntMap Int))
-> IntMap (IntMap Int) -> IntMap (IntMap Int)
forall {b}. (Int -> b -> b) -> b -> b
each \Int
j IntMap (IntMap Int)
dists ->
    case (Int -> Int -> IntMap (IntMap Int) -> Maybe Int
forall {b}. Int -> Int -> IntMap (IntMap b) -> Maybe b
lk Int
i Int
k IntMap (IntMap Int)
dists, Int -> Int -> IntMap (IntMap Int) -> Maybe Int
forall {b}. Int -> Int -> IntMap (IntMap b) -> Maybe b
lk Int
k Int
j IntMap (IntMap Int)
dists) of
        (Just Int
d1, Just Int
d2) ->
            (IntMap Int -> IntMap Int -> IntMap Int)
-> Int -> IntMap Int -> IntMap (IntMap Int) -> IntMap (IntMap Int)
forall a. (a -> a -> a) -> Int -> a -> IntMap a -> IntMap a
IntMap.insertWith ((Int -> Int -> Int) -> IntMap Int -> IntMap Int -> IntMap Int
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
IntMap.unionWith Int -> Int -> Int
forall a. Ord a => a -> a -> a
min) Int
i (Int -> Int -> IntMap Int
forall a. Int -> a -> IntMap a
IntMap.singleton Int
j (Int
d1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d2)) IntMap (IntMap Int)
dists
        (Maybe Int, Maybe Int)
_ -> IntMap (IntMap Int)
dists
    where
        each :: (Int -> b -> b) -> b -> b
each Int -> b -> b
g b
z = (b -> Int -> b) -> b -> [Int] -> b
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Int -> b -> b) -> b -> Int -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> b -> b
g) b
z [Int]
keys
        lk :: Int -> Int -> IntMap (IntMap b) -> Maybe b
lk Int
i Int
j IntMap (IntMap b)
m =
         do IntMap b
m' <- Int -> IntMap (IntMap b) -> Maybe (IntMap b)
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
i IntMap (IntMap b)
m
            Int -> IntMap b -> Maybe b
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
j IntMap b
m'