# import pytorch related dependencies
import torch
from torch import nn
import numpy as np
import torchvision as torchvision
import torchvision.transforms as transforms
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
import gradio as gr

# model setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
classes = ['actinic keratoses', 'basal cell carcinoma', 'benign keratosis-like lesions',
           'dermatofibroma', 'melanoma', 'melanocytic nevi', 'vascular lesions']
model = torchvision.models.mobilenet_v3_large(pretrained=False)
model.classifier[3] = nn.Linear(1280, 7)
state_dict_trained = torch.load('checkpoints/ham10k_checkpoint_mobile_0.82_epoch24.pt',
                                map_location=torch.device('cpu'))
model.load_state_dict(state_dict_trained["model_state_dict"])
model.eval()

# image pre-processing
norm_mean = (0.4914, 0.4822, 0.4465)
norm_std = (0.2023, 0.1994, 0.2010)
transform = transforms.Compose([
    transforms.CenterCrop((400, 400)),
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std)
])


# convert tensot to numpy array
def tensor2npimg(tensor, mean, std):
    # inverse of normalization
    tensor = tensor.clone()
    mean_tensor = torch.as_tensor(list(mean), dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1)
    std_tensor = torch.as_tensor(list(std), dtype=tensor.dtype, device=tensor.device).view(-1, 1, 1)
    tensor.mul_(std_tensor).add_(mean_tensor)
    # convert tensor to numpy format for plt presentation
    npimg = tensor.numpy()
    npimg = np.transpose(npimg, (1, 2, 0))  # C*H*W => H*W*C
    return npimg


# draw Grad-CAM on image
# target layer could be any layer before the final attention block
# Some common choices are:
#   FasterRCNN: model.backbone
#   Resnet18 and 50: model.layer4[-1]
#   VGG and densenet161: model.features[-1]
#   mnasnet1_0: model.layers[-1]
#   ViT: model.blocks[-1].norm1
#   SwinT: model.layers[-1].blocks[-1].norm1
def image_grad_cam(model, input_tensor, input_float_np, target_layers):
    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
    grayscale_cam = cam(input_tensor=input_tensor, aug_smooth=True, eigen_smooth=True)
    grayscale_cam = grayscale_cam[0, :]
    return show_cam_on_image(input_float_np, grayscale_cam, use_rgb=True)


# config the predict function for Gradio, input type of image is numpy.nparray
def predict(gt, input_img):
    leasion_tensor = transform(input_img)
    input_float_np = tensor2npimg(leasion_tensor, norm_mean, norm_std)
    leasion_tensor = leasion_tensor.unsqueeze(dim=0)
    # predict
    with torch.no_grad():
        outputs = model(leasion_tensor)
    outputs = torch.exp(outputs)
    # probabilities of all classes
    pred_softmax = torch.softmax(outputs, dim=1).cpu().numpy()[0]
    # class with hightest probability
    pred = torch.argmax(outputs, dim=1).cpu().numpy()
    # diagnostic suggestions
    if pred == 1 or pred == 4:
        suggestion = "CHECK WITH YOUR MD!"
    else:
        suggestion = "Nothing to be worried about."
    # grad_cam image
    target_layers = model.features[-1]
    output_img = image_grad_cam(model, leasion_tensor, input_float_np, target_layers)
    # return label dict and suggestion
    return {classes[i]: float(pred_softmax[i]) for i in range(len(classes))}, suggestion, output_img


gr.Interface(
        fn=predict,
        inputs=[gr.Text(label="Ground Truth"), gr.Image(shape=(400, 400), type="pil", label="Image")],
        outputs=[gr.Label(label="Predict Result"), gr.Text(label="Recommendation", interactive=False), gr.Image(label="GradCAM")],
        examples=[['actinic keratoses', 'images/akiec.jpg'],
                  ['basal cell carcinoma', 'images/bcc.jpg'],
                  ['benign keratosis-like lesions', 'images/bkl.jpg'],
                  ['dermatofibroma', 'images/df.jpg'],
                  ['melanoma', 'images/mel.jpg'],
                  ['melanoma', 'images/mel2.jpg'],
                  ['melanoma', 'images/mel3.jpg'],
                  ['melanocytic nevi', 'images/nv.jpg'],
                  ['melanocytic nevi', 'images/nv2.jpg']],
        title="Skin Lesion Classifier"
      ).launch()