from ultralytics import YOLO
import torch
import cv2
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from pytorch_grad_cam import EigenCAM
from pytorch_grad_cam.utils.image import show_cam_on_image, scale_cam_image
import gradio as gr

# Global Color Palette
COLORS = np.random.uniform(0, 255, size=(80, 3))

def parse_detections(results):
    boxes, colors, names = [], [], []
    for result in results:
        # Accessing boxes directly from the result
        for box in result.boxes:
            xmin, ymin, xmax, ymax = box.xyxy[0].int().tolist()  # Convert to list of integers
            category = int(box.cls[0].item())  # Class index
            name = result.names[category]  # Get class name from names
            boxes.append((xmin, ymin, xmax, ymax))
            colors.append(COLORS[category])  # Ensure COLORS is defined elsewhere in your code
            names.append(name)

    return boxes, colors, names

def draw_detections(boxes, colors, names, img):
    for box, color, name in zip(boxes, colors, names):
        xmin, ymin, xmax, ymax = box
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, 2)
        cv2.putText(img, name, (xmin, ymin - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2,
                    lineType=cv2.LINE_AA)
    return img


def generate_cam_image(model, target_layers, tensor, rgb_img, boxes):
    cam = EigenCAM(model, target_layers)
    grayscale_cam = cam(tensor)[0, :, :]
    img_float = np.float32(rgb_img) / 255
    cam_image = show_cam_on_image(img_float, grayscale_cam, use_rgb=True)
    renormalized_cam = np.zeros(grayscale_cam.shape, dtype=np.float32)
    for x1, y1, x2, y2 in boxes:
        renormalized_cam[y1:y2, x1:x2] = scale_cam_image(grayscale_cam[y1:y2, x1:x2].copy())
    renormalized_cam = scale_cam_image(renormalized_cam)
    renormalized_cam_image = show_cam_on_image(img_float, renormalized_cam, use_rgb=True)

    return cam_image, renormalized_cam_image


def xai_yolov8n(image):
    model = YOLO('yolov8n.pt')  # Load YOLOv8n pre-trained weights
    model.eval()

    # Check if GPU is available and use it
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    target_layers = [model.model.model[-2]]  # Grad-CAM target layer
    
    # Process the image through the model
    results = model([image])
    
    # If results are a list, extract the first element (detected results)
    if isinstance(results, list):
        results = results[0]  # Extracting the first result (if list)
    
    # Ensure that outputs are in tensor form
    logits = results.pred[0]  # Get the prediction tensor from the results

    # Parse the detections
    boxes, colors, names = parse_detections([results])  # Ensure results are passed as a list
    detections_img = draw_detections(boxes, colors, names, image.copy())
    
    # Prepare image for Grad-CAM
    img_float = np.float32(image) / 255
    transform = transforms.ToTensor()
    tensor = transform(img_float).unsqueeze(0).to(device)  # Ensure tensor is on the right device
    
    # Generate CAM images
    cam_image, renormalized_cam_image = generate_cam_image(model, target_layers, tensor, image, boxes)

    # Combine original image, CAM image, and renormalized CAM image
    final_image = np.hstack((image, cam_image, renormalized_cam_image))

    # Return final image and a caption
    caption = "Results using YOLOv8n"
    return Image.fromarray(final_image), caption