# app.py

import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
import json
import os

# --- 1. Define int_to_char mapping and decode_prediction function ---
# This part is crucial and should accurately reflect what your model was trained on.
# We'll load int_to_char from the JSON file that was pushed to the repo.

# Get the directory where app.py is located.
# When deployed on Hugging Face Spaces, your model files will typically be in the
# same root directory as app.py if it's cloned from a model repo.
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

# Define paths to model and mapping relative to CURRENT_DIR
MODEL_PATH = os.path.join(CURRENT_DIR, "captcha_recognition_model_char.keras")
INT_TO_CHAR_PATH = os.path.join(CURRENT_DIR, "int_to_char.json")

try:
    # Load the int_to_char mapping from the JSON file
    with open(INT_TO_CHAR_PATH, "r") as f:
        str_int_to_char_mapping = json.load(f)
        # Convert keys back to integers as expected by decode_prediction
        int_to_char = {int(k): v for k, v in str_int_to_char_mapping.items()}
    print(f"int_to_char mapping loaded successfully from {INT_TO_CHAR_PATH}")
except Exception as e:
    print(f"Error loading int_to_char.json: {e}")
    # Fallback to a default or raise an error if the mapping is critical
    # For robust deployment, ensure int_to_char.json is always present and valid.
    int_to_char = {i: chr(i + ord('A')) for i in range(26)} # Example placeholder
    int_to_char.update({26 + i: str(i) for i in range(10)})
    int_to_char.update({36 + i: chr(i + ord('a')) for i in range(26)})
    int_to_char[0] = '<pad>' # Assuming 0 is pad
    print("Using a default placeholder for int_to_char due to error. Please verify original mapping.")

# Assuming fixed_solution_length is known from your model design.
# You might need to retrieve this from your model's config if it's not truly fixed,
# but for most captcha models, it's a fixed value.
fixed_solution_length = 5 # <--- IMPORTANT: Adjust this if your actual fixed_solution_length is different!

def decode_prediction(prediction_output, int_to_char_mapping):
    """Decodes the integer-encoded prediction back to a string."""
    # The prediction output from a Keras model is a NumPy array.
    # It usually has shape (batch_size, fixed_solution_length, num_classes)
    predicted_indices = np.argmax(prediction_output, axis=-1)[0] # Get indices for the first image in batch

    # Convert indices back to characters using the mapping
    predicted_chars = [int_to_char_mapping.get(idx, '') for idx in predicted_indices]

    # Join the characters to form the solution string, excluding padding
    solution = "".join([char for char in predicted_chars if char != '<pad>'])

    return solution

# --- 2. Load the pre-trained Keras model ---
# This function will run once when the Gradio app starts.
def load_model():
    try:
        model = tf.keras.models.load_model(MODEL_PATH)
        print(f"Model loaded successfully from {MODEL_PATH}")
        return model
    except Exception as e:
        print(f"Error loading the model from {MODEL_PATH}: {e}")
        # For deployment, this should ideally not fail.
        # Ensure your model is correctly pushed as SavedModel.
        return None

model = load_model()

# --- 3. Define the prediction function for Gradio ---
def predict_captcha(image: Image.Image) -> str:
    if model is None:
        return "Error: Model not loaded. Please check logs."

    # Preprocess the input image to match model's expected input
    # Ensure this matches the preprocessing done during training!
    img = image.resize((200, 50)) # Model input width, height (from previous discussion)
    img_array = np.array(img).astype(np.float32)
    img_array = np.expand_dims(img_array, axis=0) # Add batch dimension

    # Uncomment and adjust if you applied normalization during training
    # img_array = img_array / 255.0

    # Make prediction
    prediction = model.predict(img_array, verbose=0)

    # Decode the prediction
    decoded_solution = decode_prediction(prediction, int_to_char)

    return decoded_solution

# --- 4. Create the Gradio Interface ---
iface = gr.Interface(
    fn=predict_captcha,
    inputs=gr.Image(type="pil", label="Upload Captcha Image"),
    outputs=gr.Textbox(label="Predicted Captcha"),
    title="Captcha Recognition",
    description="Upload a captcha image (200x50 pixels expected) to get the predicted text.",
    examples=[
        # You can add example image paths here for the Gradio demo.
        # These images should be present in your Hugging Face Space repository.
        # e.g., "./example_captcha_1.png", "./example_captcha_2.png"
    ],
    allow_flagging="never", # Optional: Disable flagging data
    live=False # Set to True for real-time inference as you draw/upload
)

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