from queue import PriorityQueue

lines = [line.split(".") for line in open("19.input")]

blueprints = [list(map(int,[line[0].split(" ")[6], line[1].split(" ")[5], line[2].split(" ")[5], line[2].split(" ")[8], line[3].split(" ")[5], line[3].split(" ")[8]])) for line in lines]


def add(a, b):
    return tuple(a_ + b_ for (a_, b_) in zip(list(a), list(b)))

def sub(a, b):
    return tuple(a_ - b_ for (a_, b_) in zip(list(a), list(b)))

def lte(a, b):
    return all(a_ <= b_ for (a_, b_) in zip(list(a), list(b)))


def possible_purchases(res, robo, rem, costs, max_costs):
    """Returns set of possible purchase tuples"""

    result = set()

    for robo_idx, cost in enumerate(costs):
        if lte(cost, res):
            result.add(robo_idx)

    if 3 in result:
        return {3}


    for i in range(3):
        if res[i] + robo[i] * rem >= max_costs[i] * rem:
            result -= {i}

    # For part 2, this can be changed to < instead of <= for more speed. ¯\_(ツ)_/¯
    # I guess I was just lucky with my input...
    if len(result) <=  2:
        result.add(-1)

    return result


def find_max(bp_idx, blueprint, steps):
    c_ore, c_clay, c_ob_ore, c_ob_clay, c_geo_ore, c_geo_ob = blueprint

    q = PriorityQueue()

    states = {}
    results = {}
    purc_cache = {}

    res = (0, 0, 0, 0)
    robo = (1, 0, 0, 0)

    costs = [
            (c_ore, 0, 0, 0),
            (c_clay, 0, 0, 0),
            (c_ob_ore, c_ob_clay, 0, 0),
            (c_geo_ore, 0, c_geo_ob, 0),
            ]

    max_costs = [max(c[i] for c in costs) for i in range(4)]

    start = (res, robo, steps)

    q.put(start)

    i = 0

    while not q.empty():
        current = q.get()
        #print(len(states), current)
        res, robo, rem = current
        if res[0] > i:
            i = res[0]
            print(i)

        if robo[0] >= c_geo_ore and robo[2] >= c_geo_ob:
            print("Early exit")
            states[current] = res[3] + robo[3] * rem + rem * rem // 2
            continue

        if rem == 0:
            continue

        for robo_idx in possible_purchases(res, robo, rem, costs, max_costs):
            if robo_idx != -1:
                new_purc = [0, 0, 0, 0]
                new_purc[robo_idx] = 1
                next_st = (add(sub(res, costs[robo_idx]), robo), add(robo, tuple(new_purc)), rem - 1)
            else:
                next_st = (add(res, robo), robo, rem - 1)

            if next_st not in states:
                states[next_st] = next_st[0][3]
                q.put(next_st)

    return max(s for s in states.values())


def part1():
    total = 0

    for i, bp in enumerate(blueprints):
        score = find_max(i, bp, 24)
        print(i + 1, score)
        total += (i + 1) * score

    print("Part 1:", total)

part1()

def part2():
    product = 1

    for i, bp in enumerate(blueprints[:3]):
        score = find_max(i, bp, 32)
        print(i + 1, score)
        product *= score

    print("Part 2:", product)

part2()