import threading
import time

import numpy as np
from PIL import Image

pool = []
conv_counter = 0


def image_convolution(data, conv_matrix):
    global conv_counter, pool
    y_data_pixels, x_data_pixels, _ = data.shape
    for y_data in range(1, y_data_pixels - 1):
        t = threading.Thread(target=image_convolution_row, args=(data, conv_matrix, y_data))
        pool.append(t)
        t.start()
    while any(t.is_alive() for t in pool):
        time.sleep(1)
        print(f"Loading with {len(pool)} threads on convolution {conv_counter}...")
    for thread in pool:
        thread.join()
    pool = []
    conv_counter += 1


def image_convolution_row(data, conv_matrix, y_data):
    y_data_pixels, x_data_pixels, _ = data.shape
    for x_data in range(1, x_data_pixels - 1):
        t = threading.Thread(target=image_convolution_pixel, args=(data, conv_matrix, y_data, x_data))
        pool.append(t)
        t.start()


def image_convolution_pixel(data, conv_matrix, y_data, x_data):
    for color in range(3):
        new_val = 0
        for y_conv in range(0, 3):
            for x_conv in range(0, 3):
                # new_val += conv_matrix[y_conv, x_conv] * data[y_data + y_conv - 1, x_data + x_conv - 1]
                continue
        global grid_new
        new_val = max(0, min(new_val, 255))
        grid_new[y_data, x_data, color] = new_val


image = Image.open("koko.jpg")

grid = np.array(image, dtype=np.uint8)
grid_new = np.zeros(grid.shape)
y_pixels, x_pixels, _ = grid.shape

convolutions = [
    np.array([
        [0, 1, 0],
        [0, 0, 0],
        [0, -1, 0]
    ]),
    np.array([
        [1, 0, 0],
        [0, 0, 0],
        [0, 0, -1]
    ]),
    np.array([
        [0, 0, 0],
        [1, 0, -1],
        [0, 0, 0]
    ]),
    np.array([
        [0, 0, 1],
        [0, 0, 0],
        [-1, 0, 0]
    ])
]

for y in range(y_pixels):
    for x in range(x_pixels):
        grid[y][x] = (grid[y][x] - (256 / 2)) % 256

for convolution in convolutions:
    image_convolution(grid, convolution)
    grid = grid_new

image = Image.fromarray(grid)
image.show()