from collections import deque
from itertools import product
from tools.aoc import AOCDay
from typing import Any


class Valve:
    def __init__(self, name: str, flowrate: int):
        self.name: str = name
        self.flowrate: int = flowrate
        self.tunnels: set = set()


class Tunnel:
    def __init__(self, target: Valve, length: int = 1):
        self.target: Valve = target
        self.length: int = length

    def __str__(self):
        return f"Tunnel(target={self.target.name}, length={self.length})"

    def __repr__(self):
        return str(self)


def get_openable_valve_tunnels(valve: Valve, open_valves: set, time_remaining: int) -> set:
    tunnels = set()
    queue = deque()
    visited = set()
    queue.append((0, valve))
    while queue:
        d, v = queue.popleft()
        if v.name in visited:
            continue
        visited.add(v.name)
        if v.name not in open_valves and d + 2 <= time_remaining:
            tunnels.add((d, v))
        for x in v.tunnels:
            queue.append((d + x.length, x.target))

    return tunnels


def get_max_flow(valve: Valve, open_valves: list, time_remaining: int = 30, depth: int = 0) -> int:
    max_flow = 0

    ov = {t for t in valve.tunnels if t.target not in open_valves and t.length < time_remaining - 2}
    if time_remaining <= 0 or not ov:
        return 0

    for tunnel in ov:
        this_open_flow = tunnel.target.flowrate * (time_remaining - tunnel.length - 1)
        this_flow = get_max_flow(tunnel.target, open_valves + [tunnel.target], time_remaining - tunnel.length - 1, depth + 1)
        if this_flow + this_open_flow > max_flow:
            max_flow = this_flow + this_open_flow


    return max_flow


def get_max_flow_double(valve1: Valve, valve2: Valve, open_valves: list, DP: dict, time_remaining_1: int = 26, time_remaining_2: int = 26, depth: int = 0) -> int:
    dp_key = valve1.name + valve2.name + "%02d" % time_remaining_1 + "%02d" % time_remaining_2 + "".join(open_valves)
    dp_key2 = valve2.name + valve1.name + "%02d" % time_remaining_1 + "%02d" % time_remaining_2 + "".join(open_valves)
    if dp_key in DP:
        return DP[dp_key]

    if time_remaining_1 <= 0 and time_remaining_2 <= 0:
        return 0

    ov1 = {t for t in valve1.tunnels if t.target.name not in open_valves and t.length < time_remaining_1 - 2}
    ov2 = {t for t in valve2.tunnels if t.target.name not in open_valves and t.length < time_remaining_2 - 2}

    if not ov1 and not ov2:
        return 0

    if not ov1:
        ov1 = {Tunnel(valve1, 99)}

    if not ov2:
        ov2 = {Tunnel(valve2, 99)}

    permut = product(ov1, ov2)

    max_flow = 0
    for v1, v2 in permut:
        if v1.target.name == v2.target.name:
            continue

        this_open_flow = 0

        t1 = valve1
        if v1.length + 2 <= time_remaining_1:
            this_open_flow += v1.target.flowrate * (time_remaining_1 - v1.length - 1)
            t1 = v1.target

        t2 = valve2
        if v2.length + 2 <= time_remaining_2:
            this_open_flow += v2.target.flowrate * (time_remaining_2 - v2.length - 1)
            t2 = v2.target

        this_flow_rate = get_max_flow_double(
            t1,
            t2,
            sorted(open_valves + [v1.target.name, v2.target.name]),
            DP,
            time_remaining_1 - v1.length - 1,
            time_remaining_2 - v2.length - 1,
            depth + 1
        )

        if this_flow_rate + this_open_flow > max_flow:
            max_flow = this_flow_rate + this_open_flow

    DP[dp_key] = max_flow
    DP[dp_key2] = max_flow
    return max_flow


class Day(AOCDay):
    inputs = [
        [
            (1651, "input16_test"),
            (1850, "input16"),
        ],
        [
            (1707, "input16_test"),
            (2306, "input16"),
        ]
    ]

    def get_valve_graph(self) -> Valve:
        valves = {}
        tmp_tunnels = {}
        for line in self.getInput():
            p = line.split(" ")
            valve_name = p[1]
            flowrate = int(p[4][5:-1])
            tunnels = "".join(p[9:]).split(",")
            valves[valve_name] = Valve(valve_name, flowrate)
            tmp_tunnels[valve_name] = tunnels

        for name, tunnels in tmp_tunnels.items():
            valves[name].tunnels = {Tunnel(valves[t]) for t in tunnels}

        for valve in valves.values():
            tunnels = set()
            queue = deque()
            visited = set()
            queue.append((0, valve))
            while queue:
                d, v = queue.popleft()
                if v in visited:
                    continue
                visited.add(v)
                if v != valve and v.flowrate > 0:
                    tunnels.add(Tunnel(v, d))
                for x in v.tunnels:
                    queue.append((d + 1, x.target))

            tmp_tunnels[valve.name] = tunnels

        for name, tunnels in tmp_tunnels.items():
            valves[name].tunnels = tunnels

        return valves["AA"]

    def part1(self) -> Any:
        return get_max_flow(self.get_valve_graph(), [])

    def part2(self) -> Any:
        root = self.get_valve_graph()
        return get_max_flow_double(root, root, [], {})


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