import gradio as gr
import torch
import torch.nn as nn
import numpy as np
from torchvision import transforms
from transformers import ViTForImageClassification
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from PIL import Image

# Device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Transforms
test_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

mean = np.array([0.485, 0.456, 0.406])
std  = np.array([0.229, 0.224, 0.225])

# Load model
model = ViTForImageClassification.from_pretrained(
    'google/vit-base-patch16-224',
    num_labels=2,
    ignore_mismatched_sizes=True
)
model.load_state_dict(torch.load('vit_pneumonia_best.pth', map_location=device))
model = model.to(device)
model.eval()

# ViT wrapper for Grad-CAM
class ViTWrapper(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    def forward(self, x):
        return self.model(x).logits

def reshape_transform(tensor, height=14, width=14):
    result = tensor[:, 1:, :].reshape(tensor.size(0), height, width, tensor.size(2))
    result = result.transpose(2, 3).transpose(1, 2)
    return result

wrapped_model = ViTWrapper(model)
target_layers = [model.vit.encoder.layer[-1].layernorm_before]
cam = GradCAM(model=wrapped_model, target_layers=target_layers,
              reshape_transform=reshape_transform)

# Prediction function
def predict(pil_image):
    img_tensor = test_transforms(pil_image).unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(img_tensor).logits
        probs  = torch.softmax(output, dim=1)[0]

    grayscale_cam = cam(input_tensor=img_tensor)[0]

    img_numpy = test_transforms(pil_image).permute(1, 2, 0).numpy()
    img_numpy = std * img_numpy + mean
    img_numpy = np.clip(img_numpy, 0, 1).astype(np.float32)

    heatmap = show_cam_on_image(img_numpy, grayscale_cam, use_rgb=True)

    label = {
        "NORMAL":    float(probs[0]),
        "PNEUMONIA": float(probs[1])
    }

    return label, Image.fromarray(heatmap)

# Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Chest X-Ray"),
    outputs=[
        gr.Label(label="Prediction"),
        gr.Image(type="pil", label="Grad-CAM Heatmap")
    ],
    title="🫁 Pneumonia Detector",
    description="Upload a chest X-ray image to detect pneumonia with explainable AI (Grad-CAM)",
)

demo.launch()