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

# ===================== LOAD MODELS =====================

blood_model = None
gender_model = None

def load_models():
    global blood_model, gender_model

    if blood_model is None:
        try:
            blood_model = tf.keras.models.load_model("blood_model.h5")
        except Exception as e:
            print(f"Blood model error: {e}")
            blood_model = None

    if gender_model is None:
        try:
            gender_model = tf.keras.models.load_model("gender_model.keras")
        except Exception as e:
            print(f"Gender model error: {e}")
            gender_model = None

    return blood_model, gender_model


# ===================== PREPROCESS =====================

def preprocess_image(image):
    image = np.array(image)
    img = cv2.resize(image, (224, 224))
    img = img / 255.0
    img = np.expand_dims(img, axis=0)
    return img


# ===================== LABELS =====================

BLOOD_GROUPS = ["A+", "A-", "B+", "B-", "AB+", "AB-", "O+", "O-"]
GENDERS = ["Male", "Female"]


# ===================== PREDICTION =====================

def analyze_fingerprint(image):

    if image is None:
        return "Please upload image", "", ""

    try:
        blood_model, gender_model = load_models()

        if blood_model is None or gender_model is None:
            return "Model files missing", "", ""

        img = preprocess_image(image)

        # Blood prediction
        blood_pred = blood_model.predict(img)[0]
        blood_idx = np.argmax(blood_pred)

        blood_text = f"Predicted Blood Group: {BLOOD_GROUPS[blood_idx]} ({blood_pred[blood_idx]:.2%})"
        blood_scores = "\n".join(
            [f"{bg}: {score:.2%}" for bg, score in zip(BLOOD_GROUPS, blood_pred)]
        )

        # Gender prediction
        gender_pred = gender_model.predict(img)[0]
        gender_idx = np.argmax(gender_pred)

        gender_text = f"Predicted Gender: {GENDERS[gender_idx]} ({gender_pred[gender_idx]:.2%})"

        return blood_text, blood_scores, gender_text

    except Exception as e:
        return f"Error: {str(e)}", "", ""


# ===================== UI =====================

with gr.Blocks(title="Forensic Fingerprint Analysis") as demo:
    gr.Markdown("# 🔍 Forensic Fingerprint Analysis")
    gr.Markdown("Upload a fingerprint image to predict blood group and gender.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Fingerprint")
            btn = gr.Button("Analyze")

        with gr.Column():
            blood_output = gr.Textbox(label="Blood Group")
            blood_scores = gr.Textbox(label="All Scores")
            gender_output = gr.Textbox(label="Gender")

    btn.click(
        analyze_fingerprint,
        inputs=image_input,
        outputs=[blood_output, blood_scores, gender_output],
    )


# ===================== RUN =====================

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)