import gradio as gr
import os
from core.predict import ImageClassifier
from PIL import Image
import cv2

# Set up model path
cwd = os.getcwd()
model_path = os.path.join(cwd, "model", "lenet5_digit_classifier.pth")

# Initialize classifier
classifier = ImageClassifier(model_path=model_path)

# Function to extract and save date digits from cheque image
def save_dates_digits(image_path):
    cheque_image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if cheque_image is None:
        raise ValueError("Image could not be read. Check the path or format.")

    inverted_image = 255 - cheque_image
    cv2.imwrite("inverted_cheque.jpg", inverted_image)

    # Coordinates for cropping date region
    x, y, w, h = 1790, 100, 460, 80
    cropped_date = inverted_image[y:y+h, x:x+w]

    # Digit cropping parameters
    start_x = 5
    start_y = 0
    digit_width = 60
    digit_height = 70
    digit_labels = ['D1', 'D2', 'M1', 'M2', 'Y1', 'Y2', 'Y3', 'Y4']

    os.makedirs("./Dates_New", exist_ok=True)
    for i, label in enumerate(digit_labels):
        digit_img = cropped_date[start_y:start_y + digit_height, start_x + i * digit_width : start_x + (i + 1) * digit_width]
        save_path = os.path.join("./Dates_New", f"{label}.jpg")
        cv2.imwrite(save_path, digit_img)

# Gradio interface function
def classify_image(image):
    image_path = "uploaded_image.jpg"
    image.save(image_path)

    # Step 1: Extract digit images
    save_dates_digits(image_path)

    # Step 2: Predict each digit
    digit_labels = ['D1', 'D2', 'M1', 'M2', 'Y1', 'Y2', 'Y3', 'Y4']
    digit_folder = "./Dates_New"
    predicted_digits = []
    digit_images = []

    for label in digit_labels:
        digit_path = os.path.join(digit_folder, f"{label}.jpg")
        if os.path.exists(digit_path):
            digit_label, _ = classifier.predict(digit_path)
            predicted_digits.append(str(digit_label))
            digit_images.append((Image.open(digit_path), f"{label}: {digit_label}"))
        else:
            predicted_digits.append("?")
            digit_images.append((None, f"{label}: ?"))

    # Step 3: Reconstruct date
    day = "".join(predicted_digits[0:2])
    month = "".join(predicted_digits[2:4])
    year = "".join(predicted_digits[4:8])
    reconstructed_date = f"{day}/{month}/{year}"

    # Step 4: Load inverted cheque image
    inverted_cheque_path = "inverted_cheque.jpg"
    inverted_cheque = Image.open(inverted_cheque_path) if os.path.exists(inverted_cheque_path) else None

    # Step 5: Return all outputs
    return reconstructed_date, image, inverted_cheque, digit_images


# Gradio UI setup
demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil", label="Upload Cheque Image"),
    outputs=[
        gr.Textbox(label="Predicted Date (DD/MM/YYYY)"),
        gr.Image(label="Original Image"),
        gr.Image(label="Inverted Cheque Image"),
        gr.Gallery(label="Digit Predictions", columns=8, height=100)
    ],
    title="Cheque Date Prediction using LeNet-5 CNN",
    description="Upload a cheque image to extract and predict handwritten date using a CNN model. Each digit is shown below with its prediction."
)

# Launch the app
if __name__ == "__main__":
    demo.launch()