import warnings
warnings.filterwarnings('ignore')
import torch, numpy as np, os
from torch import nn
from transformers import AutoModelForImageClassification, AutoConfig, AutoImageProcessor
import matplotlib.pyplot as plt
from PIL import Image
import saliency.core as saliency
import io
import gradio as gr
import PIL

model_choice = 3
model_names = ["nvidia/mit-b0",'facebook/convnext-base-224', 'microsoft/resnet-18', 'microsoft/swin-tiny-patch4-window7-224']
model_name = model_names[model_choice]
device     = 'cuda' if torch.cuda.is_available() else 'cpu'

class Model(nn.Module):
    def __init__(self, MODEL_NAME=model_name):
        super().__init__()
        self.config = AutoConfig.from_pretrained(MODEL_NAME, finetuning_task="image-classification")
        self.model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
        self.class_len = self.config.num_labels
        self.id2label = self.config.id2label
        self.label2id = self.config.label2id

    def forward(self, x):
        if isinstance(x, np.ndarray):
            x = torch.from_numpy(x)
        if len(x.shape) == 3:
            x = x.unsqueeze(0)
        if x.shape[-1] == 3:
            x = x.permute(0, 3, 1, 2)
        x = x.to(device)
        x = self.model(x)
        return x.logits

def conv_layer_forward_hook(module, input, output):
    """Method from Examples_pytorch.ipynb for the gradcam library https://github.com/PAIR-code/saliency."""
    global last_conv_layer_outputs
    last_conv_layer_outputs[saliency.base.CONVOLUTION_LAYER_VALUES] = torch.movedim(output, 3, 1).detach().cpu().numpy()
def conv_layer_backward_hook(module, grad_input, grad_output):
    """Method from Examples_pytorch.ipynb for the gradcam library https://github.com/PAIR-code/saliency."""
    global last_conv_layer_outputs
    last_conv_layer_outputs[saliency.base.CONVOLUTION_OUTPUT_GRADIENTS] = torch.movedim(grad_output[0], 3, 1).detach().cpu().numpy()

auto_transformer, class_to_id, id_to_class, last_conv_layer, last_conv_layer_outputs = None, None, None, None, None


def swap_models(name):
    global model, auto_transformer, class_to_id, id_to_class, last_conv_layer, last_conv_layer_outputs
    auto_transformer = AutoImageProcessor.from_pretrained(name)
    model = Model(MODEL_NAME=name)
    model = model.to(device).eval()
    # register the hooks for the last convolution layer for Grad-Cam
    named_modules = dict(model.model.named_modules())
    last_conv_layer_name = None
    for name, module in named_modules.items():
        if isinstance(module, torch.nn.Conv2d):
            last_conv_layer_name = name

    last_conv_layer = named_modules[last_conv_layer_name]
    last_conv_layer_outputs = {}

    last_conv_layer.register_forward_hook(conv_layer_forward_hook)
    last_conv_layer.register_backward_hook(conv_layer_backward_hook)
    class_to_id = {v:k for k,v in model.model.config.id2label.items()}
    id_to_class = {k:v for k,v in model.model.config.id2label.items()}

swap_models(model_name)

def saliency_graph(img1, steps=25):
    img1 = auto_transformer(img1)
    img1 = np.squeeze(np.array(img1.pixel_values))
    if img1.shape[0] < img1.shape[1]:
        img1 = np.moveaxis(img1, 0, -1)
    img1 = (img1 - np.min(img1)) / (np.max(img1) - np.min(img1))

    class_idx_str = 'class_idx_str'
    def gradcam_call(images, call_model_args=None, expected_keys=None):
        if not isinstance(images, np.ndarray) and not isinstance(images, torch.Tensor) and not isinstance(images, PIL.Image.Image):
            # return two blank images
            im1 = np.zeros((224, 224, 3))
            im2 = np.zeros((224, 224, 3))
            return im1, im2
        
        if len(images.shape) == 3:
            images = np.expand_dims(images, 0)
        images = torch.tensor(images, dtype=torch.float32)
        images = images.requires_grad_(True)
        target_class_idx = call_model_args[class_idx_str]
        y_pred = model(images)
        
        if saliency.base.INPUT_OUTPUT_GRADIENTS in expected_keys:
            out =  y_pred[:, target_class_idx]
            # move actual color channel to the 1st dimension
            #images = torch.movedim(images, 3, 1)
            grads = torch.autograd.grad(out, images, grad_outputs=torch.ones_like(out))
            grads = grads[0].detach().cpu().numpy()
            return {saliency.base.INPUT_OUTPUT_GRADIENTS: grads}
        else:
            hot = torch.zeroes_like(y_pred)
            hot[:, target_class_idx] = 1
            model.zero_grad()
            y_pred.backward(gradient=hot, retain_graph=True)
            return last_conv_layer_outputs

    im = img1.astype(np.float32)
    base = np.zeros(img1.shape)

    pred = model(torch.from_numpy(im))
    class_pred = pred.argmax(dim=1).item()
    call_model_args = {class_idx_str: class_pred}
    gradients = saliency.IntegratedGradients()

    s = gradients.GetSmoothedMask(im, gradcam_call, call_model_args, x_steps=steps, x_baseline=base, batch_size=25)

    smoothgrad_mask_grayscale = saliency.VisualizeImageGrayscale(s)

    with torch.no_grad():
        output = model.forward(img1)
        output = torch.nn.functional.softmax(output, dim=1)
        output = output.cpu().numpy()
    top_5 = [(id_to_class[int(i)], output[0][i]) for i in np.argsort(output)[0][-5:][::-1]]


    # Render the saliency masks.
    fig, ax = plt.subplots(1, 1, figsize=(10, 10))
    ax.barh([x[0] for x in top_5], [x[1] for x in top_5])
    ax.set_title('Top 5 Predictions')
    buf = io.BytesIO()
    fig.savefig(buf, format='jpg')
    buf.seek(0)
    fig_img = Image.open(buf)
    plt.close(fig)
    return smoothgrad_mask_grayscale, fig_img

# gradio Interface
def gradio_interface(img):
    smoothgrad_mask_grayscale, fig_img = saliency_graph(img, steps=20)
    return smoothgrad_mask_grayscale, fig_img

with gr.Blocks(live=True) as iface:
    #examples = gr.Examples(examples=["ex1.jpg", "ex2.jpg", "ex3.jpg", "ex4.jpg"], label="Examples", inputs="image", examples_per_page=4)
    gr.Markdown("This function finds the most critical pixels in an image for predicting a class by looking at the pixels models attend to. The best models will ideally make predictions by highlighting the expected object. Poorly generalizable models will often rely on environmental cues instead and forego looking at the most important pixels. Highlighting the most important pixels helps explain/build trust about whether a given model uses the correct features to make its prediction.")
    with gr.Row():
        with gr.Column():
            test_image = gr.Image(label="Input Image", live=True)
            input_btn = gr.Button(label="Classify image")
            model_select_dropdown = gr.Radio(model_names, label="Model to test", interactive=True, default=3)
        with gr.Column():
            output = gr.Image(label="Pixels used for classification")
            output2 = gr.Image(label="Top 5 Predictions")

    input_btn.click(gradio_interface, test_image, outputs=[output, output2])
    model_select_dropdown.change(swap_models, inputs=[model_select_dropdown])
    examples = gr.Examples(
        examples = [os.path.join('./', x) for x in os.listdir('./') if x.endswith('.jpg')],
        inputs=gr.Image(),
        label="Examples",
        fn=gradio_interface,
        cache_examples=True,
        run_on_click=True,
        postprocess=True,
        preprocess=True,
        outputs=[output, output2])


iface.launch()