from lib import *

EXAMPLE = """
Blueprint 1:
  Each ore robot costs 4 ore.
  Each clay robot costs 2 ore.
  Each obsidian robot costs 3 ore and 14 clay.
  Each geode robot costs 2 ore and 7 obsidian.

Blueprint 2:
  Each ore robot costs 2 ore.
  Each clay robot costs 3 ore.
  Each obsidian robot costs 3 ore and 8 clay.
  Each geode robot costs 3 ore and 12 obsidian.
"""

def solve(input: Input, second=False):
    if not second:
        res = 0
    else:
        res = 1
    bps = []
    for line in input.text.replace(".", ".\n").replace(":", ":\n").splitlines():
        if "Blueprint" in line:
            bps.append([])
        elif "ore robot" in line:
            cost = str_to_ints(line)
            cost += [0, 0]
            bps[-1].append(cost)
        elif "clay robot" in line:
            cost = str_to_ints(line)
            cost += [0, 0]
            bps[-1].append(cost)
        elif "obsidian robot" in line:
            cost = str_to_ints(line)
            cost += [0,]
            bps[-1].append(cost)
        elif "geode robot" in line:
            cost = str_to_ints(line)
            cost.insert(1, 0)
            bps[-1].append(cost)

    if second:
        bps = bps[:3]
        time = 32
    else:
        time = 24

    end_states = []
    for i, bp in enumerate(bps):
        # ((ore bots, clay bots, obs bots, geo bots), (ore, clay, obs, geo))
        start = ((1, 0, 0, 0), (0, 0, 0, 0))
        states = [start]
        seen = set(states)

        for _ in range(time):
            new_states = []
            while states:
                bots, ress = states.pop()

                add_ress = [0, 0, 0, 0]
                for boti, count in enumerate(bots):
                    add_ress[boti] += count

                all_built = True
                for boti, cost in enumerate(bp):
                    if ress[0] >= cost[0] and ress[1] >= cost[1] and ress[2] >= cost[2]:
                        new_ress = list(ress)
                        new_ress[0] -= cost[0]
                        new_ress[1] -= cost[1]
                        new_ress[2] -= cost[2]
                        new_ress = tuple(map(sum, zip(new_ress, add_ress)))
                        new_bots = list(bots)
                        new_bots[boti] += 1
                        new_state = (tuple(new_bots), new_ress)
                        if not new_state in seen:
                            new_states.append(new_state)
                            seen.add(new_state)
                    else:
                        all_built = False

                # XXX: our search space is too large here it is possible to
                # optimze by not storing reduntant paths (paths where we acrue
                # more of a resource than we need), but I don't know how to
                # make it more efficient right now.
                if not all_built:
                    new_ress = tuple(map(sum, zip(ress, add_ress)))
                    new_state = (bots, new_ress)
                    if not new_state in seen:
                        new_states.append(new_state)
                        seen.add(new_state)

            # prune to keep search space "reasonable"
            states = sorted(new_states, key=lambda s: list(reversed(s[0])), reverse=True)[:100000]

        if not second:
            r = max(states, key=lambda s: s[1][3])
            q = (i + 1) * r[1][3]
            # print(i + 1, r, q)
            res += q
        else:
            r = max(states, key=lambda r: r[1][3])
            res *= r[1][3]
    return res

def main():
    DAY_INPUT = "d19.txt"
    e1 = solve(Input(EXAMPLE))
    print("Example  1:", e1)
    assert e1 == 33
    print("Solution 1:", solve(Input(DAY_INPUT)))
    print("Example  2:", solve(Input(EXAMPLE), True))
    print("Solution 2:", solve(Input(DAY_INPUT), True))
    return

if __name__ == "__main__":
    main()