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 numpy as np
from utils.goat import call_inference
import io
import warnings

# 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, device=device)

# Load additional models
models = ["Organika/sdxl-detector", "cmckinle/sdxl-flux-detector"]
feature_extractor_3 = AutoFeatureExtractor.from_pretrained(models[0], device=device)
model_3 = AutoModelForImageClassification.from_pretrained(models[0]).to(device)
feature_extractor_4 = AutoFeatureExtractor.from_pretrained(models[1], device=device)
model_4 = AutoModelForImageClassification.from_pretrained(models[1]).to(device)

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

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

def convert_pil_to_bytes(image, format='JPEG'):
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format=format)
    img_byte_arr = img_byte_arr.getvalue()
    return img_byte_arr

@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}
        print(result_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"AI, Confidence: {result_1['artificial']:.4f}"
        elif result_1['real'] >= confidence_threshold:
            label_1 = f"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}
        print(result_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"AI, Confidence: {result_2['AI Image']:.4f}"
        elif result_2['Real Image'] >= confidence_threshold:
            label_2 = f"Real, 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 with softmax
    try:
        inputs_3 = feature_extractor_3(img_pil, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs_3 = model_3(**inputs_3)
            logits_3 = outputs_3.logits
            probabilities_3 = softmax(logits_3.cpu().numpy()[0])
        result_3 = {
            labels_3[0]: float(probabilities_3[0]),  # AI
            labels_3[1]: float(probabilities_3[1])   # Real
        }
        print(result_3)
        # Ensure the result dictionary contains all class names
        for class_name in labels_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"AI, Confidence: {result_3['AI']:.4f}"
        elif result_3['Real'] >= confidence_threshold:
            label_3 = f"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 with softmax
    try:
        inputs_4 = feature_extractor_4(img_pil, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs_4 = model_4(**inputs_4)
            logits_4 = outputs_4.logits
            probabilities_4 = softmax(logits_4.cpu().numpy()[0])
        result_4 = {
            labels_4[0]: float(probabilities_4[0]),  # AI
            labels_4[1]: float(probabilities_4[1])   # Real
        }
        print(result_4)
        # Ensure the result dictionary contains all class names
        for class_name in labels_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"AI, Confidence: {result_4['AI']:.4f}"
        elif result_4['Real'] >= confidence_threshold:
            label_4 = f"Real, Confidence: {result_4['Real']:.4f}"
        else:
            label_4 = "Uncertain Classification"
    except Exception as e:
        label_4 = f"Error: {str(e)}"
    
    try:
        img_bytes = convert_pil_to_bytes(img_pil)
        response5_raw = call_inference(img_bytes)
        response5 = response5_raw.json()
        print(response5)
        label_5 = f"Result: {response5}"
    except Exception as e:
        label_5 = f"Error: {str(e)}"
    
    # Combine results
    combined_results = {
        "SwinV2/detect": label_1,
        "ViT/AI-vs-Real": label_2,
        "Swin/SDXL": label_3,
        "Swin/SDXL-FLUX": label_4,
        "GOAT": label_5
    }
    return img_pil, combined_results

# Define a function to generate the HTML content
def generate_results_html(results):
    html_content = f"""
    <link href="https://stackpath.bootstrapcdn.com/bootstrap/4.3.1/css/bootstrap.min.css" rel="stylesheet">
    <div class="container">
        <div class="row mt-4">
            <div class="col">
                <h5>SwinV2/detect</h5>
                <p>{results.get("SwinV2/detect", "N/A")}</p>
            </div>
            <div class="col">
                <h5>ViT/AI-vs-Real</h5>
                <p>{results.get("ViT/AI-vs-Real", "N/A")}</p>
            </div>
            <div class="col">
                <h5>Swin/SDXL</h5>
                <p>{results.get("Swin/SDXL", "N/A")}</p>
            </div>
            <div class="col">
                <h5>Swin/SDXL-FLUX</h5>
                <p>{results.get("Swin/SDXL-FLUX", "N/A")}</p>
            </div>
            <div class="col">
                <h5>GOAT</h5>
                <p>{results.get("GOAT", "N/A")}</p>
            </div>
        </div>
    </div>
    """
    return html_content

# Modify the predict_image function to return the HTML content
def predict_image_with_html(img, confidence_threshold):
    img_pil, results = predict_image(img, confidence_threshold)
    html_content = generate_results_html(results)
    return img_pil, html_content

# Define the Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# AI Generated Image Classification")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Image to Analyze", sources=['upload'], type='pil')
            confidence_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Confidence Threshold")
            inputs = [image_input, confidence_slider]
        with gr.Column():
            image_output = gr.Image(label="Processed Image")
            # Custom HTML component to display results in 5 columns
            results_html = gr.HTML(label="Model Predictions")
            outputs = [image_output, results_html]
    
    gr.Button("Predict").click(fn=predict_image_with_html, inputs=inputs, outputs=outputs)

# Launch the interface
iface.launch()