advent_of_code_2022_day_16.nim

import deques, sets, tables, strscans, strutils, sequtils, strformat, threadpool
{.experimental: "parallel".}

type 
    Valve = ref object
        flowRate: int
        next: seq[string]
        paths: seq[tuple[node: string, length: int]]

const inputPattern = "Valve $w has flow rate=$i; tunnel$* lead$* to valve$* $+"

proc loadValves*(input: seq[string]): Table[string, Valve] =
    for line in input:
        let (ok, node, flowRate, _, _, _, next) = line.scanTuple(inputPattern)

        if ok:
            result[node] = Valve(flowRate: flowRate, next: next.split(", "))

proc shortestPath(graph: Table[string, Valve], start: string): Table[string, string] =
    var previous: Table[string, string]
    var visited: HashSet[string]
    var queue: Deque[tuple[node: string, distance: int]]
    queue.addLast((start, 0))
    visited.incl(start)

    while queue.len > 0:
        let (node, distance) = queue.popFirst()

        for neighbour in graph[node].next:
            if not visited.contains(neighbour):
                previous[neighbour] = node
                queue.addLast((neighbour, distance + 1))
                visited.incl(neighbour)

    previous

proc getValvePathsFor(graph: Table[string, Valve], 
                      start: string,
                      fullPath: Table[string, string]): seq[tuple[node: string, length: int]] =

    for dst in fullPath.keys.toSeq.filterIt(it != start and graph[it].flowRate > 0):
        var path =  @[dst]
        var prev = fullPath[dst]

        while prev != start:
            path.add(prev)
            prev = fullPath[prev]

        result.add((dst, path.len))

proc getValvePaths*(graph: Table[string, Valve], start: string): Table[string, Valve] =
    result = graph
    for node in graph.keys:
        result[node].paths = getValvePathsFor(graph, node, shortestPath(graph, node))

proc walkMaxOutput*(graph: Table[string, Valve], 
                    start: string, 
                    timeLeft: int, 
                    depth: int,
                    available: HashSet[string]): (int, seq[string]) =
    var available = available
    var time = timeLeft
    var current = start
    var output: int
    var flowRate: int
    var finalPath: seq[string]

    while time > 0:
        finalPath.add(current)
        
        if graph[current].flowRate > 0 and available.contains(current):
            output.inc(flowRate)
            available.excl(current)
            flowRate.inc(graph[current].flowRate)
            time.dec
        else:
            var maxOutput: int
            var maxPath: tuple[node: string, length: int]
            for path in graph[current].paths.filterIt(available.contains(it.node)):
                let (subOutput, _) = walkMaxOutput(graph, path.node, time - path.length, depth + 1, available)
                if subOutput > maxOutput:
                    maxOutput = subOutput
                    maxPath = path
            
            if maxPath.length > 0:
                time.dec(maxPath.length)
                if time >= 0:
                    output.inc(maxPath.length * flowRate)
                else:
                    output.inc((maxPath.length + time) * flowRate)
                current = maxPath.node
            else:
                time.dec
                output.inc(flowRate)

        if depth == 0:
            echo "==="
            echo fmt"Time: {timeLeft - time}"
            echo fmt"Current: {current}"
            echo fmt"FlowRate: {flowRate}"
            echo fmt"Output: {output}"

    (output, finalPath)

# taken from https://github.com/MichaelCook/cribbage/blob/master/cribbage.nim
iterator choose[T](a: openarray[T], num_choose: int): seq[T] =
  var
    chosen = newSeqOfCap[T](num_choose)
    i = 0
    i_stack = newSeqOfCap[int](num_choose)
  
  while true:
    if chosen.len == num_choose:
      yield chosen
      discard chosen.pop()
      i = i_stack.pop() + 1
    elif i != a.len:
      chosen.add(a[i])
      i_stack.add(i)
      inc i
    elif i_stack.len > 0:
      discard chosen.pop()
      i = i_stack.pop() + 1
    else:
      break

proc getCombinations(graph: Table[string, Valve]): seq[(HashSet[string], HashSet[string])] =
    let nodes = graph.keys.toSeq.filterIt(graph[it].flowRate > 0)
    let both = nodes.toHashSet

    var sets: HashSet[HashSet[HashSet[string]]]
    for length in 1..(nodes.len div 2):
        for left in choose(nodes, length):
            let left = left.toHashSet 
            let right = both - left
            sets.incl([left, right].toHashSet)
    
    for sPair in sets.toSeq:
        let pair = sPair.toSeq
        result.add((pair[0], pair[1]))
    
    echo fmt"Combinations count: {result.len}"

let graph = getValvePaths(loadValves("day-16/input.txt".lines.toSeq), "AA")

# part 1

let (output, path) = walkMaxOutput(graph, "AA", 30, 0, graph.keys.toSeq.toHashSet)

echo fmt"Max single output: {output} for path {path}"

# part 2

proc computePair(graph: Table[string, Valve], left, right: HashSet[string]): int =
    echo fmt"Checking max output for {left}..."
    let (outputLeft, _) = walkMaxOutput(graph, "AA", 26, 1, left + ["AA"].toHashSet)
    echo fmt"Checking max output for {right}..."
    let (outputRight, _) = walkMaxOutput(graph, "AA", 26, 1, right + ["AA"].toHashSet)

    result = outputLeft + outputRight

    echo fmt"Combined output for {left} and {right}: {result}"

let combinations = getCombinations(graph)
var outputs = newSeq[int](combinations.len)
parallel:
    for idx in 0..<combinations.len:
        echo fmt"Running job {idx}..."
        let (left, right) = combinations[idx]
        outputs[idx] = spawn computePair(graph, left, right)

echo fmt"Max combined output: {outputs.max}"