import os
import cv2
import numpy as np
import tensorflow as tf
import pickle
import tkinter as tk
from tkinter import filedialog

# Set global parameters (update these as per your model's training settings)
IMAGE_WIDTH = 128
IMAGE_HEIGHT = 64
MAX_TEXT_LENGTH = 32

# Function to preprocess the image
def preprocess_image(image_path, image_width, image_height): 
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        print(f"Error reading image: {image_path}")
        return None
    img = cv2.resize(img, (image_width, image_height))
    img = img / 255.0
    img = img.T  # Transpose if your model requires it
    img = np.expand_dims(img, axis=-1)
    img = np.expand_dims(img, axis=0)
    return img

# Function to predict the word from the image using your model and tokenizer
def predict_word(model, image_path, tokenizer, max_text_length):
    img = preprocess_image(image_path, IMAGE_WIDTH, IMAGE_HEIGHT)
    if img is None:
        return ""
    
    # Get the model prediction
    prediction = model.predict(img)
    predicted_indices = np.argmax(prediction, axis=-1)[0]
    
    # Create a reverse mapping from index to character
    reverse_index = {index: char for char, index in tokenizer.word_index.items()}
    
    # Convert prediction indices into a string (ignoring padding zeros)
    predicted_word = ''.join([reverse_index.get(index, '') for index in predicted_indices if index != 0])
    return predicted_word

# Load your trained Keras model
model = tf.keras.models.load_model("handwritten_text_recognition_model.keras")

# Load the tokenizer (make sure 'tokenizer.pickle' exists in your project directory)
with open("tokenizer.pickle", "rb") as handle:
    tokenizer = pickle.load(handle)

# Create a simple Tkinter file dialog to select an image
root = tk.Tk()
root.withdraw()  # Hide the main window

# Open file dialog for image selection
image_path = filedialog.askopenfilename(
    title="Select an Image", 
    filetypes=[("Image Files", "*.png;*.jpg;*.jpeg")]
)

if image_path:
    predicted_word = predict_word(model, image_path, tokenizer, MAX_TEXT_LENGTH)
    print("Predicted word:", predicted_word)
else:
    print("No image selected.")