import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torchvision.models import vgg16, VGG16_Weights
from PIL import Image
import gradio as gr
import numpy as np
import cv2
from gradcam import GradCAM, apply_heatmap
import base64
from io import BytesIO

# Categories for classification
categories = ["COVID", "Lung_Opacity", "Normal", "Viral Pneumonia"]

# Load Pretrained VGG16 Model
print("Loading Model...")
device = torch.device("cpu")
vgg_model = vgg16(weights=VGG16_Weights.IMAGENET1K_V1)
for param in vgg_model.features.parameters():
    param.requires_grad = False
vgg_model.classifier[6] = nn.Linear(4096, 4)
vgg_model.load_state_dict(torch.load("vgg_xray_model.pth", map_location=device))
vgg_model = vgg_model.to(device)
vgg_model.eval()
print("Model Loaded Successfully!")

# Initialize Grad-CAM with the last convolutional layer
gradcam = GradCAM(vgg_model, vgg_model.features[-1])

# Image Preprocessing
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# Prediction function with Grad-CAM
def predict(image):
    # Preprocess image (add batch dimension)
    image_tensor = transform(image).unsqueeze(0).to(device)
    
    # Perform inference first
    with torch.no_grad():
        output = vgg_model(image_tensor)
        _, predicted = torch.max(output, 1)
        prediction = categories[predicted.item()]
    
    # Determine status
    status = "COVID" if prediction in ["COVID", "Lung_Opacity"] else "Non-COVID"
    
    try:
        # Generate Grad-CAM heatmap with gradients
        with torch.set_grad_enabled(True):
            image_tensor.requires_grad_()
            heatmap = gradcam.generate_heatmap(image_tensor, predicted.item())
        
        # Generate overlay with bounding boxes
        overlay = apply_heatmap(
            image, 
            heatmap,
            min_contour_area=100,
            threshold=0.4
        )
        
    except Exception as e:
        print(f"Heatmap generation failed: {str(e)}")
        # Fallback to original image
        overlay = image.copy()

    # Return both JSON and Image outputs
    return (
        {"prediction": prediction, "status": status},  # First output (JSON)
        overlay                                        # Second output (Image)
    )

# Gradio Interface
print("Initializing Gradio Interface...")

interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.JSON(label="Diagnosis"), gr.Image(type="pil", label="Heatmap")],
    title="X-Ray Classifier with Grad-CAM",
    allow_flagging="never"
)

print("Launching Gradio App...")
interface.launch(server_name="0.0.0.0", server_port=7860)
print("Gradio App Running!")