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
from tqdm import tqdm
from PIL import Image
import uuid
import tempfile
import pandas as pd
from numpy import exp
import numpy as np
from sklearn.metrics import ConfusionMatrixDisplay
import urllib.request

# Define models
models = [
    "umm-maybe/AI-image-detector",
    "Organika/sdxl-detector",
    "cmckinle/sdxl-flux-detector",
]

pipe0 = pipeline("image-classification", f"{models[0]}")
pipe1 = pipeline("image-classification", f"{models[1]}")
pipe2 = pipeline("image-classification", f"{models[2]}")

fin_sum = []
uid = uuid.uuid4()

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

# Function to extract images from zip
def extract_zip(zip_file):
    temp_dir = tempfile.mkdtemp()  # Temporary directory
    with zipfile.ZipFile(zip_file, 'r') as z:
        z.extractall(temp_dir)
    return temp_dir

# Function to classify images in a folder
def classify_images(image_dir, model_pipeline):
    images = []
    labels = []
    preds = []
    for folder_name, ground_truth_label in [('real', 1), ('ai', 0)]:
        folder_path = os.path.join(image_dir, folder_name)
        if not os.path.exists(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")
                pred = model_pipeline(img)
                pred_label = np.argmax([x['score'] for x in pred])
                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}")
    return labels, preds, images

# Function to generate evaluation metrics
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, ax = plt.subplots()
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=["AI", "Real"])
    disp.plot(cmap=plt.cm.Blues, ax=ax)
    plt.close(fig)

    fig_roc, ax_roc = plt.subplots()
    ax_roc.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
    ax_roc.plot([0, 1], [0, 1], color='gray', linestyle='--')
    ax_roc.set_xlim([0.0, 1.0])
    ax_roc.set_ylim([0.0, 1.05])
    ax_roc.set_xlabel('False Positive Rate')
    ax_roc.set_ylabel('True Positive Rate')
    ax_roc.set_title('Receiver Operating Characteristic (ROC) Curve')
    ax_roc.legend(loc="lower right")
    plt.close(fig_roc)

    return accuracy, roc_score, report, fig, fig_roc

# Gradio function for batch image processing
def process_zip(zip_file):
    extracted_dir = extract_zip(zip_file.name)
    labels, preds, images = classify_images(extracted_dir, pipe0)  # You can switch to pipe1 or pipe2
    accuracy, roc_score, report, cm_fig, roc_fig = evaluate_model(labels, preds)
    shutil.rmtree(extracted_dir)  # Clean up extracted files
    return accuracy, roc_score, report, cm_fig, roc_fig

# Single image classification functions
def image_classifier0(image):
    labels = ["AI", "Real"]
    outputs = pipe0(image)
    results = {}
    for idx, result in enumerate(outputs):
        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
    fin_sum.append(results)
    return results

def image_classifier1(image):
    labels = ["AI", "Real"]
    outputs = pipe1(image)
    results = {}
    for idx, result in enumerate(outputs):
        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
    fin_sum.append(results)
    return results

def image_classifier2(image):
    labels = ["AI", "Real"]
    outputs = pipe2(image)
    results = {}
    for idx, result in enumerate(outputs):
        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
    fin_sum.append(results)
    return results

def aiornot0(image):
    labels = ["AI", "Real"]
    mod = models[0]
    feature_extractor0 = AutoFeatureExtractor.from_pretrained(mod)
    model0 = AutoModelForImageClassification.from_pretrained(mod)
    input = feature_extractor0(image, return_tensors="pt")
    with torch.no_grad():
        outputs = model0(**input)
        logits = outputs.logits
        probability = softmax(logits)  # Apply softmax on logits
        px = pd.DataFrame(probability.numpy())
    prediction = logits.argmax(-1).item()
    label = labels[prediction]

    html_out = f"""
    <h1>This image is likely: {label}</h1><br><h3>
    Probabilities:<br>
    Real: {float(px[1][0]):.4f}<br>
    AI: {float(px[0][0]):.4f}"""
    
    results = {
        "Real": float(px[1][0]),
        "AI": float(px[0][0])
    }
    fin_sum.append(results)
    return gr.HTML.update(html_out), results

def aiornot1(image):
    labels = ["AI", "Real"]
    mod = models[1]
    feature_extractor1 = AutoFeatureExtractor.from_pretrained(mod)
    model1 = AutoModelForImageClassification.from_pretrained(mod)
    input = feature_extractor1(image, return_tensors="pt")
    with torch.no_grad():
        outputs = model1(**input)
        logits = outputs.logits
        probability = softmax(logits)  # Apply softmax on logits
        px = pd.DataFrame(probability.numpy())
    prediction = logits.argmax(-1).item()
    label = labels[prediction]

    html_out = f"""
    <h1>This image is likely: {label}</h1><br><h3>
    Probabilities:<br>
    Real: {float(px[1][0]):.4f}<br>
    AI: {float(px[0][0]):.4f}"""
    
    results = {
        "Real": float(px[1][0]),
        "AI": float(px[0][0])
    }
    fin_sum.append(results)
    return gr.HTML.update(html_out), results

def aiornot2(image):
    labels = ["AI", "Real"]
    mod = models[2]
    feature_extractor2 = AutoFeatureExtractor.from_pretrained(mod)
    model2 = AutoModelForImageClassification.from_pretrained(mod)
    input = feature_extractor2(image, return_tensors="pt")
    with torch.no_grad():
        outputs = model2(**input)
        logits = outputs.logits
        probability = softmax(logits)  # Apply softmax on logits
        px = pd.DataFrame(probability.numpy())
    prediction = logits.argmax(-1).item()
    label = labels[prediction]

    html_out = f"""
    <h1>This image is likely: {label}</h1><br><h3>
    Probabilities:<br>
    Real: {float(px[1][0]):.4f}<br>
    AI: {float(px[0][0]):.4f}"""
    
    results = {
        "Real": float(px[1][0]),
        "AI": float(px[0][0])
    }
    fin_sum.append(results)
    return gr.HTML.update(html_out), results

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

def tot_prob():
    try:
        fin_out = sum([result["Real"] for result in fin_sum]) / len(fin_sum)
        fin_sub = 1 - fin_out
        out = {
            "Real": f"{fin_out:.4f}",
            "AI": f"{fin_sub:.4f}"
        }
        return out
    except Exception as e:
        print(e)
        return None

def fin_clear():
    fin_sum.clear()
    return None

def upd(image):
    rand_im = uuid.uuid4()
    image.save(f"{rand_im}-vid_tmp_proc.png")
    out = Image.open(f"{rand_im}-vid_tmp_proc.png")
    return out

# Set up Gradio app
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():
        # Tab for single image detection
        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")
                mes = gr.HTML("""""")

            with gr.Group():
                with gr.Row():
                    fin = gr.Label(label="Final Probability")
                with gr.Row():
                    for i, model in enumerate(models):
                        with gr.Box():
                            gr.HTML(f"""<b>Testing on Model {i}: <a href='https://huggingface.co/{model}'>{model}</a></b>""")
                            globals()[f'outp{i}'] = gr.HTML("""""")
                            globals()[f'n_out{i}'] = gr.Label(label="Output")

            btn.click(fin_clear, None, fin, show_progress=False)
            load_btn.click(load_url, in_url, [inp, mes])

            btn.click(aiornot0, [inp], [outp0, n_out0]).then(
                aiornot1, [inp], [outp1, n_out1]).then(
                aiornot2, [inp], [outp2, n_out2]).then(
                tot_prob, None, fin, show_progress=False)

            btn.click(image_classifier0, [inp], [n_out0]).then(
                image_classifier1, [inp], [n_out1]).then(
                image_classifier2, [inp], [n_out2]).then(
                tot_prob, None, fin, show_progress=False)

        # Tab for batch processing
        with gr.Tab("Batch Image Processing"):
            zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
            output_acc = gr.Label(label="Accuracy")
            output_roc = gr.Label(label="ROC Score")
            output_report = gr.Textbox(label="Classification Report", lines=10)
            output_cm = gr.Plot(label="Confusion Matrix")
            output_roc_plot = gr.Plot(label="ROC Curve")

            batch_btn = gr.Button("Process Batch")

    # Connect batch processing
    batch_btn.click(process_zip, zip_file, [output_acc, output_roc, output_report, output_cm, output_roc_plot])

app.launch(show_api=False, max_threads=24)