import torch
import dataset
import albumentations
from utils import get_misclassified_data
from albumentations.pytorch import ToTensorV2
from visualize import display_cifar_misclassified_data
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 resnet import ResNet18
import gradio as gr

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

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

# dataloader arguments - something you'll fetch these from cmdprmt
dataloader_args = dict(shuffle=True, batch_size=128, num_workers=4, pin_memory=True) if cuda else dict(shuffle=True, batch_size=64)

test_loader = dataset.get_test_data_loader(**dataloader_args)

classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

# Get the misclassified data from test dataset
misclassified_data = get_misclassified_data(model, device, test_loader)

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, target_layer_number = -1,
              top_predictions=3, is_misclassified_images=True, num_misclassified_images=10):
    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))
    transforms = albumentations.Compose(
        # Normalize
        [albumentations.Normalize([0.49139968, 0.48215841, 0.44653091], 
                                  [0.24703223, 0.24348513, 0.26158784]), 
        # Convert to tensor                          
        ToTensorV2()])
    input_img = transforms(image = input_img)['image']
    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:
        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)
    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])

    if is_misclassified_images:
        # Plot the misclassified data
        misclassified_images = display_cifar_misclassified_data(misclassified_data, number_of_samples=num_misclassified_images)
    else:
        misclassified_images = None
    
    return classes[prediction[0].item()], visualization, top_n_confidences, misclassified_images

title = "CIFAR10 trained on ResNet18 Model with GradCAM"
description = "A simple Gradio interface to infer on ResNet model, and get GradCAM results"
examples = [["cat.jpg", True, 0.5, -1, 3, True, 10], 
            ["dog.jpg", True, 0.5, -1, 3, True, 10], 
            ["bird.jpg", True, 0.5, -1, 3, True, 10], 
            ["car.jpg", True, 0.5, -1, 3, True, 10], 
            ["deer.jpg", True, 0.5, -1, 3, True, 10], 
            ["frog.jpg", True, 0.5, -1, 3, True, 10], 
            ["horse.jpg", True, 0.5, -1, 3, True, 10], 
            ["plane.jpg", True, 0.5, -1, 3, True, 10], 
            ["ship.jpg", True, 0.5, -1, 3, True, 10], 
            ["truck.jpg", True, 0.5, -1, 3, True, 10]]
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.Slider(-2, -1, value = -2, step=1, label="Which Layer?"),
        gr.Slider(2, 10, value=3, step=1, label="Number of Top Classes"),
        gr.Checkbox(label="Show Misclassified Images"),
        gr.Slider(5, 40, value=10, step=5, label="Number of Misclassified Images")
        ], 
    outputs = [
        "text", 
        gr.Image(width=256, height=256, label="Output"),
        gr.Label(label="Top Classes"),
        gr.Plot(label="Misclassified Images")
        ],
    title = title,
    description = description,
    examples = examples,
)
demo.launch()