import spaces
import gradio as gr
from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassification, AutoFeatureExtractor, AutoModelForImageClassification
from torchvision import transforms
import torch
from PIL import Image
import pandas as pd
import warnings
import math

# Suppress warnings
warnings.filterwarnings("ignore", category=UserWarning, message="Using a slow image processor as `use_fast` is unset")

# Ensure using GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load the first model and processor
image_processor_1 = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy", use_fast=True)
model_1 = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
model_1 = model_1.to(device)
clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device)

# Load the second model
model_2_path = "Heem2/AI-vs-Real-Image-Detection"
clf_2 = pipeline("image-classification", model=model_2_path)

# Load additional models
models = ["Organika/sdxl-detector", "cmckinle/sdxl-flux-detector"]
pipe0 = pipeline("image-classification", model=models[0])
pipe1 = pipeline("image-classification", model=models[1])

# Define class names for all models
class_names_1 = ['artificial', 'real']
class_names_2 = ['AI Image', 'Real Image']
class_names_3 = ['AI', 'Real']
class_names_4 = ['AI', 'Real']

def softmax(vector):
    e = math.exp(vector - vector.max())  # for numerical stability
    return e / e.sum()

@spaces.GPU(duration=10)
def predict_image(img, confidence_threshold):
    # Ensure the image is a PIL Image
    if not isinstance(img, Image.Image):
        raise ValueError(f"Expected a PIL Image, but got {type(img)}")
    
    # Convert the image to RGB if not already
    if img.mode != 'RGB':
        img_pil = img.convert('RGB')
    else:
        img_pil = img
    
    # Resize the image
    img_pil = transforms.Resize((256, 256))(img_pil)
    
    # Predict using the first model
    try:
        prediction_1 = clf_1(img_pil)
        result_1 = {pred['label']: pred['score'] for pred in prediction_1}
        
        # Ensure the result dictionary contains all class names
        for class_name in class_names_1:
            if class_name not in result_1:
                result_1[class_name] = 0.0
        
        # Check if either class meets the confidence threshold
        if result_1['artificial'] >= confidence_threshold:
            label_1 = f"Label: artificial, Confidence: {result_1['artificial']:.4f}"
        elif result_1['real'] >= confidence_threshold:
            label_1 = f"Label: real, Confidence: {result_1['real']:.4f}"
        else:
            label_1 = "Uncertain Classification"
    except Exception as e:
        label_1 = f"Error: {str(e)}"
    
    # Predict using the second model
    try:
        prediction_2 = clf_2(img_pil)
        result_2 = {pred['label']: pred['score'] for pred in prediction_2}
        
        # Ensure the result dictionary contains all class names
        for class_name in class_names_2:
            if class_name not in result_2:
                result_2[class_name] = 0.0
        
        # Check if either class meets the confidence threshold
        if result_2['AI Image'] >= confidence_threshold:
            label_2 = f"Label: AI Image, Confidence: {result_2['AI Image']:.4f}"
        elif result_2['Real Image'] >= confidence_threshold:
            label_2 = f"Label: Real Image, Confidence: {result_2['Real Image']:.4f}"
        else:
            label_2 = "Uncertain Classification"
    except Exception as e:
        label_2 = f"Error: {str(e)}"
    
    # Predict using the third model
    try:
        prediction_3 = pipe0(img_pil)
        result_3 = {}
        for idx, result in enumerate(prediction_3):
            result_3[class_names_3[idx]] = float(result['score'])
        
        # Ensure the result dictionary contains all class names
        for class_name in class_names_3:
            if class_name not in result_3:
                result_3[class_name] = 0.0
        
        # Check if either class meets the confidence threshold
        if result_3['AI'] >= confidence_threshold:
            label_3 = f"Label: AI, Confidence: {result_3['AI']:.4f}"
        elif result_3['Real'] >= confidence_threshold:
            label_3 = f"Label: Real, Confidence: {result_3['Real']:.4f}"
        else:
            label_3 = "Uncertain Classification"
    except Exception as e:
        label_3 = f"Error: {str(e)}"
    
    # Predict using the fourth model
    try:
        prediction_4 = pipe1(img_pil)
        result_4 = {}
        for idx, result in enumerate(prediction_4):
            result_4[class_names_4[idx]] = float(result['score'])
        
        # Ensure the result dictionary contains all class names
        for class_name in class_names_4:
            if class_name not in result_4:
                result_4[class_name] = 0.0
        
        # Check if either class meets the confidence threshold
        if result_4['AI'] >= confidence_threshold:
            label_4 = f"Label: AI, Confidence: {result_4['AI']:.4f}"
        elif result_4['Real'] >= confidence_threshold:
            label_4 = f"Label: Real, Confidence: {result_4['Real']:.4f}"
        else:
            label_4 = "Uncertain Classification"
    except Exception as e:
        label_4 = f"Error: {str(e)}"
    
    # Combine results
    combined_results = {
        "SwinV2": label_1,
        "AI-vs-Real-Image-Detection": label_2,
        "Organika/sdxl-detector": label_3,
        "cmckinle/sdxl-flux-detector": label_4
    }
    
    return combined_results

# Define the Gradio interface
image = gr.Image(label="Image to Analyze", sources=['upload'], type='pil')  # Ensure the image type is PIL
confidence_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Confidence Threshold")
label = gr.JSON(label="Model Predictions")

# Launch the interface
iface = gr.Interface(
    fn=predict_image,
    inputs=[image, confidence_slider],
    outputs=label,
    title="AI Generated Classification"
)
iface.launch()