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 torch.utils.data import DataLoader
import itertools
import matplotlib.pyplot as plt

from custom_resnet import Custom_ResNet
import utils as utils

model = Custom_ResNet()
model.load_state_dict(torch.load("results/custom_resnet_trained.pth", map_location=torch.device('cpu')), strict=False)
model.eval()

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

cifar_valid = utils.Cifar10SearchDataset('.', train=False, download=True, transform=utils.augmentation_custom_resnet('Valid'))

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]
)

def inference(wants_gradcam, n_gradcam, target_layer_number, transparency, wants_misclassified, n_misclassified, input_img = None, n_top_classes=10):
    
    if wants_gradcam: 
      
      outputs_inference_gc = []
      cifar_valid_loader = DataLoader(cifar_valid, batch_size=1, shuffle = True)
      count_gradcam = 1

      for data, target in cifar_valid_loader:

        data, target = data.to('cpu'), target.to('cpu')
        target_layers = [model.layer2[target_layer_number]]

        cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
        grayscale_cam = cam(input_tensor=data, targets=None)
        grayscale_cam = grayscale_cam[0, :]

        org_img = inv_normalize(data).squeeze(0).numpy()
        org_img = np.transpose(org_img, (1, 2, 0))
        visualization = np.array(show_cam_on_image(org_img, grayscale_cam, use_rgb=True, image_weight=transparency))
        outputs_inference_gc.append(visualization)

        count_gradcam += 1
        if count_gradcam > n_gradcam:
          break 
    else:
      outputs_inference_gc = None

    if wants_misclassified: 
      outputs_inference_mis = []

      cifar_valid_loader = DataLoader(cifar_valid, batch_size=1, shuffle = True)
      count_mis = 1

      for data, target in cifar_valid_loader:

        data, target = data.to('cpu'), target.to('cpu')

        outputs = model(data)
        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 target.numpy()[0] != prediction.numpy()[0]: 
          
            count_mis += 1

            org_img = inv_normalize(data).squeeze(0).numpy()
            org_img = np.transpose(org_img, (1, 2, 0))
            
            fig = plt.figure()
            fig.add_subplot(111)

            plt.imshow(org_img)
            plt.title(f'Target: {classes[target.numpy()[0]]}\nPred: {classes[prediction.numpy()[0]]}')
            plt.axis('off')
            
            fig.canvas.draw()

            fig_img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
            fig_img = fig_img.reshape(fig.canvas.get_width_height()[::-1] + (3,))

            plt.close(fig)

            outputs_inference_mis.append(fig_img)

        if count_mis > n_misclassified:
            break 

    else:
      outputs_inference_mis = None

    if input_img is not None:
        transform=utils.augmentation_custom_resnet('Valid')
        org_img = input_img
        input_img = transform(image=input_img)
        input_img = input_img['image'].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)

        confidences = {k: v for k, v in sorted(confidences.items(), key=lambda item: item[1], reverse=True)}
        confidences = dict(itertools.islice(confidences.items(), n_top_classes))
    else:
      confidences = None
    

    return outputs_inference_gc, outputs_inference_mis, confidences

    
title = "CIFAR10 trained on Custom ResNet Model with GradCAM"
description = "A simple Gradio interface to infer on ResNet model, and get GradCAM results"
examples = [[None, None, None, None, None, None, 'examples/test_'+str(i)+'.jpg', None] for i in range(10)]

demo = gr.Interface(inference, 
                    inputs = [gr.Checkbox(False, label='Do you want to see GradCAM outputs?'),
                              gr.Slider(0, 10, value = 0, step=1, label="How many?"),
                              gr.Slider(-2, -1, value = -2, step=1, label="Which target layer?"), 
                              gr.Slider(0, 1, value = 0, label="Opacity of GradCAM"), 
                              gr.Checkbox(False, label='Do you want to see misclassified images?'),
                              gr.Slider(0, 10, value = 0, step=1, label="How many?"),
                              gr.Image(shape=(32, 32), label="Input image"), 
                              gr.Slider(0, 10, value = 0, step=1, label="How many top classes you want to see?")
                              ], 
                    outputs = [
                              gr.Gallery(label="GradCAM Outputs", show_label=True, elem_id="gallery").style(columns=[2], rows=[2], object_fit="contain", height="auto"), 
                              gr.Gallery(label="Misclassified Images", show_label=True, elem_id="gallery").style(columns=[2], rows=[2], object_fit="contain", height="auto"), 
                              gr.Label(num_top_classes=None)
                              ],
                    title = title, 
                    description = description, 
                    examples = examples
                    )
demo.launch()