import colorsys
import plotly.graph_objects as go
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
import time


def apply_convolution_operator(operator, image):
    if isinstance(image, str):
        image = Image.open(image)
    if isinstance(image, Image.Image):
        x_pixel, y_pixel = image.size

        # turn into editable list of colors
        data = np.asarray(image, dtype=np.uint8)
        # storage for the convoluted image
        data_new = np.zeros([y_pixel, x_pixel, 3], dtype=np.uint8)

        for i in range(1, y_pixel - 1):
            for j in range(1, x_pixel - 1):
                for k in range(3):
                    new_value = (
                            operator[0, 0] * data[i - 1, j - 1, k] +
                            operator[0, 1] * data[i - 1, j, k] +
                            operator[0, 2] * data[i - 1, j + 1, k] +
                            operator[1, 0] * data[i, j - 1, k] +
                            operator[1, 1] * data[i, j, k] +
                            operator[1, 2] * data[i, j + 1, k] +
                            operator[2, 0] * data[i + 1, j - 1, k] +
                            operator[2, 1] * data[i + 1, j, k] +
                            operator[2, 2] * data[i + 1, j + 1, k]
                    )
                    new_value = max(0, min(new_value, 255))
                    data_new[i, j, k] = new_value
        return Image.fromarray(data_new)
    else:
        raise TypeError("Not a valid instance of Image.Image")


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

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

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

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

identity = np.array([
    [0, 0, 0],
    [0, 1, 0],
    [0, 0, 0]
])

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

"""
url = "https://cdn.prod.www.spiegel.de/images/9f9be1b3-ac0e-49f9-acbd-a809a102f9b6_w960_r1.778_fpx39_fpy30.jpg"
response = requests.get(url)

if response.status_code == 200:
    # image = Image.open(BytesIO(response.content))
    image = Image.open("next_generation.jpg")
    operators = [laplace]
    edited_image = image
    for operator in operators:
        edited_image = apply_convolution_operator(operator, edited_image)
    edited_image = edited_image.convert("RGB")  # Convert to RGB mode
    edited_image.show()
    edited_image.save("edited.jpg")
else:
    print(f"Failed to download image. Status code: {response.status_code}")
"""


class Heatmap_Simulator:
    k = 1
    h = 2
    delta_f = 0.005
    cookie_size = 0.5

    def __init__(self, y_size=28, x_size=28, tmp_min=0, tmp_max=1):
        init_heatmap = np.zeros((y_size, x_size))
        self.x_size = x_size
        self.y_size = y_size
        self.tmp_min = tmp_min
        self.tmp_max = tmp_max
        size = self.cookie_size / 2
        for y in range(int(y_size * (0.5 - size)), int(y_size * (0.5 + size))):
            for x in range(int(x_size * (0.5 - size)), int(x_size * (0.5 + size))):
                init_heatmap[y][x] = 1
        self.heatmap = init_heatmap

    def simulate(self, cycles=1):
        for c in range(cycles):
            new_heatmap = np.zeros((self.y_size, self.x_size))
            for y in range(self.y_size):
                for x in range(self.x_size):
                    u = self.heatmap[y, x]
                    value = u + self.delta_f * (self.k / (1 ** 2)) * self._consider_convolution(x, y)
                    value = max(self.tmp_min, min(self.tmp_max, value))
                    new_heatmap[y][x] = value
            self.heatmap = new_heatmap

    def _consider_convolution(self, x, y):
        top = self.heatmap[y - 1, x] if y > 0 else 0
        bottom = self.heatmap[y + 1, x] if y < self.heatmap.shape[0] - 1 else 0
        left = self.heatmap[y, x - 1] if x > 0 else 0
        right = self.heatmap[y, x + 1] if x < self.heatmap.shape[1] - 1 else 0
        center = self.heatmap[y, x]
        return top + bottom + left + right + center * (-4)

    def export(self):
        if self.heatmap is not None:
            return Heatmap_Image(self.heatmap, 0, 1)
        else:
            AttributeError("Heatmap is not initialized.")


class Visual_Heatmap_Simulator(Heatmap_Simulator):
    def __init__(self, y_size=28, x_size=28, tmp_min=0, tmp_max=1, delay=0.5):
        super().__init__(y_size, x_size, tmp_min, tmp_max)
        self.delay = delay

    def simulate_visual(self, cycles=100, skip=5):
        for c in range(cycles):
            if c != 0:
                self.simulate(1)
            if c % skip == 0:
                self._update_plot()
                print(f"Figure updated: {c}.")

    def _update_plot(self):
        plt.imshow(self.heatmap, cmap='viridis', interpolation='nearest')
        plt.draw()
        plt.pause(self.delay)


simulator = Visual_Heatmap_Simulator(28, 28, -0.5, 1)
simulator.simulate_visual(10000, 50)


class Heatmap_Image:
    def __init__(self, heatmap, min_val, max_val):
        self.min_val = min_val
        self.max_val = max_val
        x_size, y_size = heatmap.shape
        colormap = np.zeros((y_size, x_size, 3), dtype=np.uint8)
        for y in range(y_size):
            for x in range(x_size):
                colormap[y][x] = self._val_to_color(heatmap[y][x])
        self.image = Image.fromarray(colormap, 'RGB')

    def _val_to_color(self, value):
        normalized_value = (value - self.min_val) / (self.max_val - self.min_val)
        hsv_color = (0.66 - 0.66 * normalized_value, 1.0, 1.0)  # Blue to Red
        rgb_color = colorsys.hsv_to_rgb(*hsv_color)
        scaled_rgb_color = tuple(int(c * 255) for c in rgb_color)
        return scaled_rgb_color

    def show(self):
        self.image.show()

    def save(self, file):
        self.image.save(file)