import gradio as gr
import numpy as np
import cv2
import onnxruntime as ort


# Load ONNX model
session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])


def preprocess_video(video_path):
    cap = cv2.VideoCapture(video_path)

    frames = []
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # Sample 16 frames uniformly
    indices = np.linspace(0, total_frames - 1, 16).astype(int)

    for idx in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ret, frame = cap.read()
        if not ret:
            continue
        frame = cv2.resize(frame, (224, 224))
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame = frame / 255.0
        frames.append(frame)

    cap.release()

    # Pad if less than 16 frames
    while len(frames) < 16:
        frames.append(frames[-1])

    frames = np.array(frames, dtype=np.float32)
    frames = np.transpose(frames, (0, 3, 1, 2))   # (16, H, W, C) → (16, C, H, W)
    frames = np.expand_dims(frames, axis=0)         # → (1, 16, C, H, W)

    return frames


def softmax(x):
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()


def predict(video):
    if video is None:
        return "No video uploaded.", "N/A"

    input_data = preprocess_video(video)

    outputs = session.run(None, {"input": input_data})
    logits = outputs[0]
    probs = softmax(logits[0])

    label = "REAL" if probs[1] > probs[0] else "FAKE"
    confidence = f"REAL: {probs[0]:.2%}  |  FAKE: {probs[1]:.2%}"

    return label, confidence


interface = gr.Interface(
    fn=predict,
    inputs=gr.Video(label="Upload Video"),
    outputs=[
        gr.Textbox(label="Prediction"),
        gr.Textbox(label="Confidence Scores")
    ],
    title="AI Generated Video Detector",
    description="Upload a video to classify it as REAL or FAKE."
)

interface.launch()