import torch, torchvision
from torchvision import transforms
import numpy as np
import gradio as gr
from PIL import Image
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from resnet18 import ResNet18
import gradio as gr

model = ResNet18()
model.load_state_dict(torch.load("resnet_model.pth", map_location=torch.device('cpu')), strict=False)

inv_normalize = transforms.Normalize(
    mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],
    std=[1/0.23, 1/0.23, 1/0.23]
)
classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

def resize_image_pil(image, new_width, new_height):

    # Convert to PIL image
    img = Image.fromarray(np.array(image))
    
    # Get original size
    width, height = img.size

    # Calculate scale
    width_scale = new_width / width
    height_scale = new_height / height 
    scale = min(width_scale, height_scale)

    # Resize
    resized = img.resize((int(width*scale), int(height*scale)), Image.NEAREST)
    
    # Crop to exact size
    resized = resized.crop((0, 0, new_width, new_height))

    return resized

def inference(input_img, is_grad_cam=True, transparency = 0.5,layer='layer2', target_layer_number = -1, top_predictions=3):
    input_img = resize_image_pil(input_img, 32, 32)

    input_img = np.array(input_img)
    org_img = input_img
    input_img = input_img.reshape((32, 32, 3))
    transform = transforms.ToTensor()
    input_img = transform(input_img)
    input_img = input_img
    input_img = input_img.unsqueeze(0)
    outputs = model(input_img)
    softmax = torch.nn.Softmax(dim=0)
    o = softmax(outputs.flatten())
    confidences = {classes[i]: float(o[i]) for i in range(10)}
    _, prediction = torch.max(outputs, 1)
    if is_grad_cam:
        if layer == 'layer1':
          target_layers = [model.layer1[target_layer_number]]
          cam = GradCAM(model=model, target_layers=target_layers)
          grayscale_cam = cam(input_tensor=input_img, targets=None)
          grayscale_cam = grayscale_cam[0, :]
          visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency)
        
        if layer == 'layer2':
          target_layers = [model.layer2[target_layer_number]]
          cam = GradCAM(model=model, target_layers=target_layers)
          grayscale_cam = cam(input_tensor=input_img, targets=None)
          grayscale_cam = grayscale_cam[0, :]
          visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency)

        if layer == 'layer3':
          target_layers = [model.layer3[target_layer_number]]
          cam = GradCAM(model=model, target_layers=target_layers)
          grayscale_cam = cam(input_tensor=input_img, targets=None)
          grayscale_cam = grayscale_cam[0, :]
          visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency)

        if layer == 'layer4':
          target_layers = [model.layer4[target_layer_number]]
          cam = GradCAM(model=model, target_layers=target_layers)
          grayscale_cam = cam(input_tensor=input_img, targets=None)
          grayscale_cam = grayscale_cam[0, :]
          visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency)
    else:
        visualization = None

    # Sort the confidences dictionary based on confidence values
    sorted_confidences = dict(sorted(confidences.items(), key=lambda item: item[1], reverse=True))
    
    # Pick the top n predictions
    top_n_confidences = dict(list(sorted_confidences.items())[:top_predictions])
    
    return classes[prediction[0].item()], visualization, top_n_confidences

title = "CIFAR10 trained on ResNet18 Model with GradCAM"
description = "A simple Gradio interface to infer on ResNet model, and get GradCAM results"
examples = [
    ["examples/" + "cat.jpeg", True, 0.5, -1, 3],
    ["examples/" + "horse.jpg", True, 0.5, -1, 3],
    ["examples/" + "bird.webp", True, 0.5, -1, 3],
    ["examples/" + "dog1.jpg", True, 0.5, -1, 3],
    ["examples/" + "frog1.webp", True, 0.5, -1, 3],
    ["examples/" + "deer.webp", True, 0.5, -1, 3],
    ["examples/" + "airplane.png", True, 0.5, -1, 3],
    ["examples/" + "shipp.jpg",True, 0.5, -1, 3],
    ["examples/" + "car.jpg", True, 0.5, -1, 3],
    ["examples/"+ "truck1.jpg", True, 0.5, -1, 3],
]
demo = gr.Interface(
    inference, 
    inputs = [
        gr.Image(width=256, height=256, label="Input Image"),
        gr.Checkbox(label="Show GradCAM"),
        gr.Slider(0, 1, value = 0.5, label="Overall Opacity of Image"), 
        gr.Dropdown(["layer1", "layer2", "layer3", "layer4"], label="Layer"),
        gr.Slider(-2, -1, value = -2, step=1, label="Which Layer?"),
        gr.Slider(2, 10, value=3, step=1, label="Number of Top Classes")
        ], 
    outputs = [
        "text", 
        gr.Image(width=256, height=256, label="Output"),
        gr.Label(label="Top Classes")
        ],
    title = title,
    description = description,
    examples = examples,
)
demo.launch()