"""Just compiling with numba."""
import numba
import numpy as np


@numba.njit(numba.void(numba.int64[:,:]))
def apply_gravity(terrain):
    """
    $ python -m pyperf timeit --fast -s 'from examples.sandpile_numpy import main' 'main(10_000, False)'
    Mean +- std dev: 19.2 ms +- 0.3 ms
    $ python -m pyperf timeit --fast -s 'from examples.sandpile_numpy import main' 'main(100_000, False)'
    Mean +- std dev: 1.83 sec +- 0.04 sec

    2**16 grains : 1.6s
    2**17 grains : 3.9s
    2**18 grains : 13s
    2**19 grains : 51s
    2**20 grains : 3m5
    """
    shape = np.shape(terrain)
    while True:
        did_someting = False
        for x, y in np.ndindex(shape):
            if terrain[x, y] >= 4:
                div, terrain[x, y] = divmod(terrain[x][y], 4)
                terrain[x - 1][y] += div
                terrain[x + 1][y] += div
                terrain[x][y + 1] += div
                terrain[x][y - 1] += div
                did_someting = True

        if not did_someting:
            return


def main(height, show=True):
    width = int(height**0.5) + 1
    terrain = np.zeros((width, width), dtype=np.int64)
    terrain[width // 2, width // 2] = height
    apply_gravity(terrain)
    if show:
        import sys
        np.set_printoptions(threshold=sys.maxsize, linewidth=999)
        print(terrain)