import gradio as gr
import torch
import torchvision.transforms as transforms
import numpy as np
from PIL import Image
import cv2
import os
from sklearn.svm import SVC
from scipy.fftpack import fft2, fftshift
import pretrainedmodels

# 🔌 Transform pour CNN
transform = transforms.Compose([
    transforms.Resize((299, 299)),
    transforms.ToTensor(),
    transforms.Normalize([0.5]*3, [0.5]*3)
])

class ProxyXceptionNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = pretrainedmodels.__dict__["xception"](num_classes=1000, pretrained="imagenet")
        self.model.eval()

    def predict(self, image_pil):
        img = transform(image_pil).unsqueeze(0)
        with torch.no_grad():
            out = self.model(img)
            prob = torch.softmax(out, dim=1).numpy()
        # Simule une décision deepfake basée sur la classe dominante : purement démonstratif !
        pred_class = np.argmax(prob)
        return int(pred_class % 2 == 0)  # pseudo-détection arbitraire

cnn_model = ProxyXceptionNet()

# ⚖️ SVM avec FFT (entraînement simple sur données factices)
fft_train = np.random.rand(20, 1000)
y_train = [0]*10 + [1]*10
svm = SVC(probability=True)
svm.fit(fft_train, y_train)

# 📈 FFT Feature Extraction
def extract_fft_features_pil(image):
    img = image.convert("L").resize((256, 256))
    img_np = np.array(img)
    f = fft2(img_np)
    fshift = fftshift(f)
    magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1)
    return magnitude_spectrum.flatten()[:1000]

# 🔄 Fonction principale Gradio
def predict(image, method):
    if method == "CNN (Xception)":
        result = cnn_model.predict(image)
        label = "Deepfake ❌" if result == 1 else "Réel ✅"
    else:
        features = extract_fft_features_pil(image)
        pred = svm.predict([features])[0]
        label = "Deepfake ❌" if pred == 1 else "Réel ✅"

    return label

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Image à analyser"),
        gr.Radio(["CNN (Xception)", "FFT + SVM"], label="Méthode de détection")
    ],
    outputs=gr.Textbox(label="Résultat"),
    title="Détection de Deepfakes (Image)",
    description="Upload une image et choisis la méthode pour tester si c'est un deepfake."
)

demo.launch(share=True)