import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# Load model
print("Loading DeepGuard model...", flush=True)
model = tf.keras.models.load_model("deepfake_model.keras")
print("Model loaded successfully!", flush=True)


def generate_fft(image):
    """Generate FFT magnitude spectrum visualization."""
    try:
        img_gray = image.convert('L').resize((128, 128))
        img_array = np.array(img_gray, dtype=np.float32)
        f_transform = np.fft.fft2(img_array)
        f_shift = np.fft.fftshift(f_transform)
        magnitude = np.abs(f_shift)
        magnitude = np.log1p(magnitude)
        magnitude = ((magnitude - magnitude.min()) / (magnitude.max() - magnitude.min()) * 255).astype(np.uint8)
        return Image.fromarray(magnitude)
    except Exception as e:
        print(f"FFT Error: {e}")
        return None


def predict(image):
    """Main prediction function."""
    if image is None:
        return "Please upload an image", None
    
    try:
        img = image.convert('RGB').resize((128, 128))
        img_array = np.array(img) / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        
        prediction = float(model.predict(img_array, verbose=0)[0][0])
        
        if prediction > 0.5:
            label = "FAKE (AI Generated)"
            confidence = prediction * 100
        else:
            label = "REAL (Authentic)"
            confidence = (1 - prediction) * 100
        
        result_text = f"""## Detection Result: {label}

**Confidence:** {confidence:.2f}%

**Raw Model Score:** {prediction:.6f}

---

### Interpretation
- Score > 0.5: Model detects GAN artifacts = FAKE
- Score < 0.5: No GAN artifacts detected = REAL

### Important Note
This model is trained on StyleGAN-generated faces and may not accurately detect images from modern diffusion models (Stable Diffusion, Midjourney, DALL-E).
"""
        
        fft_image = generate_fft(image)
        return result_text, fft_image
        
    except Exception as e:
        return f"Error: {str(e)}", None


# Simple Interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Face Image"),
    outputs=[
        gr.Markdown(label="Detection Result"),
        gr.Image(label="FFT Frequency Analysis")
    ],
    title="DeepGuard - AI Face Authenticator",
    description="Detect AI-generated (deepfake) faces using deep learning. Upload a face image for analysis.",
    article="""
### Model Info
- **Architecture:** XceptionTransfer
- **Training Data:** 140,000 real and GAN-generated faces
- **Accuracy:** 88% on test dataset

### FFT Interpretation
| Pattern | Meaning |
|---------|---------|
| Bright center | Normal low-frequency content |
| Radiating spokes | Edge directions in image |
| Grid artifacts | Potential GAN fingerprint |

### Limitations
This model is trained on StyleGAN faces only. It may not detect Stable Diffusion, Midjourney, or DALL-E images.
""",
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()