from functools import lru_cache
from collections import defaultdict


CUBE_THRESHOLD = 10000

def base_cubes(xlen, ylen, zlen) -> set:
    r = set()
    for x in range(xlen):
        for y in range(ylen):
            for z in range(zlen):
                r.add((x, y, z))
    return r


@lru_cache
def neighbors(x, y, z) -> list:
    dirs = [
        (1, 0, 0),
        (-1, 0, 0),
        (0, 1, 0),
        (0, -1, 0),
        (0, 0, 1),
        (0, 0, -1),
    ]
    return [(x + dx, y + dy, z + dz) for dx, dy, dz in dirs]


@lru_cache
def count_cubes_per_layer(x, y, z) -> list:
    """ First iteration. Naively add outer layers. """
    r = []
    clayer = base_cubes(x, y, z)
    players = set(clayer)
    while True:
        nlayer = set()
        for c in clayer:
            for nb in neighbors(*c):
                if nb not in players:
                    nlayer.add(nb)
        clayer = nlayer
        players |= clayer
        layer_count = len(nlayer)
        r.append(layer_count)
        if layer_count > CUBE_THRESHOLD:
            break
    return r


def count_cubes_per_layer2(x, y, z) -> list:
    """ Second iteration. Realized that the delta between the layers grows by
    the same value. """
    r = []
    clayer = base_cubes(x, y, z)
    players = set(clayer)
    layer_0 = 0
    layer_1 = 0
    for i in range(2):
        nlayer = set()
        for c in clayer:
            for nb in neighbors(*c):
                if nb not in players:
                    nlayer.add(nb)
        clayer = nlayer
        players |= clayer
        layer_count = len(nlayer)
        if i == 0:
            layer_0 = layer_count
        elif i == 1:
            layer_1 = layer_count
        r.append(layer_count)
        if layer_count > CUBE_THRESHOLD:
            break

    layer = r[-1]
    layer_delta = (layer_1 - layer_0) + 8
    while layer <= CUBE_THRESHOLD:
        layer += layer_delta
        r.append(layer)
        layer_delta += 8
    return r


def count_cubes_per_layer3(x, y, z) -> list:
    """ Third iteration. Calculate first two layers directly, and then use
    constant delta approach from second iteration. """
    r = []
    layer_0 = 2 * (x * y + x * z + y * z)
    layer_1 = layer_0 + 4 * (x + y + z)
    r = [layer_0, layer_1]
    layer = r[-1]
    layer_delta = (layer_1 - layer_0) + 8
    while layer <= CUBE_THRESHOLD:
        layer += layer_delta
        r.append(layer)
        layer_delta += 8
    return r


def euler_126():
    # We assume that the least n for which C(n) is equal to `target`
    # is under `layer_min_threshold`. We got this my trial and error.
    layer_min_threshold = 20000
    target = 1000

    layers = defaultdict(int)
    for x in range(1, layer_min_threshold):
        for y in range(1, x + 1):

            # If we exceed the threshold, cut the loop short.
            if 2 * x * y  > layer_min_threshold:
                break

            for z in range(1, y + 1):
                layer_0 = 2 * (x * y + x * z + y * z)
                layer_1 = layer_0 + 4 * (x + y + z)

                layers[layer_0] += 1
                layers[layer_1] += 1

                layer = layer_1
                layer_delta = (layer_1 - layer_0) + 8

                # If we exceed the threshold, cut the loop short.
                while layer < layer_min_threshold:
                    layer += layer_delta
                    layer_delta += 8
                    layers[layer] += 1

    # Find the actual least value n for which C(n) == target.
    layermin = 10**9
    for n in layers.keys():
        if layers[n] == target and n < layermin:
            layermin = n
    return layermin


if __name__ == "__main__":
    solution = euler_126()
    print("e126.py: " + str(solution))
    assert(solution == 18522)