from orrinjelo.utils.decorators import timeit
import numpy as np

def parse(lines):
    return np.array([[int(c) for c in line.strip()] for line in lines])

visited = []
def flash(a, x, y):
    global visited
    if (x,y) in visited:
        return
    for dx in range(-1,2):
        for dy in range(-1,2):
            if dx == 0 and dy == 0:
                continue
            if x+dx < 0 or x+dx >= a.shape[0]:
                continue
            if y+dy < 0 or y+dy >= a.shape[1]:
                continue
            a[x+dx, y+dy] += 1
            visited.append((x,y))
            if a[x+dx, y+dy] > 9:
                flash(a, x+dx, y+dy)

def progress(a):
    global visited
    a += 1

    x,y = np.where(a > 9)
    visited = []
    for i in range(len(x)):
        flash(a,x[i],y[i])

    count = np.sum(a > 9)
    # print('a:\n', a)

    a[a > 9] = 0

    return a, count

@timeit("Day 11 Part 1")
def part1(input_str, use_rust=False):
    octomap = parse(input_str)
    total_count = 0

    for i in range(100):
        octomap, count = progress(octomap)
        total_count += count

    return total_count

@timeit("Day 11 Part 2")
def part2(input_str, use_rust=False):
    octomap = parse(input_str)
    step = 0
    while True:
        step += 1
        octomap, count = progress(octomap)
        if count == octomap.shape[0]*octomap.shape[1]:
            break

    return step


# = Test ================================================

inputlist = [
    '5483143223',
    '2745854711',
    '5264556173',
    '6141336146',
    '6357385478',
    '4167524645',
    '2176841721',
    '6882881134',
    '4846848554',
    '5283751526',
]

def test_part1():
    # import matplotlib.pyplot as plt
    # plt.imshow(parse(inputlist))
    # plt.show()
    
    assert part1(inputlist) == 1656


def test_part2():
    assert part2(inputlist) == 195

import pygame
import sys


def plot(input_str):
    # octomap = parse(input_str)
    octomap = np.random.randint(0,9,(100,100))

    pygame.init()
    clock = pygame.time.Clock()

    scale = 5

    screen = pygame.display.set_mode((octomap.shape[0]*scale,octomap.shape[1]*scale))
    surface = pygame.Surface((octomap.shape[0]*scale,octomap.shape[1]*scale))

    frame = 0

    history = []
    for i in range(500):
        print('Generating frame #', i)
        octomap, _ = progress(octomap)
        history.append(np.copy(octomap))
    input()

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                 pygame.quit(); sys.exit();

        # erase the screen
        screen.fill((255,0,0))

        try:
            octomap = history[frame]
        except:
            frame = 0

        for i in range(octomap.shape[0]):
            for j in range(octomap.shape[1]):
                if octomap[i,j] == 0:
                    brightness = 255
                else:
                    brightness = int(255.0 * octomap[i,j]/10.0)

                print(i*scale, j*scale, brightness)
                pygame.draw.rect(
                    screen, 
                    (brightness,brightness,brightness),
                    pygame.Rect(i*scale, j*scale, scale, scale)
                )

        pygame.display.update()
        # surface.blit(screen, (0,0))

        clock.tick(30)
        frame += 1
[((41, 1, 41, 24), 'orrinjelo.utils.decorators.timeit', 'timeit', ({(41, 8, 41, 23): '"""Day 11 Part 1"""'}, {}), "('Day 11 Part 1')", False, 'from orrinjelo.utils.decorators import timeit\n'), ((52, 1, 52, 24), 'orrinjelo.utils.decorators.timeit', 'timeit', ({(52, 8, 52, 23): '"""Day 11 Part 2"""'}, {}), "('Day 11 Part 2')", False, 'from orrinjelo.utils.decorators import timeit\n'), ((29, 10, 29, 25), 'numpy.where', 'np.where', ({(29, 19, 29, 24): 'a > 9'}, {}), '(a > 9)', True, 'import numpy as np\n'), ((34, 12, 34, 25), 'numpy.sum', 'np.sum', ({(34, 19, 34, 24): 'a > 9'}, {}), '(a > 9)', True, 'import numpy as np\n'), ((97, 14, 97, 46), 'numpy.random.randint', 'np.random.randint', ({(97, 32, 97, 33): '0', (97, 34, 97, 35): '9', (97, 36, 97, 45): '(100, 100)'}, {}), '(0, 9, (100, 100))', True, 'import numpy as np\n'), ((99, 4, 99, 17), 'pygame.init', 'pygame.init', ({}, {}), '()', False, 'import pygame\n'), ((100, 12, 100, 31), 'pygame.time.Clock', 'pygame.time.Clock', ({}, {}), '()', False, 'import pygame\n'), ((104, 13, 104, 85), 'pygame.display.set_mode', 'pygame.display.set_mode', ({(104, 37, 104, 84): '(octomap.shape[0] * scale, octomap.shape[1] * scale)'}, {}), '((octomap.shape[0] * scale, octomap.shape[1] * scale))', False, 'import pygame\n'), ((105, 14, 105, 77), 'pygame.Surface', 'pygame.Surface', ({(105, 29, 105, 76): '(octomap.shape[0] * scale, octomap.shape[1] * scale)'}, {}), '((octomap.shape[0] * scale, octomap.shape[1] * scale))', False, 'import pygame\n'), ((117, 21, 117, 39), 'pygame.event.get', 'pygame.event.get', ({}, {}), '()', False, 'import pygame\n'), ((143, 8, 143, 31), 'pygame.display.update', 'pygame.display.update', ({}, {}), '()', False, 'import pygame\n'), ((113, 23, 113, 39), 'numpy.copy', 'np.copy', ({(113, 31, 113, 38): 'octomap'}, {}), '(octomap)', True, 'import numpy as np\n'), ((119, 17, 119, 30), 'pygame.quit', 'pygame.quit', ({}, {}), '()', False, 'import pygame\n'), ((119, 32, 119, 42), 'sys.exit', 'sys.exit', ({}, {}), '()', False, 'import sys\n'), ((140, 20, 140, 63), 'pygame.Rect', 'pygame.Rect', ({(140, 32, 140, 39): '(i * scale)', (140, 41, 140, 48): '(j * scale)', (140, 50, 140, 55): 'scale', (140, 57, 140, 62): 'scale'}, {}), '(i * scale, j * scale, scale, scale)', False, 'import pygame\n')]