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
import torch
from torchvision import transforms 
from model import CustomResNet
from utils.utils import wrong_predictions
from utils.dataloader import get_dataloader
import random 
from collections import OrderedDict
import os

test_o = get_dataloader()
# test_o=next(iter(test_o))   


examples_dir = os.path.join(os.getcwd(), 'examples')
examples = [[os.path.join(examples_dir, img), 0.5] for img in os.listdir(examples_dir)]


model = CustomResNet()
model.load_state_dict(torch.load('modelp.ckpt', map_location='cpu')['state_dict']) #, strict = False)
# model = model.cpu()

classes = ['plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
norm_mean=(0.4914, 0.4822, 0.4465) 
norm_std=(0.2023, 0.1994, 0.2010)
misclassified_images, all_predictions = wrong_predictions(model,test_o, norm_mean, norm_std, classes, 'cpu')

# layers = ['layer_1', 'layer_3']
# layers = [model.layer_1, model.layer_2, model.layer_3]
def inference(input_img, transparency, layer_num, top_classes):
    input_img_ori = input_img.copy()
    transform = transforms.ToTensor()
    # transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(
    #     mean=[0.485,0.456,0.406],
    #     std=[0.229, 0.224, 0.255]
    # )])
    inv_normalize = transforms.Normalize(
        mean=[-0.485/0.229, -0.456/0.224, -0.406/0.255],
        std=[1/0.229, 1/0.224, 1/0.255]
    )
    input_img = transform(input_img)
    # input_img = input_img.to(device)
    input_img = input_img.unsqueeze(0)
    outputs = model(input_img)
    _, prediction = torch.max(outputs, 1)
    softmax = torch.nn.Softmax(dim=0)
    outputs = softmax(outputs.flatten())
    # print(outputs)
    confidences = {classes[i]: float(outputs[i]) for i in range(10)}
    confidences = OrderedDict(sorted(confidences.items(), key=lambda x:x[1], reverse=True))
    # print(confidences)
    filtered_confidences ={}# OrderedDict()
    for i, (key, val) in enumerate(confidences.items()):
        if i ==  top_classes:
            break
        filtered_confidences[key] = val


    if layer_num == 1:
        target_layers = [model.layer_1]
    elif layer_num == 2:
        target_layers = [model.layer_2]
    else:
        target_layers = [model.layer_3]
    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
    grayscale_cam = cam(input_tensor=input_img, targets=None)
    grayscale_cam = grayscale_cam[0, :]
    img = input_img.squeeze(0)
    img = inv_normalize(img)
    rgb_img = np.transpose(img, (1, 2, 0))
    rgb_img = np.array(np.clip(rgb_img,0,1), np.float32)
    visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=transparency)

    # visualization = input_img_ori
    return filtered_confidences, visualization
    # return filtered_confidences, superimposed_img

def get_misclassified_images(num):
    outputimgs = []
    # misclassified_images = wrong_predictions(model,test_o, norm_mean, norm_std, classes, 'cpu')
    for i in range(int(num)):
        # misclassified_images[0][0].cpu().numpy()
        inv_normalize = transforms.Normalize(
            mean=[-0.485/0.229, -0.456/0.224, -0.406/0.255],
            std=[1/0.229, 1/0.224, 1/0.255]
        )
        inv_tensor = np.array(inv_normalize(misclassified_images[random.randint(2,98)][0]).cpu().permute(1,2,0)*255, dtype='uint8')
        outputimgs.append(inv_tensor)
    return outputimgs


def get_gradcam_images(num, transparency, layer_num):
    outcoms=[]
    for i in range(int(num)):
        input_img = all_predictions[random.randint(2,98)][0]
        inv_normalize = transforms.Normalize(
            mean=[-0.485/0.229, -0.456/0.224, -0.406/0.255],
            std=[1/0.229, 1/0.224, 1/0.255]
        )
        input_img = input_img.unsqueeze(0)
        if layer_num == 1:
            target_layers = [model.layer_1]
        elif layer_num == 2:
            target_layers = [model.layer_2]
        else:
            target_layers = [model.layer_3]

        cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
        grayscale_cam = cam(input_tensor=input_img, targets=None)
        grayscale_cam = grayscale_cam[0, :]
        img = input_img.squeeze(0)
        img = inv_normalize(img)
        rgb_img = np.transpose(img, (1, 2, 0))
        rgb_img = np.array(np.clip(rgb_img,0,1), np.float32)
        visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=transparency)
        outcoms.append(visualization)
    return outcoms

# demo = gr.Interface(inference, [gr.Image(shape=(32, 32)), gr.Slider(0, 1)], ["text", gr.Image(shape=(32, 32)).style(width=128, height=128)])
inference_new_image = gr.Interface(
    inference, 
    inputs = [gr.Image(shape=(32, 32), label="Input Image"), gr.Slider(0, 1, value = 0.3, label="transparency?"), gr.Slider(1, 3, value = 1,step=1, label="layer?"),
              gr.Slider(1, 10, value = 3, step=1, label="top classes?")], 
    
    outputs = [gr.Label(),gr.Image(shape=(32, 32), label="Model Prediction").style(width=300, height=300)],
    title = 'gradio app',
    description = 'for dl purposes',
    examples = examples,
)

misclassified_interface = gr.Interface(
    get_misclassified_images, 
    inputs = [gr.Number(value=10, label="images number")], 
    
    outputs = [gr.Gallery(label="misclassified images")],
    title = 'gradio app',
    description = 'for dl purposes'
)

gradcam_images = gr.Interface(
    get_gradcam_images, 
    inputs = [gr.Number(value=10, label="images number"), gr.Slider(0, 1, value = 0.3, label="transparency?"), gr.Slider(1, 3, value = 1,step=1, label="layer?")], 
    
    outputs = [gr.Gallery(label="gradcam images")],
    title = 'gradio app',
    description = 'for dl purposes'
)


demo = gr.TabbedInterface([inference_new_image, misclassified_interface, gradcam_images], tab_names=["Input image", "Misclassified Images", "grad cam images"], 
                          title="customresnet gradcam")
demo.launch()