import numpy as np
from cpmpy import *
np.set_printoptions(edgeitems=30, linewidth=100000)

# Read input
fileObj = open('input', "r")
words = fileObj.read().splitlines()
fileObj.close()
blueprints = []
for w in words:
    spl = w.split(' ')
    blueprints.append(list(map(int, [spl[6], spl[12], spl[18], spl[21], spl[27], spl[30]])))

ORE = 0
CLAY = 1
OBSIDIAN = 2
GEODE = 3

NR_OF_STEPS = 32  # 24 for part 1


qualityLevels = []
for index in range(3):  # range(len(blueprints) for part 1
    bp = blueprints[index]
    a, b, c, d, e, f = bp
    costArray = cpm_array([
        [0, 0, 0, 0],
        [a, 0, 0, 0],
        [b, 0, 0, 0],
        [c, d, 0, 0],
        [e, 0, f, 0],
    ])
    # Lengths are NR_OF_STEPS+1 so we can have starting values at Minute 0
    botsPerStep = intvar(0, 9999, shape=(4, NR_OF_STEPS + 1), name="bots")  # Array of amount of bots per step per type
    resourcesPerStep = intvar(0, 9999, shape=(4, NR_OF_STEPS + 1), name="resources")  # Array of amount of resources per step per type

    building = intvar(-1, 3, shape=(1, NR_OF_STEPS+1), name="building")  # -1 = not building anything, bigger x means "building a robot that can harvest resource x"
    blueprintModel = Model(
    )
    for r in range(4):  # For every resource type...
        blueprintModel += (botsPerStep[r, 0] == (1 if r == ORE else 0))  # Start with 1 ore robot
        blueprintModel += (resourcesPerStep[r, 0] == 0)  # Start with 0 resources

        for s in range(1, NR_OF_STEPS+1):  # For every step...
            blueprintModel += (resourcesPerStep[r, s] == resourcesPerStep[r, s-1] >= costArray[building[0, s] + 1, r])  # Don't spend resources before you have them
            blueprintModel += (resourcesPerStep[r, s] == resourcesPerStep[r, s-1] + botsPerStep[r, s-1] - costArray[building[0, s]+1, r])  # Bots of type r make 1 resource per step, subtract resources you spent on robots
            blueprintModel += (botsPerStep[r, s] == (botsPerStep[r, s-1]) + (building[0, s] == r))  # Keep your robots, + 1 if you built one

    blueprintModel.maximize(resourcesPerStep[3][NR_OF_STEPS])  # Objective: maximize the number of geodes in the end
    if blueprintModel.solve():
        print("solved blueprint %s with max geodes %s" % ((index+1), resourcesPerStep[3][NR_OF_STEPS].value()))
        # print(costArray)
        # print(bots, resources and building per step)
        # print(np.row_stack((botsPerStep.value(),resourcesPerStep.value(), building.value())))
        qualityLevels.append((index+1) * resourcesPerStep[3][NR_OF_STEPS].value())
    else:
        print("No solution found")
print(qualityLevels)
print(sum(qualityLevels))