import tensorflow as tf
import numpy as np
from tensorflow.keras.utils import CustomObjectScope
from PIL import Image
import gradio as gr
import os
from huggingface_hub import hf_hub_download

# Function to calculate absolute difference
def absolute_difference(x):
    return tf.abs(x[0] - x[1])

# Download model from Hugging Face Hub
def load_model_from_hub():
    try:
        print("Downloading model from Hugging Face Hub...")
        model_path = hf_hub_download(
            repo_id="Nithikorn/Signature_Detection",
            filename="final_signature_model1.keras" 
        )
        
        print(f"✅ Model downloaded successfully at: {model_path}")
        
        # Load model with CustomObjectScope
        with CustomObjectScope({'absolute_difference': absolute_difference}):
            model = tf.keras.models.load_model(model_path)
            
        print("✅ Model loaded successfully!")
        return model
    except Exception as e:
        print(f"❌ Error loading model: {str(e)}")
        raise e

# Load model
try:
    model = load_model_from_hub()
except Exception as e:
    print(f"❌ Cannot load model: {str(e)}")
    model = None

def preprocess_image(image):
    image = image.convert("RGB")  # Convert to RGB
    image = image.resize((224, 224))  # Resize
    image = np.array(image, dtype=np.float32) / 255.0  # Normalize
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image

# Signature comparison function
def predict_signature(image1, image2):
    print("Starting prediction...")
    
    if image1 is None or image2 is None:
        print("No images found")
        return "❌ Please upload both images"
    
    if model is None:
        return "❌ Unable to load model. Please check if you've uploaded the model to Hugging Face Hub"
        
    try:
        print("Processing image 1...")
        img1_processed = preprocess_image(image1)
        print("Processing image 2...")
        img2_processed = preprocess_image(image2)
        
        print("Sending data to model...")
        prediction = model.predict([img1_processed, img2_processed])
        
        print(f"Received result: {prediction}, shape: {prediction.shape}")
        
        # Convert result to numeric value
        if prediction.shape == (1, 1):
            prediction_value = float(prediction[0][0])
        elif prediction.shape == (1,):
            prediction_value = float(prediction[0])
        else:
            prediction_value = float(prediction.flatten()[0])
        
        print(f"Value obtained: {prediction_value}")
        
        threshold = 0.43
        result = "✅ Genuine Signature" if prediction_value >= threshold else "❌ Forged Signature"
        
        print(f"Verification result: {result}")
        return f"Similarity Score: {prediction_value:.4f}\n{result}"
        
    except Exception as e:
        print(f"Error occurred: {str(e)}")
        import traceback
        traceback.print_exc()  # Show full error details
        return f"❌ Error occurred: {str(e)}"

# Create Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Signature Verification System")
    
    with gr.Row():
        img1 = gr.Image(type="pil", label="Signature 1", show_label=True)
        img2 = gr.Image(type="pil", label="Signature 2", show_label=True)
    
    btn_predict = gr.Button("Verify Signatures")
    output = gr.Textbox(label="Result")
    
    btn_predict.click(predict_signature, inputs=[img1, img2], outputs=output)

# Run Gradio app
demo.launch()