import heapq
import itertools
import re
from tools.aoc import AOCDay
from typing import Any


class Valve:
    def __init__(self, name: str, flowrate: int) -> None:
        self.__name = name
        self.__flowrate = flowrate
        self.neighbours = {}

    @property
    def flowrate(self) -> int:
        return self.__flowrate

    @property
    def name(self) -> str:
        return self.__name

    def shortest_path_to(self, other: 'Valve'):
        if other in self.neighbours:
            return self.neighbours[other]

        current = self
        v = set()
        v.add(current)
        q = [(v, k) for k, v in self.neighbours.items()]
        heapq.heapify(q)

        while q:
            dist, current = heapq.heappop(q)
            if current == other:
                return dist

            if current in v:
                continue

            v.add(current)
            for n, d in current.neighbours.items():
                heapq.heappush(q, (dist + d, n))

    def __lt__(self, other: 'Valve') -> bool:
        return self.__name < other.__name

    def __repr__(self) -> str:
        return "Valve(%s;%s)" % (self.__name, self.__flowrate)

    def __str__(self) -> str:
        return self.__repr__()


def get_most_pressure_release_solo(root: Valve, remaining_minutes: int = 30, visited: set = None) -> int:
    if visited is None:
        visited = set()
    visited.add(root)

    my_flowrate = 0
    if root.flowrate > 0:
        remaining_minutes -= 1
        my_flowrate = root.flowrate * remaining_minutes

    max_flowrate = 0
    for n, d in root.neighbours.items():
        if n in visited:
            continue

        if remaining_minutes <= d + 1:
            continue

        n_flowrate = get_most_pressure_release_solo(n, remaining_minutes - d, visited.copy())
        if n_flowrate > max_flowrate:
            max_flowrate = n_flowrate

    return max_flowrate + my_flowrate


def get_mode_pressure_release_double(root_me: Valve, root_ele: Valve, r_min_me: int = 26, r_min_ele: int = 26, visited: set = None) -> int:
    dp_key = root_me.name + root_ele.name + "%02d" % r_min_me + "%02d" % r_min_ele + "".join(v.name for v in sorted(visited))
    dp_key2 = root_ele.name + root_me.name + "%02d" % r_min_me + "%02d" % r_min_ele + "".join(v.name for v in sorted(visited))

    if dp_key in DP:
        return DP[dp_key]

    my_flowrate = 0
    if r_min_me > 1 and root_me not in visited:
        visited.add(root_me)
        r_min_me -= 1
        my_flowrate += root_me.flowrate * r_min_me

    if r_min_ele > 1 and root_ele not in visited:
        visited.add(root_ele)
        r_min_ele -= 1
        my_flowrate += root_ele.flowrate * r_min_ele

    if not my_flowrate and root_me.name != 'AA':
        return 0

    max_flowrate = 0
    rem_list = [x for x in root_me.neighbours.keys() if x not in visited]
    if len(rem_list) == 1:
        if root_me.neighbours[rem_list[0]] < root_ele.neighbours[rem_list[0]]:
            max_flowrate = get_mode_pressure_release_double(rem_list[0], root_ele, r_min_me - root_me.neighbours[rem_list[0]], r_min_ele, visited.copy())
        else:
            max_flowrate = get_mode_pressure_release_double(root_me, rem_list[0], r_min_me, r_min_ele - root_ele.neighbours[rem_list[0]], visited.copy())
    else:
        max_prod = len(rem_list) ** 2
        cur_prod = 0
        for v1, v2 in itertools.product(rem_list, repeat=2):
            if len(visited) == 1:
                cur_prod += 1
                print("Iter (%d/%d)" % (cur_prod, max_prod))
            if v1 == v2:
                continue

            me, ele = v1, v2
            if r_min_me <= root_me.neighbours[me] + 1:
                me = root_me

            if r_min_ele <= root_ele.neighbours[ele] + 1:
                ele = root_ele

            if me == root_me and ele == root_ele:
                continue

            n_flowrate = get_mode_pressure_release_double(
                me, ele,
                r_min_me - (root_me.neighbours[me] if me in root_me.neighbours else 1),
                r_min_ele - (root_ele.neighbours[ele] if ele in root_ele.neighbours else 1),
                visited.copy()
            )
            if n_flowrate > max_flowrate:
                max_flowrate = n_flowrate

            if n_flowrate == 1671:
                break

    DP[dp_key] = my_flowrate + max_flowrate
    DP[dp_key2] = my_flowrate + max_flowrate
    return my_flowrate + max_flowrate


DP = {}


class Day(AOCDay):
    input_regexp = re.compile(r'Valve ([A-Z][A-Z]) has flow rate=([0-9]+); tunnels? leads? to valves? (.*)')
    inputs = [
        [
            (1651, "input16_test"),
            (1947, "input16_dennis"),
            (1850, "input16"),
        ],
        [
            (1707, "input16_test"),
            (2556, "input16_dennis"),
            (2306, "input16"),
        ]
    ]

    def parse_input(self) -> Valve:
        n_cache = {}
        valves = {}
        for line in self.getInput():
            name, flowrate, neighbours = self.input_regexp.match(line).groups()
            valves[name] = Valve(name, int(flowrate))
            n_cache[name] = neighbours.split(", ")

        for valve, neighbours in n_cache.items():
            for n in neighbours:
                valves[valve].neighbours[valves[n]] = 1

        for valve in valves.values():
            if valve.flowrate != 0:
                continue

            for n_set in valve.neighbours.keys():
                del n_set.neighbours[valve]
                for n_get in valve.neighbours.keys():
                    if n_get == n_set:
                        continue

                    if n_get in n_set.neighbours:
                        continue

                    n_set.neighbours[n_get] = valve.neighbours[n_set] + valve.neighbours[n_get]

        for v_set in valves.values():
            for v_get in valves.values():
                if v_set == v_get or v_get in v_set.neighbours or v_get.flowrate == 0:
                    continue

                v_set.neighbours[v_get] = v_set.shortest_path_to(v_get)

        for n in list(valves['AA'].neighbours.keys()):
            if n.flowrate == 0:
                del valves['AA'].neighbours[n]

        return valves['AA']

    def part1(self) -> Any:
        return get_most_pressure_release_solo(self.parse_input())

    def part2(self) -> Any:
        global DP
        DP = {}
        root_valve = self.parse_input()
        return get_mode_pressure_release_double(root_valve, root_valve, visited={root_valve})


if __name__ == '__main__':
    day = Day(2022, 16)
    day.run(verbose=True)