import torchvision
import tkinter as tk
from PIL import Image, ImageTk, ImageDraw, ImageFont


class ViewFrame(tk.Frame):
    def __init__(self, master, data):
        super().__init__(master)
        self.master = master

        self.canvas = tk.Canvas(self, width=28*20, height=28*20, bg='black')
        self.canvas.pack()

        self.canvas.bind('<Motion>', self.on_hover)
        self.canvas.bind('<Leave>', self.on_leave)
        self.canvas.config(cursor='none')

        self.data = data
        self.index = 0

        self.show_numbers = False

        tk.Label(self,
            text='Annotation :',
            font=('Helvetica', 12),
            anchor='w').pack(padx=10, expand=True, fill=tk.X)

        self.output_label = tk.Label(self, text='?', font=('Helvetica', 48))
        self.output_label.pack()

        self.controls = tk.Frame(self)
        self.controls.pack(fill=tk.X, padx=10, pady=10)

        self.next_button = tk.Button(self.controls, text='Next', command=self.next)
        self.next_button.pack(side=tk.RIGHT, expand=True, fill=tk.X)

        self.previous_button = tk.Button(self.controls, text='Previous', command=self.previous)
        self.previous_button.pack(side=tk.LEFT, expand=True, fill=tk.X)

        self.font = ImageFont.load_default()

        self.update_images()
        self.show_image()
    
    def on_hover(self, event):
        self.show_numbers = True
        self.show_image()
    
    def on_leave(self, event):
        self.show_numbers = False
        self.show_image()
    
    def next(self):
        self.index += 1
        self.index %= len(self.data)
        self.update_images()
        self.show_image()

    def previous(self):
        self.index -= 1
        self.index %= len(self.data)
        self.update_images()
        self.show_image()

    def update_images(self):
        self.image, self.label = self.data[self.index]
        self.image_resized = self.image.resize((28*20, 28*20))
        self.image_numbered = self.image.resize((28*20, 28*20), Image.BOX)

        image_draw = ImageDraw.Draw(self.image_numbered)
        for x in range(self.image.width):
                for y in range(self.image.height):
                    center_x = x * 20 + 10
                    center_y = y * 20 + 10
                    value = self.image.getpixel((x, y))
                    text = str(value)
                    _, _, w, h = image_draw.textbbox((0, 0), text)
                    color = 'white' if value < 128 else 'black'
                    image_draw.text((center_x-w/2, center_y-h/2), text, fill=color, font=self.font)

    def show_image(self):
        self.canvas.delete('all')

        image_copy = self.image_resized.copy()

        if self.show_numbers:
            mouse_x = self.canvas.winfo_pointerx() - self.canvas.winfo_rootx()
            mouse_y = self.canvas.winfo_pointery() - self.canvas.winfo_rooty()
            
            mask = Image.new('RGBA', image_copy.size, (0, 0, 0, 0))
            mask_draw = ImageDraw.Draw(mask)
            mask_draw.ellipse((mouse_x-100, mouse_y-100, mouse_x+100, mouse_y+100), fill='white')

            image_copy.paste(self.image_numbered, (0, 0), mask)
            image_copy_draw = ImageDraw.Draw(image_copy)
            image_copy_draw.ellipse((mouse_x-100, mouse_y-100, mouse_x+100, mouse_y+100), outline='gray', width=2)
    
        self.photo_image = ImageTk.PhotoImage(image_copy)
        self.canvas.create_image(3, 3, image=self.photo_image, anchor=tk.NW)
        self.master.title(f'Image {self.index + 1} / {len(self.data)}')
        self.output_label.config(text=str(self.label))


if __name__ == '__main__':
    data = torchvision.datasets.MNIST('data/', train=True, download=True)

    root = tk.Tk()
    view_frame = ViewFrame(root, data)
    view_frame.pack()
    root.eval('tk::PlaceWindow . center')
    root.mainloop()


