import os
import cv2
import numpy as np
import gradio as gr
import tensorflow as tf
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
from huggingface_hub import hf_hub_download
from mtcnn import MTCNN

# Download and load models
xcp_path = hf_hub_download(repo_id="Zeyadd-Mostaffa/deepfake-image-detector", filename="xception_model.h5")
eff_path = hf_hub_download(repo_id="Zeyadd-Mostaffa/deepfake-image-detector", filename="efficientnet_model.h5")
xcp_model = load_model(xcp_path)
eff_model = load_model(eff_path)

# Face detector
detector = MTCNN()

def extract_face(image):
    faces = detector.detect_faces(image)
    if not faces:
        return None
    x, y, w, h = faces[0]['box']
    x, y = max(0, x), max(0, y)
    return image[y:y+h, x:x+w]

def predict(image):
    face = extract_face(image)
    if face is None:
        return "No face detected", None

    # Xception
    xcp_img = cv2.resize(face, (299, 299))
    xcp_tensor = xcp_pre(xcp_img.astype(np.float32))[np.newaxis, ...]
    xcp_pred = xcp_model.predict(xcp_tensor, verbose=0).flatten()[0]

    # EfficientNet
    eff_img = cv2.resize(face, (224, 224))
    eff_tensor = eff_pre(eff_img.astype(np.float32))[np.newaxis, ...]
    eff_pred = eff_model.predict(eff_tensor, verbose=0).flatten()[0]

    # Ensemble
    avg_pred = (xcp_pred + eff_pred) / 2
    label = "Real" if avg_pred > 0.5 else "Fake"

    # Log probabilities
    print(f"Xception: {xcp_pred:.4f}, EfficientNetB4: {eff_pred:.4f}, Ensemble Avg: {avg_pred:.4f}")

    # Return label with confidence
    result = f"{label} (Avg: {avg_pred:.3f}, XCP: {xcp_pred:.3f}, EFF: {eff_pred:.3f})"
    return result, face

interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=[
        gr.Label(label="Prediction"),
        gr.Image(type="numpy", label="Detected Face")
    ],
    title="Deepfake Image Detector (Ensemble: Xception + EfficientNetB4)",
    description="Upload an image. The model detects the face, classifies it as real or fake using an ensemble of Xception and EfficientNetB4."
)

interface.launch()