import tkinter as tk
import torch
import torchvision.transforms as transforms
from PIL import Image, ImageDraw, ImageTk

from network import *
from preprocessing import preprocess


NETWORK = SimpleNet  # SimpleNet, DeepNet or ConvNet


class DrawFrame(tk.Frame):
    def __init__(self, master, network_factory):
        super().__init__(master)
        self.master = master
        self.canvas = tk.Canvas(self, width=28*15, height=28*15, bg='black')
        self.canvas.pack()
        self.canvas.bind('<Button-1>', self.start_draw)
        self.canvas.bind('<ButtonRelease-1>', self.stop_draw)
        self.canvas.config(cursor='pencil')

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

        self.strike_width = 40

        self.width_control = tk.Scale(self, from_=1, to=7, orient=tk.HORIZONTAL)
        self.width_control.pack(fill=tk.X, expand=True, padx=10, pady=10)
        self.width_control.set(self.strike_width // 10)
        self.width_control.bind('<ButtonRelease-1>', self.set_width)

        tk.Label(self,
            text='Entrée du réseau :',
            font=('Helvetica', 12),
            anchor='w').pack(padx=10, expand=True, fill=tk.X)

        self.image_preview_empty = ImageTk.PhotoImage(Image.new('L', (28, 28), 'black'))
        self.preview = tk.Label(self, image=self.image_preview_empty)
        self.preview.pack()

        tk.Label(self,
            text='Prédiction :',
            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(padx=10)

        self.bars_canvas = tk.Canvas(self, width=28*15, height=100, bg='white')
        self.bars_canvas.pack()
        self.set_values([0] * 10)

        self.image = Image.new('L', (28*15, 28*15), 'black')

        self.network = network_factory()
        name = network_factory.__name__
        try:
            self.network.load_state_dict(torch.load(f'data/model-{name}.pth'))
        except FileNotFoundError:
            print(f'Pas de modèle trouvé pour {name}')


        self.clear_button = tk.Button(self, text='Tout effacer', command=self.clear)
        self.clear_button.pack(fill=tk.X, pady=10, padx=10)
    
    def set_values(self, values):
        self.bars_canvas.delete('all')
        for i, value in enumerate(values):
            start_x = i * 40 + 20
            end_x = start_x + 30
            start_y = 100 - value * 100
            end_y = 200
            self.bars_canvas.create_rectangle(
                start_x, start_y, end_x, end_y,
                fill='black')

    def set_width(self, event):
        self.strike_width = self.width_control.get() * 10
    
    def start_draw(self, event):
        self.last_x, self.last_y = event.x, event.y
        self.canvas.bind('<B1-Motion>', self.draw)
        self.draw(event)

    def draw(self, event):
        x, y = event.x, event.y
        draw = ImageDraw.Draw(self.image)
        draw.line((self.last_x, self.last_y, x, y), fill='white', width=self.strike_width)
        draw.ellipse(
            (x-self.strike_width//2, y-self.strike_width//2, x+self.strike_width//2, y+self.strike_width//2),
            fill='white')
        self.last_x, self.last_y = x, y
        self.photo_image = ImageTk.PhotoImage(self.image)
        self.canvas.create_image(3, 3, image=self.photo_image, anchor=tk.NW)
        self.predict()

    def stop_draw(self, event):
        self.canvas.unbind('<B1-Motion>')
    
    def clear(self):
        self.canvas.delete('all')
        self.image = Image.new('L', (28*15, 28*15), 'black')
        self.preview.config(image=self.image_preview_empty)
        self.output_label.config(text='?')
        self.set_values([0] * 10)

    def predict(self):
        image = preprocess(self.image)
        self.image_preview = ImageTk.PhotoImage(image)
        self.preview.config(image=self.image_preview)

        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.0,), (1.0,))
        ])
        image = transform(image)
        
        self.network.eval()
        with torch.no_grad():
            output = self.network(image.unsqueeze(0).float())
            probs = torch.exp(output).view(-1).tolist()
            self.set_values(probs)
            pred = output.data.max(1, keepdim=True)[1].item()
            self.output_label.config(text=str(pred))
        
        

if __name__ == '__main__':
    root = tk.Tk()
    root.title(f'Reconnaissance de chiffres - {NETWORK.__name__}')
    DrawFrame(root, NETWORK).pack()
    root.eval('tk::PlaceWindow %s center' % root.winfo_pathname(root.winfo_id()))
    root.mainloop()