# import scipy.ndimage as nd
# import imageio.v3 as iio
# import numpy as np
import math
import os

from PIL import Image, ImageFilter
import pygame as pg

IMAGE_PATH = 'sample-images/sunflower.jpg'
FONT_PATH = 'fonts/scientifica.ttf'

BLUR_STRENGTH_1 = 1
BLUR_STRENGTH_2 = 1.2

DOG_THRESHOLD = 8

WHITE = (255,255,255)

TEXT_WIDTH = 6 # on font size 12 monocraft medium
TEXT_HEIGHT = 11

SCALE_FACTOR = 2
RESIZE_FACTOR = 2

SCALED_WIDTH = math.floor(TEXT_WIDTH / SCALE_FACTOR)
SCALED_HEIGHT = math.floor(TEXT_HEIGHT / SCALE_FACTOR)

DEBUG = True

if not os.path.exists('./debug/') and DEBUG:
    os.makedirs('./debug')

# '█@?OPoci. '
CHARS = list(reversed([
    u'█',
    '@',
    '?',
    'O',
    'P',
    'o',
    'c',
    'i',
    '.',
    ' '
]))

def main():

    image = Image.open(IMAGE_PATH)
    image = image.resize((image.width * RESIZE_FACTOR, image.height * RESIZE_FACTOR))

    L_image = image.convert('L')

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


    print(image.size)
    dog: Image = difference_of_gaussians(blur_1, blur_2)

    dog.save('dog.png')

    sb, gradient = sobel(dog)

    w, h = sb.size

    sb.save('sobel.png')

    text_grid_width = w / SCALED_WIDTH
    text_grind_height = h / SCALED_HEIGHT

    pg.init()
    font = pg.font.Font(FONT_PATH, 12)
    window = pg.display.set_mode((text_grid_width * TEXT_WIDTH, text_grind_height * TEXT_HEIGHT))

    y_offset = 0

    for y in range(math.floor(h / SCALED_HEIGHT)):
        x_offset = 0

        for x in range(math.floor(w / SCALED_WIDTH)):
            histogram = {}

            # Collect the most common char for a group of pixels
            colors = []

            for y2 in range(SCALED_HEIGHT):
                for x2 in range(SCALED_WIDTH):

                    real_x = x2 + x_offset
                    real_y = y2 + y_offset

                    gradient_v = gradient[real_y * w + real_x]

                    char = ' '

                    if gradient_v:
                        char = match_gradient(gradient_v)
                    else:
                        char = CHARS[round(L_image.getpixel((real_x, real_y))/(255/len(CHARS))) - 1]

                    if char in histogram:
                        histogram[char] += 1
                    else:
                        histogram[char] = 1

                    colors.append(image.getpixel((real_x, real_y)))

            color_avg = average_colors(colors)

            # get most common
            most_common = None
            score = float('-inf')
            for char in histogram:
                if histogram[char] > score:
                    score = histogram[char]
                    most_common = char

            rendered_char = font.render(str(most_common), True, color_avg)
            window.blit(rendered_char, (x * TEXT_WIDTH, y * TEXT_HEIGHT))

            x_offset += SCALED_WIDTH

        y_offset += SCALED_HEIGHT

    pg.display.update()
    pg.image.save(window, "render.png")
    pg.quit()

    if DEBUG:
        sb.save('debug/sobel-dog.png')
        dog.save('debug/dog.png')

        # Draw only the ASCII edges
        pg.init()
        font = pg.font.Font(FONT_PATH, 12)
        window = pg.display.set_mode((text_grid_width * TEXT_WIDTH, text_grind_height * TEXT_HEIGHT))

        y_offset = 0

        for y in range(math.floor(h / SCALED_HEIGHT)):
            x_offset = 0

            for x in range(math.floor(w / SCALED_WIDTH)):
                histogram = {}

                # Collect the most common char for a group of pixels
                colors = []

                for y2 in range(SCALED_HEIGHT):
                    for x2 in range(SCALED_WIDTH):

                        real_x = x2 + x_offset
                        real_y = y2 + y_offset

                        gradient_v = gradient[real_y * w + real_x]

                        char = ' '

                        if gradient_v:
                            char = match_gradient(gradient_v)
                        else:
                            char = ' '

                        if char in histogram:
                            histogram[char] += 1
                        else:
                            histogram[char] = 1

                        colors.append(image.getpixel((real_x, real_y)))

                color_avg = average_colors(colors)

                # get most common
                most_common = None
                score = float('-inf')
                for char in histogram:
                    if histogram[char] > score:
                        score = histogram[char]
                        most_common = char

                rendered_char = font.render(most_common, True, color_avg)
                window.blit(rendered_char, (x * TEXT_WIDTH, y * TEXT_HEIGHT))

                x_offset += SCALED_WIDTH

            y_offset += SCALED_HEIGHT

        pg.display.update()
        pg.image.save(window, "debug/edges-render.png")
        pg.quit()




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) -> tuple[Image.Image, list]:
    if img.mode == 'L':
        # return sobel_L(img)
        pass
    img = img.convert('RGB')

    width, height = img.size

    newimg = Image.new("RGB", (width, height), "white")
    gradient = [None for x in range(0, width * height)]
    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
            r, g, b = img.getpixel((x - 1, y - 1))

            # 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
            r, g, b = img.getpixel((x-1, y))

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

            r, g, b = img.getpixel((x-1, y+1))

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

            # middle pixels
            r, g, b = img.getpixel((x, y-1))

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

            r, g, b = img.getpixel((x, y+1))

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

            # right column
            r, g, b = img.getpixel((x+1, y-1))

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

            r, g, b = img.getpixel((x+1, y))

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

            r, g, b = img.getpixel((x+1, y+1))

            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)
            gradient_v = math.atan2(Gy, Gx)
            # print(gradient_v)

            # draw the length in the edge image
            #newpixel = img.putpixel((length,length,length))
            newimg.putpixel((x,y),(length,length,length))
            if length < 20:
                gradient[y * width + x] = gradient_v

    return (newimg, gradient)

def sobel_L(img: Image) -> tuple[Image.Image, list]:
    pass

# TODO @0x01FE : refactor plz increment by 30 & 60
def match_gradient(n: float) -> str:

    n = abs(n)

    if n < math.radians(30):
        return '|'
    elif n < math.radians(60):
        return '/'
    elif n < math.radians(120):
        return '-'
    elif n < math.radians(150):
        return '\\'
    elif n < math.radians(210):
        return '|'
    elif n < math.radians(240):
        return '/'
    elif n < math.radians(300):
        return '-'
    elif n < math.radians(330):
        return '\\'
    else:
        return '|'

def average_colors(colors: list):
    avg = [0 for x in range(len(colors[0]))]

    for color in colors:
        for i, num in enumerate(color):
            avg[i] += num

    for i, n in enumerate(avg):
        avg[i] /= len(colors)

    return avg

if __name__ == '__main__':
    main()