import torch
from PIL import Image
import os
import torchvision.transforms as transforms
from Model.OCR_Model import OCRModel
import torchvision.transforms.functional as F
import torch.nn.functional as C
from guizero import App, Text, MenuBar
from tkinter import filedialog

def prediction_decode(output):
    probabilities = C.softmax(output, dim=1)
    conf, index_t = torch.max(probabilities, dim=1)
    predicted_index = index_t.item()
    conf_p = conf.item() * 100
    labels = [
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 
    'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 
    'U', 'V', 'W', 'X', 'Y', 'Z',
    'a', 'b', 'd', 'e', 'f', 'g', 'h', 'n', 'q', 'r', 't'
    ]
    predicted_char = labels[predicted_index]

    return predicted_char, conf_p

def predict():
    file_path = filedialog.askopenfilename(
        title="Open Image",
        filetypes=[("Image Files", "*.png *.jpg *.jpeg *.bmp"), ("All Files", "*.*")]
        )
    if file_path:
        print("[Status] Opening Image...")
        try:
            image = Image.open(file_path).convert("L")

            transform = transforms.Compose([
                transforms.Resize((28, 28)),
                transforms.ToTensor(),
                transforms.Normalize(mean=(0.1751,), std=(0.3332,))
            ])

            x = transform(image).unsqueeze(0)
            x = torch.transpose(x, 2, 3) 

            print(f"[AI] AI is thinking...")
            with torch.no_grad():
                predicted = model(x)

            print(f"[AI] Decoding prediction...")
            predicted_digit, conf = prediction_decode(predicted)
            result.value = f"I feel {conf:.2f}% confident that I saw the digit {predicted_digit}"
            print(f"[AI] I feel {conf:.2f}% confident that I read the digit {predicted_digit}")
        except Exception as e:
            result.value = f"Error :/"
            print(f"[Error] {e}")
    else:
        print("[Status] Cancelled")

print("[Status] Loading Model...")
script_dir = os.path.dirname(os.path.abspath(__file__))
model_path = os.path.join(script_dir, "OCR_Model.pt")
state_dic = torch.load(model_path, weights_only=True)
model = OCRModel()
model.load_state_dict(state_dic)
model.eval()
print("[Info] Model Loaded")    

app = App("Optical Character Recognizer(Digits)", width=500, height=250)
info = Text(app, text="Welcome to Optical Character Recognizer.  Upload file for recognition.")
info2 = Text(app, text="As of now, this app cannot recognize characters as full sentences, \n like this.  Such changes is for the future to be added.")
menu = MenuBar(app, 
               toplevel=["File"],
               options=[ 
                   [ ["Open", predict] ] 
                ])
result = Text(app, text=" ")

app.display()