#!/usr/bin/env python

import numpy as np

properties = np.array([[3, 0, 0, -3, 2],
                       [-3, 3, 0, 0, 9],
                       [-1, 0, 4, 0, 1],
                       [0, 0, -2, 2, 8]])

def get_scores(count_calories=False):
    max_score = 0
    for sugar in range(100):
        for sprinkles in range(100-sugar):
            for candy in range(100-sugar-sprinkles):
                chocolate = 100-sugar-sprinkles-candy
                counts = np.array([sugar, sprinkles, candy, chocolate])
                scores = properties.T @ counts
                if count_calories and scores[-1] != 500:
                    continue
                if np.any(scores[:-1] <= 0):
                    continue
                else:
                    score = np.prod(scores[:-1])
                    max_score = max(max_score, score)
    return max_score

# part 1
print(get_scores())

# part 2
print(get_scores(True))