import gradio as gr
import torch
from torchvision import transforms as T
from PIL import Image
import numpy as np
import base64
from io import BytesIO
import matplotlib.pyplot as plt
from src.train import load_model, loss_fn

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model = load_model("./models/autoencoder_mnist.pth")
model.eval()

transform = T.Compose([
    T.Grayscale(num_output_channels=1),
    T.Resize((28, 28)),
    T.ToTensor()
])

def predict(image):
    original = image.convert("L")
    tensor = transform(original).unsqueeze(0).to(DEVICE)
    original_resized = original.resize((200, 200), Image.NEAREST)

    with torch.no_grad():
        decoded = model(tensor)
        score = loss_fn(decoded, tensor).item()
        diff = (decoded - tensor).squeeze().cpu().numpy() ** 2
        reconstructed = decoded.squeeze().cpu().numpy()

    # Anomaly interpretation
    if score > 0.02:
        verdict = "🚨 Anomaly detected!"
    else:
        verdict = "✅ Everything looks normal."

    # Create heatmap
    fig, ax = plt.subplots()
    ax.imshow(diff, cmap="hot")
    ax.axis("off")
    buf1 = BytesIO()
    plt.savefig(buf1, format="png", bbox_inches="tight", pad_inches=0)
    plt.close(fig)
    buf1.seek(0)
    heatmap = Image.open(buf1)

    # Reconstructed image
    fig2, ax2 = plt.subplots()
    ax2.imshow(reconstructed, cmap="gray")
    ax2.axis("off")
    buf2 = BytesIO()
    plt.savefig(buf2, format="png", bbox_inches="tight", pad_inches=0)
    plt.close(fig2)
    buf2.seek(0)
    reconstructed_img = Image.open(buf2)

    return f"{score:.5f}", verdict, original_resized, reconstructed_img, heatmap


iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload MNIST-style Image"),
    outputs=[
        gr.Textbox(label="Anomaly Score"),
        gr.Textbox(label="Verdict"),
        gr.Image(label="Original", height=200),
        gr.Image(label="Reconstruction", height=200),
        gr.Image(label="Heatmap", height=200),
    ],
    allow_flagging="never",
    title="🧠 Anomaly Detection",
    description="Upload a grayscale image (28x28) to get anomaly score."
)

iface.launch(iface.launch(server_name="0.0.0.0", server_port=7860))