import gradio as gr
import cv2
import numpy as np
from PIL import Image
import os
from utils.model_loader import load_blood_group_model, load_gender_model
from utils.predictor import predict_blood_group, predict_gender
from utils.preprocessing import preprocess_fingerprint

# Load models on startup
blood_model = load_blood_group_model()
gender_model = load_gender_model()

def analyze_fingerprint(image):
    """Analyze uploaded fingerprint image"""
    if image is None:
        return "Please upload a fingerprint image.", "", "", ""
    
    try:
        # Preprocess for both models
        preprocessed_blood = preprocess_fingerprint(image, model_type="blood_group")
        preprocessed_gender = preprocess_fingerprint(image, model_type="gender")
        
        # Get predictions
        blood_result = predict_blood_group(blood_model, preprocessed_blood)
        gender_result = predict_gender(gender_model, preprocessed_gender)
        
        # Format results
        if blood_result["success"]:
            blood_text = f"**Predicted Blood Group:** {blood_result['blood_group']}\n**Confidence:** {blood_result['confidence']:.2%}"
            blood_details = "\n".join([f"{k}: {v:.2%}" for k, v in blood_result['all_scores'].items()])
        else:
            blood_text = "Error predicting blood group"
            blood_details = str(blood_result.get('error', 'Unknown error'))
        
        if gender_result["success"]:
            gender_text = f"**Predicted Gender:** {gender_result['gender']}\n**Confidence:** {gender_result['confidence']:.2%}"
        else:
            gender_text = "Error predicting gender"
        
        return blood_text, blood_details, gender_text, ""
    
    except Exception as e:
        return f"Error processing image: {str(e)}", "", "", ""

# Create Gradio interface
with gr.Blocks(title="Forensic Fingerprint Analysis", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔍 Forensic Fingerprint Analysis System")
    gr.Markdown("Upload a fingerprint image to predict blood group and gender using machine learning models.")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Fingerprint Image", type="pil")
            analyze_btn = gr.Button("Analyze Fingerprint", variant="primary")
        
        with gr.Column():
            gr.Markdown("### Results")
            blood_output = gr.Textbox(label="Blood Group Prediction", lines=2)
            blood_scores = gr.Textbox(label="All Blood Group Scores", lines=8)
            gender_output = gr.Textbox(label="Gender Prediction", lines=2)
    
    # Connect the function
    analyze_btn.click(
        analyze_fingerprint,
        inputs=image_input,
        outputs=[blood_output, blood_scores, gender_output, gr.Textbox(visible=False)]
    )
    
    gr.Markdown("---")
    gr.Markdown("**About:** 4th Year Forensic Research Project - Analyzing fingerprint patterns for biometric identification.")

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