import gradio as gr
import numpy as np
import tensorflow as tf
import cv2
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 mtcnn import MTCNN
import os
import warnings
warnings.filterwarnings("ignore")

# Force TF to suppress log-level warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

# Load models from local (downloaded from HF first in app setup)
xcp_model = load_model("xception_model.h5")
eff_model = load_model("efficientnet_model.h5")

# Grad-CAM for Xception
def grad_cam(model, img_array, size, preprocess_fn):
    img = cv2.resize(img_array, size)
    input_tensor = preprocess_fn(np.expand_dims(img, axis=0).astype(np.float32))
    input_tensor = tf.convert_to_tensor(input_tensor)

    with tf.GradientTape() as tape:
        conv_layer = model.get_layer(index=-5).output
        grad_model = tf.keras.models.Model([model.inputs], [conv_layer, model.output])
        conv_outputs, predictions = grad_model(input_tensor)
        loss = predictions[:, 0]

    grads = tape.gradient(loss, conv_outputs)
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    cam = tf.reduce_sum(tf.multiply(pooled_grads, conv_outputs), axis=-1).numpy()[0]

    cam = np.maximum(cam, 0)
    cam = cam / (cam.max() + 1e-8)
    cam = (cam * 255).astype(np.uint8)
    cam = cam.numpy() if hasattr(cam, 'numpy') else cam
    cam = cv2.resize(cam, size)
    heatmap = cv2.applyColorMap(cam, cv2.COLORMAP_JET)
    superimposed_img = cv2.addWeighted(cv2.cvtColor(img, cv2.COLOR_RGB2BGR), 0.6, heatmap, 0.4, 0)
    return superimposed_img

# Face detector
detector = MTCNN()

def detect_face(image):
    faces = detector.detect_faces(image)
    if not faces:
        raise ValueError("No face detected.")
    x, y, w, h = faces[0]['box']
    return image[y:y+h, x:x+w]

def predict(image):
    try:
        face = detect_face(image)

        xcp_img = cv2.resize(face, (299, 299))
        eff_img = cv2.resize(face, (224, 224))

        xcp_input = np.expand_dims(xcp_pre(xcp_img.astype(np.float32)), axis=0)
        eff_input = np.expand_dims(eff_pre(eff_img.astype(np.float32)), axis=0)

        xcp_pred = xcp_model.predict(xcp_input)[0][0]
        eff_pred = eff_model.predict(eff_input)[0][0]

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

        cam_img = grad_cam(xcp_model, face, (299, 299), xcp_pre)

        return label, cam_img
    except Exception as e:
        return "خطأ", "خطأ"

gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Face Image"),
    outputs=[gr.Label(label="Prediction"), gr.Image(label="Grad-CAM Explanation")],
    title="Deepfake Image Detector (with Grad-CAM)",
    description="Upload an image. We detect the face, classify using an ensemble (Xception + EfficientNetB4), and explain the prediction using Grad-CAM on Xception."
).launch()