import numba
import numpy as np
from time import perf_counter
import sys


def show_terrain(terrain, width):
    for x in range(width):
        for y in range(width):
            print(" ·●⬤"[int(terrain[x, y])], end="")
        print()


def apply_gravity(terrain):
    """
    $ python -m pyperf timeit --fast -s 'from examples.sandpile_numpy import main' 'main(10_000, False)'
    ...........
    Mean +- std dev: 799 ms +- 42 ms
    """
    while True:
        tops = terrain > 3
        if not np.any(tops):
            return
        terrain[tops] -= 4
        terrain[1:, :][tops[:-1, :]] += 1
        terrain[:-1, :][tops[1:, :]] += 1
        terrain[:, 1:][tops[:, :-1]] += 1
        terrain[:, :-1][tops[:, 1:]] += 1


def apply_gravity(terrain):
    """
    $ python -m pyperf timeit --fast -s 'from examples.sandpile_numpy import main' 'main(10_000, False)'
    ...........
    Mean +- std dev: 356 ms +- 35 ms
    """
    while True:
        tumbled, terrain = np.divmod(terrain, 4)
        if not np.any(tumbled):
            return terrain
        terrain[1:, :] += tumbled[:-1, :]
        terrain[:-1, :] += tumbled[1:, :]
        terrain[:, 1:] += tumbled[:, :-1]
        terrain[:, :-1] += tumbled[:, 1:]


@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
    """
    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

# @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: 100 ms +- 5 ms
#     """
#     shape = np.shape(terrain)
#     while True:
#         did_someting = False
#         for x, y in np.ndindex(shape):
#             if terrain[x, y] >= 4:
#                 terrain[x,y] -= 4
#                 terrain[x - 1][y] += 1
#                 terrain[x + 1][y] += 1
#                 terrain[x][y + 1] += 1
#                 terrain[x][y - 1] += 1
#                 did_someting = True
#
#         if not did_someting:
#             return

# @numba.njit(numba.void(numba.int64[:,:]))
# def apply_gravity(terrain):
#     """Can handle 10k sand grains in 1.5s."""
#     shape = np.shape(terrain)
#     while True:
#         did_someting = False
#         for x, y in np.ndindex(shape):
#             if terrain[x, y] >= 4000:
#                 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
#             elif terrain[x, y] >= 4:
#                 terrain[x,y] -= 4
#                 terrain[x - 1][y] += 1
#                 terrain[x + 1][y] += 1
#                 terrain[x][y + 1] += 1
#                 terrain[x][y - 1] += 1
#                 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:
        show_terrain(terrain, width)


if __name__ == "__main__":
    main(int(sys.argv[1]))
