import gradio as gr
import torch
from transformers import AutoFeatureExtractor, AutoModelForImageClassification, pipeline
import os 
from numpy import exp
import pandas as pd
from PIL import Image
import urllib.request 
import uuid
uid = uuid.uuid4()

# Reordered models as requested
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 = []

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

# Image classification function for Model 0
def image_classifier0(image):
    fin_sum.clear()  # Clear previous results
    labels = ["AI", "Real"]
    outputs = pipe0(image)
    scores = [output['score'] for output in outputs]
    soft_scores = softmax(scores)  # Ensure consistency with softmax
    results = {labels[i]: float(soft_scores[i]) for i in range(len(labels))}
    fin_sum.append(results)
    return results

# Image classification function for Model 1
def image_classifier1(image):
    fin_sum.clear()  # Clear previous results
    labels = ["AI", "Real"]
    outputs = pipe1(image)
    scores = [output['score'] for output in outputs]
    soft_scores = softmax(scores)  # Ensure consistency with softmax
    results = {labels[i]: float(soft_scores[i]) for i in range(len(labels))}
    fin_sum.append(results)
    return results

# Image classification function for Model 2
def image_classifier2(image):
    fin_sum.clear()  # Clear previous results
    labels = ["AI", "Real"]
    outputs = pipe2(image)
    scores = [output['score'] for output in outputs]
    soft_scores = softmax(scores)  # Ensure consistency with softmax
    results = {labels[i]: float(soft_scores[i]) for i in range(len(labels))}
    fin_sum.append(results)
    return results

# AI or Not function for Model 0
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])}<br>
    AI: {float(px[0][0])}"""
    
    results = {
        "Real": float(px[1][0]),
        "AI": float(px[0][0])
    }
    fin_sum.append(results)
    return gr.HTML.update(html_out), results

# AI or Not function for Model 1
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])}<br>
    AI: {float(px[0][0])}"""
    
    results = {
        "Real": float(px[1][0]),
        "AI": float(px[0][0])
    }
    fin_sum.append(results)
    return gr.HTML.update(html_out), results

# AI or Not function for Model 2
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])}<br>
    AI: {float(px[0][0])}"""
    
    results = {
        "Real": float(px[1][0]),
        "AI": float(px[0][0])
    }
    fin_sum.append(results)
    return gr.HTML.update(html_out), results

# Load URL and return image
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

# Calculate final probabilities
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}",
            "AI": f"{fin_sub}"
        }
        return out
    except Exception as e:
        print(e)
        return None

# Clear the fin_sum list
def fin_clear():
    fin_sum.clear()
    return None

# Update image
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

with gr.Blocks() as app:
    gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Test Demo - accuracy varies by model)""")
    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", visible=False)
        with gr.Row():
            with gr.Box():
                lab0 = gr.HTML(f"""<b>Testing on Original Model: <a href='https://huggingface.co/{models[0]}'>{models[0]}</a></b>""")
                nun0 = gr.HTML("""""")
            with gr.Box():
                lab1 = gr.HTML(f"""<b>Testing on SDXL Fine Tuned Model: <a href='https://huggingface.co/{models[1]}'>{models[1]}</a></b>""")
                nun1 = gr.HTML("""""")
            with gr.Box():
                lab2 = gr.HTML(f"""<b>Testing on SDXL and Flux Fine Tuned Model: <a href='https://huggingface.co/{models[2]}'>{models[2]}</a></b>""")
                nun2 = gr.HTML("""""")
        with gr.Row():
            with gr.Box():
                n_out0 = gr.Label(label="Output")
                outp0 = gr.HTML("""""")
            with gr.Box():
                n_out1 = gr.Label(label="Output")
                outp1 = gr.HTML("""""")
            with gr.Box():
                n_out2 = gr.Label(label="Output")
                outp2 = gr.HTML("""""")

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

    # Use either the aiornot functions or image_classifier consistently
    btn.click(aiornot0, [inp], [outp0, n_out0]).then(tot_prob, None, fin, show_progress=False)
    btn.click(aiornot1, [inp], [outp1, n_out1]).then(tot_prob, None, fin, show_progress=False)
    btn.click(aiornot2, [inp], [outp2, n_out2]).then(tot_prob, None, fin, show_progress=False)

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