import gradio as gr
import numpy as np
import tensorflow as tf
import cv2
import os
import dlib
from imutils import face_utils
from scipy.spatial import distance as dist

# ==================================================
# CONFIGURATION
# ==================================================
MODEL_PATH = "./models/deepfake_detector_multi_input.h5"
IMG_SIZE = (224, 224)
NUM_FRAMES = 20
DLIB_MODEL = "shape_predictor_68_face_landmarks.dat"
CASCADE_PATH = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
EAR_THRESHOLD = 0.25
EAR_CONSEC_FRAMES = 3

# ==================================================
# VERIFY DLIB LANDMARK MODEL
# ==================================================
if not os.path.exists(DLIB_MODEL):
    raise FileNotFoundError(
        "Missing 'shape_predictor_68_face_landmarks.dat'. "
        "Upload it to your Space root directory."
    )

# ==================================================
# LOAD MODELS
# ==================================================
print("[INFO] Loading deepfake detection model...")
model = tf.keras.models.load_model(MODEL_PATH)
print("[INFO] Model loaded successfully!")

# Enable GPU memory growth
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    for g in gpus:
        tf.config.experimental.set_memory_growth(g, True)

# Initialize detectors
face_cascade = cv2.CascadeClassifier(CASCADE_PATH)
dlib_detector = dlib.get_frontal_face_detector()
dlib_predictor = dlib.shape_predictor(DLIB_MODEL)

# ==================================================
# HELPER FUNCTIONS
# ==================================================
def eye_aspect_ratio(eye):
    A = dist.euclidean(eye[1], eye[5])
    B = dist.euclidean(eye[2], eye[4])
    C = dist.euclidean(eye[0], eye[3])
    return (A + B) / (2.0 * C)


def extract_blink_features(video_path):
    cap = cv2.VideoCapture(video_path)
    (lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
    (rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]

    ear_values, blink_count, closed = [], 0, 0
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    for _ in range(total_frames):
        ret, frame = cap.read()
        if not ret:
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        rects = dlib_detector(gray, 0)
        if len(rects) > 0:
            shape = dlib_predictor(gray, rects[0])
            shape = face_utils.shape_to_np(shape)
            leftEye, rightEye = shape[lStart:lEnd], shape[rStart:rEnd]
            ear = (eye_aspect_ratio(leftEye) + eye_aspect_ratio(rightEye)) / 2.0
            ear_values.append(ear)
            if ear < EAR_THRESHOLD:
                closed += 1
            else:
                if closed >= EAR_CONSEC_FRAMES:
                    blink_count += 1
                closed = 0

    cap.release()
    if not ear_values:
        return np.zeros((1, 3), dtype=np.float32)

    blink_freq = blink_count / max(len(ear_values), 1)
    ear_var = np.var(ear_values)
    return np.array([[blink_count, blink_freq, ear_var]], dtype=np.float32)


def extract_faces(video_path, num_frames=NUM_FRAMES, size=IMG_SIZE):
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    step = max(1, total // num_frames)
    frames = []

    count = 0
    while cap.isOpened() and len(frames) < num_frames:
        ret, frame = cap.read()
        if not ret:
            break
        if count % step == 0:
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = face_cascade.detectMultiScale(gray, 1.3, 5)
            if len(faces) > 0:
                x, y, w, h = sorted(faces, key=lambda b: b[2]*b[3], reverse=True)[0]
                face = frame[y:y+h, x:x+w]
            else:
                face = frame
            face = cv2.resize(face, size)
            frames.append(face / 255.0)
        count += 1
    cap.release()

    if not frames:
        raise ValueError("No frames extracted.")
    return np.array(frames)

# ==================================================
# PREDICTION PIPELINE
# ==================================================
def predict(video):
    if not video:
        return "⚠️ Please upload a video."

    try:
        faces = extract_faces(video)
        blink = extract_blink_features(video)
        blink_features = np.tile(blink, (faces.shape[0], 1))

        preds = model.predict([faces, blink_features], verbose=0)
        score = float(np.mean(preds))
        label = "🧠 FAKE" if score > 0.5 else "✅ REAL"

        blink_info = f"👁️ Blinks: {int(blink[0,0])}, Freq: {blink[0,1]:.3f}, EAR Var: {blink[0,2]:.4f}"
        return f"{blink_info}\n\n**Prediction:** {label}\nConfidence: {score:.2f}"

    except Exception as e:
        return f"❌ Error processing video: {e}"

# ==================================================
# GRADIO INTERFACE
# ==================================================
demo = gr.Interface(
    fn=predict,
    inputs=gr.Video(label="🎥 Upload a short video (≤ 20 s)"),
    outputs=gr.Markdown(),
    title="Multimodal Deepfake Detection Demo (Docker)",
    description=(
        "Uploads a video, detects faces and blinks using dlib, "
        "and combines both to classify REAL vs FAKE. "
        "Optimized with Docker for instant startup."
    ),
)

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