aoc22/day16-2.py

import queue

MAX_STEPS = 26

q = queue.PriorityQueue()
valves = {}
valvesConsider = {}
distanceMap = {}
eliminated = {}
numTotalSteps = 0

class Valve():
    def __init__(self, flowRate, tunnels, name):
        self.flowRate = flowRate
        self.tunnels = tunnels
        self.name = name

    def getFlowrate(self):
        return self.flowRate

    def genDistanceMap(self):
        nextValves = queue.PriorityQueue()
        nextValves.put((0, self))

        distanceMap[(self.name, self.name)] = 0

        while not nextValves.empty():
            curDist, valve = nextValves.get()

            for tunnel in valve.tunnels:
                if (self.name, tunnel) in distanceMap:
                    continue

                nextValve = valves[tunnel]
                nextDist = curDist + 1

                distanceMap[(self.name, nextValve.name)] = nextDist
                nextValves.put((nextDist, nextValve))

    def __lt__(self, other):
        return self.flowRate < other.flowRate and self.name < other.name


def maxRemaining(remainingTime1, remainingTime2, pos1, pos2, possibleValves):
    
    volves = possibleValves.copy()
    possibleRemaining = 0
    done = False
    newpos = pos1
    besttime = remainingTime1

    while not done:
        best = 0
        done = True
        for valve in volves:
            if remainingTime1 - distanceMap[pos1, valve] - 1 > 0:
                press = (remainingTime1 - distanceMap[pos1, valve] - 1) * possibleValves[valve]
                if press > best:
                    done = False
                    best = press
                    newpos = valve
                    besttime = remainingTime1 - distanceMap[pos1, valve] - 1
        if not done:
            volves.pop(newpos)
        possibleRemaining += best
        remainingTime1 = besttime
        pos1 = newpos
    done = False
    besttime = remainingTime2
    while not done:
        best = 0
        done = True
        for valve in volves:
            if remainingTime2 - distanceMap[pos2, valve] - 1 > 0:
                press = (remainingTime2 - distanceMap[pos2, valve] - 1) * possibleValves[valve]
                if press > best:
                    done = False
                    best = press
                    newpos = valve
                    besttime = remainingTime2 - distanceMap[pos2, valve] - 1
        if not done:
            volves.pop(newpos)
        possibleRemaining += best
        remainingTime1 = besttime
        pos2 = newpos
    
    # nextSteps = [2 * x for x in range(int(remainingTime / 2) + 1)]

    # allValveValues = list(reversed([v for _, v in possibleValves.items()]))
    # possibleRemaining = 0

    # for timeOffset in nextSteps:
    #     if timeOffset > remainingTime:
    #         break

    #     if len(allValveValues) == 0:
    #         break

    #     openFor = remainingTime - timeOffset
    #     possibleRemaining += allValveValues[0] * (remainingTime - timeOffset)
    #     allValveValues = allValveValues[1:]

    #     if len(allValveValues) == 0:
    #         break

    #     possibleRemaining += allValveValues[0] * (remainingTime - timeOffset)
    #     allValveValues = allValveValues[1:]

    return possibleRemaining



def step(position1, position2, busy1, busy2, currentTime, pressure, _valvesCons):
    global curentbest
    global numTotalSteps
    global maxRemaining

    numTotalSteps += 1

    if maxRemaining(MAX_STEPS - currentTime - busy1,MAX_STEPS - currentTime - busy2, position1, position2, _valvesCons) + pressure < curentbest:
        # There's no way we can get better than the current max pressure, so bail.
        eliminated.setdefault('maxRemaining', 0)
        eliminated['maxRemaining'] += 1
        return

    done = False
    if pressure > curentbest:
        curentbest = pressure
    if len(_valvesCons) == 0:
        return
    if busy1 == 0:
        for destination in _valvesCons:
            _CopyCons = _valvesCons.copy()
            _dist = distanceMap[(position1, destination)]
            if (currentTime + _dist + 1) >= MAX_STEPS:
                continue
            _busy1 = _dist + 1
            _pressure = pressure + ((MAX_STEPS - currentTime - _busy1) * _CopyCons[destination])
            _CopyCons.pop(destination)
            done = True
            q.put((_pressure * -1, destination, position2, _busy1, busy2, currentTime, _CopyCons))
        if done:
            return
    if busy2 == 0:
        for destination in _valvesCons:
            _CopyCons = _valvesCons.copy()
            _dist = distanceMap[(position2, destination)]
            if (currentTime + _dist + 1) >= MAX_STEPS:
                continue
            _busy2 =_dist + 1
            _pressure = pressure + ((MAX_STEPS - currentTime - _busy2)* _CopyCons[destination])
            _CopyCons.pop(destination)
            done = True
            q.put((_pressure * -1, position1, destination, busy1, _busy2, currentTime, _CopyCons))
        if done:
            return
    if busy1 <= 0 and busy2 <= 0:
        return
    elif busy2 == 0:
        idle = busy1
    elif busy1 == 0: 
        idle = busy2
    else:
        idle = min(busy1, busy2)
    _currentTime = currentTime + idle
    _busy1 = busy1 - idle
    _busy2 = busy2 - idle
    _CopyCons = _valvesCons.copy()
    if _currentTime >= MAX_STEPS:
        return
    q.put((pressure * -1, position1, position2, _busy1, _busy2, _currentTime, _CopyCons))


openvalves = []


with open('input16.txt','r') as f:
    inp = f.read().splitlines(keepends=False)

for line in inp:
    tempA, tempB = line.split(' has flow rate=')
    _valve = tempA[-2:]
    tempB = tempB.replace('; tunnel leads to valve ','; tunnels lead to valves ')
    flowRate, tempA = tempB.split('; tunnels lead to valves ')
    flowRate = int(flowRate)
    tunnels = tempA.split(', ')
    if flowRate == 0:
        openvalves.append(_valve)
    else:
        valvesConsider[_valve] = flowRate
    newValve = Valve( flowRate, tunnels, _valve)
    valves[_valve] = newValve

for _, valve in valves.items():
    valve.genDistanceMap()

for i in distanceMap:
    print(i, distanceMap[i])

position = 'AA'
currentTime = 0
pressure = 0
curentbest = 0
q.put((pressure, position, position, 0, 0, currentTime, valvesConsider))

while not q.empty():
    #print(q.qsize())
    pressure, position1, position2, busy1, busy2, currentTime, _valvesConsider = q.get()
    # print(q.qsize(), pressure, position1, position2, busy1, busy2, currentTime, numTotalSteps, eliminated)
    step(position1, position2, busy1, busy2, currentTime, pressure * -1, _valvesConsider)


print(curentbest)