import colorsys
from io import BytesIO

import requests
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt


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]
])

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


url = "https://media.newyorker.com/photos/5e3b18f932d7150008e6cf16/master/w_2240,c_limit/Neima-Picard.jpg"
response = requests.get(url)

if response.status_code == 200:
    image = Image.open(BytesIO(response.content))
    operators = [g1, g2, g3, g4]
    for operator in operators:
        image_edited = apply_convolution_operator(operator, image)
        image_edited.show()
else:
    print(f"Failed to download image. Status code: {response.status_code}")


class Heatmap_Simulator:
    k, h, delta_f = 1, 0.01, 0.05
    # Add more accurate cooling effect:
    # tmp_min, tmp_max, env = -0.5, 1, -0.5
    tmp_min, tmp_max, env = 0, 1, 0
    cookie_size = 0.5

    def __init__(self, y_size=28, x_size=28):
        init_heatmap = np.full((y_size, x_size), self.tmp_min)
        self.x_size = x_size
        self.y_size = y_size
        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] = self.tmp_max
        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 self.env
        bottom = self.heatmap[y + 1, x] if y < self.heatmap.shape[0] - 1 else self.env
        left = self.heatmap[y, x - 1] if x > 0 else self.env
        right = self.heatmap[y, x + 1] if x < self.heatmap.shape[1] - 1 else self.env
        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, delay=0.5):
        super().__init__(y_size, x_size)
        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', vmin=self.tmp_min, vmax=self.tmp_max)
        plt.draw()
        plt.pause(self.delay)


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)


"""
sim = Heatmap_Simulator()
img = sim.export()
img.save("img_0.jpg")
for i in range(1,20):
    sim.simulate(100)
    img = sim.export()
    img.save(f"img_{i}.jpg")
"""