import sys
from matrix import load_matrix_file, get_neighbors, get_size
import shared
from scipy import ndimage
import imageio
from skimage.measure import euler_number, label
from skimage.morphology import flood_fill, flood
import matplotlib.pyplot as plt
import numpy as np
from math import prod
import random


class Smoke:
    def __init__(self, name):
        self.footprint = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
        self.parse(name)

    def parse(self, name):
        _split = lambda line: list(map(int, [l for l in line]))
        self.mx = np.array(load_matrix_file(name, func=_split))
        self.mx = np.pad(self.mx, pad_width=1, mode="constant", constant_values="9")
        self.set_values()
        self.set_coords()

    def set_values(self):
        self.lows = (
            ndimage.generic_filter(self.mx, np.min, footprint=self.footprint) == self.mx
        )
        self.low_values = self.mx[self.lows][self.mx[self.lows] < 9]
        self.highs = (
            ndimage.generic_filter(self.mx, np.max, footprint=self.footprint) == self.mx
        )
        mask = self.mx != 9
        self.mask = mask.astype(int)

    def set_coords(self):
        lc = []
        for idr, row in enumerate(self.lows):
            if idr == 0 or idr == len(self.lows) - 1:
                continue
            for idc, col in enumerate(row):
                if idc == 0 or idc == len(row) - 1:
                    continue
                if col:
                    lc.append((idr, idc))
        self.low_coords = lc

    def total_low_values(self):
        return (self.low_values + 1).sum()

    def find_chunks(self):
        cell_sizes = []
        self.cells = []
        by_center = []
        for x, y in self.low_coords:
            current = np.zeros_like(self.mx)
            m = flood_fill(self.mask, (x, y), -1, connectivity=1)
            current_cell = (m == -1).astype(int)
            self.cells.append(current_cell)
            cells = np.argwhere(current_cell == 1)
            by_center.append((len(cells), (x, y), cells))
            cell_sizes.append(len(cells))
        cells_sorted = list(set(sorted(cell_sizes)))
        self.by_center = by_center
        largest = cells_sorted[-3:]
        assert len(largest) == 3
        self.largest = largest
        return prod(largest)

    def animate(self, save=False):
        breakpoint()
        largest = [x for x in self.by_center if x[0] in self.largest][-3:]
        current = np.zeros_like(self.mx)
        shuf = self.by_center[:]
        random.shuffle(shuf)

        last_file = 0

        for idx, (length, xy, c) in enumerate(shuf):
            last_file = idx
            if length in largest:
                continue

            for x, y in c:
                current[y][x] = 1

            fig, ax = plt.subplots()
            ax.imshow(current, cmap=plt.cm.gray)
            ax.axis("off")
            if save:
                plt.savefig(f"09/{idx:03}.png")
            plt.close()

        lg = np.zeros_like(self.mx)
        for _, _, c in largest:
            for x, y in c:
                lg[y][x] = 1

        fig, ax = plt.subplots()
        ax.imshow(lg)
        ax.axis("off")
        plt.savefig(f"09/largest.png")
        plt.close()

        with imageio.get_writer("09/day09.gif", mode="I") as writer:
            names = [f"09/{x:03}.png" for x in range(334)]
            for filename in names:
                try:
                    image = imageio.imread(filename)
                    writer.append_data(image)
                except FileNotFoundError:
                    pass

            for x in range(15):
                image = imageio.imread("09/largest.png")
                for x in range(5):
                    writer.append_data(image)
                image = imageio.imread(f"09/{last_file:03}.png")
                for x in range(5):
                    writer.append_data(image)


def main():
    s = Smoke(shared.get_fname(9))
    print(s.total_low_values())
    print(s.find_chunks())
    if sys.argv[-1] != "--sample":
        s.animate(1)


if __name__ == "__main__":
    main()