from collections import defaultdict
from enum import Enum
from heapq import heappush, heappop
from tools.aoc import AOCDay
from typing import Any


class Material(Enum):
    ORE = 1
    CLAY = 2
    OBSIDIAN = 3
    GEODE = 4


class Robot:
    def __init__(self, produces: Material, costs: dict, min_viable: int = 0):
        self.produces = produces
        self.costs = costs
        self.min_viable = min_viable


class Inventory:
    def __init__(self):
        self.robots = {
            Material.ORE: 1,
            Material.CLAY: 0,
            Material.OBSIDIAN: 0,
            Material.GEODE: 0,
        }
        self.materials = {
            Material.ORE: 0,
            Material.CLAY: 0,
            Material.OBSIDIAN: 0,
            Material.GEODE: 0,
        }
        self.__build_queue = []

    def gather(self) -> None:
        for mat, amount in self.robots.items():
            self.materials[mat] += amount

    def build_robot(self, robot: Robot) -> None:
        for cost_mat, cost_amount in robot.costs.items():
            self.materials[cost_mat] -= cost_amount
        self.robots[robot.produces] += 1

    def copy(self) -> 'Inventory':
        new_inv = Inventory()
        new_inv.robots = self.robots.copy()
        new_inv.materials = self.materials.copy()
        return new_inv

    def value(self) -> int:
        return self.robots[Material.ORE] \
            + 2 * self.robots[Material.CLAY] \
            + 4 * self.robots[Material.OBSIDIAN] \
            + 8 * self.robots[Material.GEODE]

    def __lt__(self, other: 'Inventory') -> bool:
        return self.materials[Material.GEODE] < other.materials[Material.GEODE]


class Blueprint:
    def __init__(self, bp_id: int, robots: dict):
        self.bp_id = bp_id
        self.robots = robots
        self.max_cost = {
            Material.ORE: max(robot.costs[Material.ORE] for _, robot in self.robots.items()),
            Material.CLAY: self.robots[Material.OBSIDIAN].costs[Material.CLAY],
            Material.OBSIDIAN: self.robots[Material.GEODE].costs[Material.OBSIDIAN],
            Material.GEODE: 1e9
        }


def pass_minute(bp: Blueprint, inv: Inventory, ignore: list) -> (Inventory, list):
    could_build = []
    for mat in Material:
        if inv.robots[mat] >= bp.max_cost[mat]:
            continue
        for robot_mat, robot_cost in bp.robots[mat].costs.items():
            if inv.materials[robot_mat] < robot_cost:
                break
        else:
            if bp.robots[mat] not in ignore:
                could_build.append(bp.robots[mat])

    inv.gather()

    return inv, could_build


def get_dfs_quality(bp: Blueprint, max_time: int = 24) -> int:
    q = []
    max_geode = 0
    max_value = defaultdict(int)
    heappush(q, (max_time, Inventory(), [], 0))
    q_count = 0
    while q:
        q_count += 1
        time_left, inv, ignore, no_build = heappop(q)
        inv, could_build = pass_minute(bp, inv, ignore)
        time_left -= 1
        inv_value = inv.value()
        if inv_value < max_value[time_left] // 2.1:
            continue
        elif inv_value > max_value[time_left]:
            max_value[time_left] = inv_value

        if time_left == 0:
            if inv.materials[Material.GEODE] > max_geode:
                max_geode = inv.materials[Material.GEODE]
            continue

        if bp.robots[Material.GEODE] in could_build:
            inv.build_robot(bp.robots[Material.GEODE])
            heappush(q, (time_left, inv, [], 0))
            continue

        for robot in could_build:
            if robot.min_viable > time_left:
                continue
            sub_inv = inv.copy()
            sub_inv.build_robot(robot)
            heappush(q, (time_left, sub_inv, [], 0))

        if no_build < bp.max_cost[Material.ORE]:
            heappush(q, (time_left, inv, could_build, no_build + 1))

    print(q_count)
    return max_geode


class Day(AOCDay):
    inputs = [
        [
            (33, "input19_test"),
            (1616, "input19"),
        ],
        [
            (3472, "input19_test"),
            (8990, "input19"),
        ]
    ]

    def get_blueprints(self) -> list:
        bp = []
        for line in self.getInput():
            parts = line.split(" ")
            blueprint_id = int(parts[1][:-1])
            ore_ore_cost = int(parts[6])
            clay_ore_cost = int(parts[12])
            obsi_ore_cost = int(parts[18])
            obsi_clay_cost = int(parts[21])
            geode_ore_cost = int(parts[27])
            geode_obsi_cost = int(parts[30])
            robots = {
                Material.ORE: Robot(Material.ORE, {Material.ORE: ore_ore_cost}, 7),
                Material.CLAY: Robot(Material.CLAY, {Material.ORE: clay_ore_cost}, 5),
                Material.OBSIDIAN: Robot(Material.OBSIDIAN, {Material.ORE: obsi_ore_cost, Material.CLAY: obsi_clay_cost}, 3),
                Material.GEODE: Robot(Material.GEODE, {Material.ORE: geode_ore_cost, Material.OBSIDIAN: geode_obsi_cost}, 1)
            }
            bp.append(Blueprint(blueprint_id, robots))

        return bp

    def part1(self) -> Any:
        blueprints = self.get_blueprints()
        score = 0
        for b in blueprints:
            quality = get_dfs_quality(b)
            score += quality * b.bp_id
            print("BP", b.bp_id, "Q", quality, "S", score)
        return score

    def part2(self) -> Any:
        blueprints = self.get_blueprints()
        score = 1
        for b in blueprints[:3]:
            quality = get_dfs_quality(b, 32)
            score *= quality
            print("BP", b.bp_id, "Q", quality, "S", score)
        return score


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