import scipy.ndimage as nd
import imageio.v3 as iio
import numpy as np
import math
from PIL import Image, ImageFilter

IMAGE_PATH = 'engine.png'

BLUR_STRENGTH_1 = 1
BLUR_STRENGTH_2 = 3

DOG_THRESHOLD = 8

sobel_x_kernel = ImageFilter.Kernel((3, 3),
                                    (
                                        1, 0, -1,
                                        2, 0, -2,
                                        1, 0, -1
                                    ))

sobel_y_kernel = ImageFilter.Kernel((3, 3),
                                    (
                                        1, 2, 1,
                                        0, 0, 0,
                                        -1, -2, -1
                                    ))


def main():

    image = Image.open(IMAGE_PATH)

    # image = image.convert('L')

    blur_1 = image.filter(ImageFilter.GaussianBlur(BLUR_STRENGTH_1))
    blur_2 = image.filter(ImageFilter.GaussianBlur(BLUR_STRENGTH_2))

    dog: Image = difference_of_gaussians(blur_1, blur_2)

    sb = sobel(image)

    sb.save('sobel-pil.png')

    sb = sobel(dog)

    sb.save('sobel-dog.png')

    # dog.save('dog.png')

    # sb_x = dog.filter(sobel_x_kernel)

    # sb_x.save('pil-sobel_x.png')

    # image_np = np.array(dog)

    # sb_x = nd.sobel(image_np, 1)
    # sb_y = nd.sobel(image_np, 0)

    # combined_sobel = np.sqrt(np.square(sb_x) + np.square(sb_y))

    # i = Image.fromarray(combined_sobel, 'L')

    # i.save('sobel.png')

    # iio.imwrite('sobel_x.png', sb_x)
    # iio.imwrite('sobel_y.png', sb_y)


def subtract_colors(t1: tuple, t2: tuple) -> tuple:

    if type(t1) != tuple:
        if (t2 - t1) >= DOG_THRESHOLD:
            return t2 - t1
        else:
            return 0

    if len(t1) != len(t2):
        print('Len of first tuple should equal second tuple probably')
        exit(1)

    ans = []
    for i, n in enumerate(t1):
        ans.append(t2[i] - n)

    return tuple(ans)

def difference_of_gaussians(blur_1: Image, blur_2: Image) -> Image:

    w, h = blur_1.size

    dog = Image.new(blur_1.mode, blur_1.size)

    for pixel_y in range(0, h):
        for pixel_x in range(0, w):

            coords = (pixel_x, pixel_y)

            dog.putpixel(coords, subtract_colors(blur_2.getpixel(coords), blur_1.getpixel(coords)))

    return dog

# Taken from
# https://enzoftware.github.io/posts/image-filter-python
def sobel(img: Image) -> Image:
    width, height = img.size

    newimg = Image.new("RGB", (width, height), "white")
    for x in range(1, width-1):  # ignore the edge pixels for simplicity (1 to width-1)
        for y in range(1, height-1): # ignore edge pixels for simplicity (1 to height-1)

            # initialise Gx to 0 and Gy to 0 for every pixel
            Gx = 0
            Gy = 0

            # top left pixel
            p = img.getpixel((x-1, y-1))
            r = p[0]
            g = p[1]
            b = p[2]

            # intensity ranges from 0 to 765 (255 * 3)
            intensity = r + g + b

            # accumulate the value into Gx, and Gy
            Gx += -intensity
            Gy += -intensity

            # remaining left column
            p = img.getpixel((x-1, y))
            r = p[0]
            g = p[1]
            b = p[2]

            Gx += -2 * (r + g + b)

            p = img.getpixel((x-1, y+1))
            r = p[0]
            g = p[1]
            b = p[2]

            Gx += -(r + g + b)
            Gy += (r + g + b)

            # middle pixels
            p = img.getpixel((x, y-1))
            r = p[0]
            g = p[1]
            b = p[2]

            Gy += -2 * (r + g + b)

            p = img.getpixel((x, y+1))
            r = p[0]
            g = p[1]
            b = p[2]

            Gy += 2 * (r + g + b)

            # right column
            p = img.getpixel((x+1, y-1))
            r = p[0]
            g = p[1]
            b = p[2]

            Gx += (r + g + b)
            Gy += -(r + g + b)

            p = img.getpixel((x+1, y))
            r = p[0]
            g = p[1]
            b = p[2]

            Gx += 2 * (r + g + b)

            p = img.getpixel((x+1, y+1))
            r = p[0]
            g = p[1]
            b = p[2]

            Gx += (r + g + b)
            Gy += (r + g + b)

            # calculate the length of the gradient (Pythagorean theorem)
            length = math.sqrt((Gx * Gx) + (Gy * Gy))

            # normalise the length of gradient to the range 0 to 255
            length = length / 4328 * 255

            length = int(length)

            # draw the length in the edge image
            #newpixel = img.putpixel((length,length,length))
            newimg.putpixel((x,y),(length,length,length))

    return newimg

if __name__ == '__main__':
    main()


# image = iio.imread(IMAGE_PATH)

# sb_x = nd.sobel(image, 1)
# sb_y = nd.sobel(image, 0)


# iio.imwrite('sobel_x.png', sb_x)
# iio.imwrite('sobel_y.png', sb_y)