import cv2
import numpy as np
import gradio as gr
from mtcnn import MTCNN
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.xception import preprocess_input as xcp_pre
from tensorflow.keras.applications.efficientnet import preprocess_input as eff_pre

# Load models
xcp_model = load_model("xception_model.h5")
eff_model = load_model("efficientnet_model.h5")

# Load face detector
detector = MTCNN()

def expand_box(x, y, w, h, scale=1.5, img_shape=None):
    """Expand face bounding box with margin."""
    cx, cy = x + w // 2, y + h // 2
    new_w, new_h = int(w * scale), int(h * scale)
    x1 = max(0, cx - new_w // 2)
    y1 = max(0, cy - new_h // 2)
    x2 = min(img_shape[1], cx + new_w // 2)
    y2 = min(img_shape[0], cy + new_h // 2)
    return x1, y1, x2, y2

def predict(image):
    faces = detector.detect_faces(image)
    if not faces:
        return "No faces detected", image

    results = []
    annotated = image.copy()

    for i, face in enumerate(faces):
        x, y, w, h = face['box']
        x, y, w, h = max(0, x), max(0, y), w, h
        x1, y1, x2, y2 = expand_box(x, y, w, h, scale=1.6, img_shape=image.shape)

        face_crop = image[y1:y2, x1:x2]

        # Preprocess for each model
        xcp_img = cv2.resize(face_crop, (299, 299))
        eff_img = cv2.resize(face_crop, (224, 224))

        xcp_tensor = xcp_pre(xcp_img.astype(np.float32))[np.newaxis, ...]
        eff_tensor = eff_pre(eff_img.astype(np.float32))[np.newaxis, ...]

        xcp_pred = xcp_model.predict(xcp_tensor, verbose=0).flatten()[0]
        eff_pred = eff_model.predict(eff_tensor, verbose=0).flatten()[0]
        avg_pred = (xcp_pred + eff_pred) / 2
        label = "Real" if avg_pred > 0.5 else "Fake"

        results.append(
            f"Face {i+1}: {label} (Avg: {avg_pred:.3f}, XCP: {xcp_pred:.3f}, EFF: {eff_pred:.3f})"
        )

        # Draw
        color = (0, 255, 0) if label == "Real" else (255, 0, 0)
        cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)
        cv2.putText(
            annotated,
            f"{label} ({avg_pred:.2f})",
            (x1, y1 - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            color,
            2,
        )

    return "\n".join(results), annotated

# Gradio Interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=[
        gr.Textbox(label="Predictions"),
        gr.Image(type="numpy", label="Annotated Image"),
    ],
    title="Deepfake Detector (Multi-Face Ensemble)",
    description="Detects all faces in an image and classifies each one as real or fake using Xception and EfficientNetB4 ensemble.",
)

interface.launch()