import matrix
import shared
import sys
import numpy as np
from scipy import ndimage
from pprint import pprint as pp


class Syntax:
    def __init__(self, name):
        mx = shared.load_file_int_matrix(name)
        self.mx = np.array(mx)
        self.flash_count = 0
        self.step_count = 0

    def check_syncronized(self):
        cpy = np.copy(self.mx)
        flat = np.ravel(cpy)
        if np.all(flat == 10):
            return True

    def step(self):
        self.step_count += 1
        self.increment_all()

        # Get a 2darray of bools of whether it's flashing
        self.flashed = (self.mx > 9).astype(int)
        flash_coords = np.argwhere(self.flashed == 1)  # get a list of any flashed

        if len(flash_coords) == 0:
            return

        self.flash_count += len(flash_coords)

        if flash_coords.any():
            for coord in flash_coords:
                self.flash({"c": coord[1], "r": coord[0]})

        if self.check_syncronized():
            return True
        self.reset_flashes()
        #matrix.pmx(self.mx, self.flashed)  # Print a pretty matrix

    def flash(self, coord: dict) -> np.ndarray:
        neighbor_coords = matrix.get_neighbor_coords(
            self.mx, c=coord["c"], r=coord["r"]
        )

        for c, val in neighbor_coords:
            self.add_clamp(**c)
            if self.flashed[c["r"]][c["c"]]:
                continue
            if self.mx[c["r"]][c["c"]] > 9:
                self.flash_count += 1
                self.flashed[c["r"]][c["c"]] = 1
                self.flash(c)

    def reset_flashes(self):
        for r, row in enumerate(self.mx):
            for c, col in enumerate(row):
                if self.mx[r][c] == 10:
                    self.mx[r][c] = 0

    def add_clamp(self, c, r):
        """ Clamp to 10, but add 1 """
        if self.mx[r][c] < 10:
            self.mx[r][c] += 1

    def increment_all(self):
        for r, row in enumerate(self.mx):
            for c, _ in enumerate(row):
                self.add_clamp(c, r)


def main():
    s = Syntax(shared.get_fname(11))
    print(f"before any steps")
    matrix.pmx(s.mx)
    print("----")
    for x in range(9999):
        if s.step():
            print(x)
            break
    matrix.pmx(s.mx)
    print(s.flash_count)
    print(s.step_count)


if __name__ == "__main__":
    main()