import gradio as gr
import torch
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
import os
import zipfile
import shutil
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, roc_auc_score, confusion_matrix, classification_report, roc_curve, auc, ConfusionMatrixDisplay
from PIL import Image
import tempfile
import numpy as np
import urllib.request

MODEL_NAME = "cmckinle/sdxl-flux-detector"
LABELS = ["AI", "Real"]

class AIDetector:
    def __init__(self):
        self.pipe = pipeline("image-classification", MODEL_NAME)
        self.feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
        self.model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)

    @staticmethod
    def softmax(vector):
        e = np.exp(vector - np.max(vector))
        return e / e.sum()

    def predict(self, image):
        inputs = self.feature_extractor(image, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = self.softmax(logits.numpy())
        
        prediction = logits.argmax(-1).item()
        label = LABELS[prediction]
        
        results = {label: float(prob) for label, prob in zip(LABELS, probabilities[0])}
        
        return label, results

def process_zip(zip_file):
    temp_dir = tempfile.mkdtemp()
    with zipfile.ZipFile(zip_file.name, 'r') as z:
        z.extractall(temp_dir)
    
    labels, preds, images = [], [], []
    detector = AIDetector()
    
    for folder_name, ground_truth_label in [('real', 1), ('ai', 0)]:
        folder_path = os.path.join(temp_dir, folder_name)
        if not os.path.exists(folder_path):
            print(f"Folder not found: {folder_path}")
            continue
        for img_name in os.listdir(folder_path):
            img_path = os.path.join(folder_path, img_name)
            try:
                img = Image.open(img_path).convert("RGB")
                _, prediction = detector.predict(img)
                pred_label = 0 if prediction["AI"] > prediction["Real"] else 1
                
                preds.append(pred_label)
                labels.append(ground_truth_label)
                images.append(img_name)
            except Exception as e:
                print(f"Error processing image {img_name}: {e}")
    
    shutil.rmtree(temp_dir)
    return evaluate_model(labels, preds)

def evaluate_model(labels, preds):
    cm = confusion_matrix(labels, preds)
    accuracy = accuracy_score(labels, preds)
    roc_score = roc_auc_score(labels, preds)
    report = classification_report(labels, preds)
    fpr, tpr, _ = roc_curve(labels, preds)
    roc_auc = auc(fpr, tpr)

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=LABELS).plot(cmap=plt.cm.Blues, ax=ax1)
    ax1.set_title("Confusion Matrix")
    
    ax2.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
    ax2.plot([0, 1], [0, 1], color='gray', linestyle='--')
    ax2.set_xlim([0.0, 1.0])
    ax2.set_ylim([0.0, 1.05])
    ax2.set_xlabel('False Positive Rate')
    ax2.set_ylabel('True Positive Rate')
    ax2.set_title('ROC Curve')
    ax2.legend(loc="lower right")
    
    plt.tight_layout()

    return accuracy, roc_score, report, fig

def load_url(url):
    try:
        urllib.request.urlretrieve(url, "temp_image.png")
        image = Image.open("temp_image.png")
        message = "Image Loaded"
    except Exception as e:
        image = None
        message = f"Image not Found<br>Error: {e}"
    return image, message

detector = AIDetector()

def create_gradio_interface():
    with gr.Blocks() as app:
        gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Test Demo - accuracy varies by model)</h4></h1></center>""")

        with gr.Tabs():
            with gr.Tab("Single Image Detection"):
                with gr.Column():
                    inp = gr.Image(type='pil')
                    in_url = gr.Textbox(label="Image URL")
                    with gr.Row():
                        load_btn = gr.Button("Load URL")
                        btn = gr.Button("Detect AI")
                    message = gr.HTML()

                with gr.Group():
                    with gr.Box():
                        gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{MODEL_NAME}'>{MODEL_NAME}</a></b>""")
                        output_html = gr.HTML()
                        output_label = gr.Label(label="Output")

            with gr.Tab("Batch Image Processing"):
                zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
                batch_btn = gr.Button("Process Batch")

                with gr.Group():
                    gr.Markdown(f"### Results for {MODEL_NAME}")
                    output_acc = gr.Label(label="Accuracy")
                    output_roc = gr.Label(label="ROC Score")
                    output_report = gr.Textbox(label="Classification Report", lines=10)
                    output_plots = gr.Plot(label="Confusion Matrix and ROC Curve")

        load_btn.click(load_url, in_url, [inp, message])
        btn.click(
            lambda img: detector.predict(img),
            inp,
            [output_html, output_label]
        )

        batch_btn.click(
            process_zip,
            zip_file,
            [output_acc, output_roc, output_report, output_plots]
        )

    return app

if __name__ == "__main__":
    app = create_gradio_interface()
    app.launch(show_api=False, max_threads=24)