import torch
from matplotlib import pyplot as plt
from PIL import Image
import numpy as np
from numpy import matlib as mb
import torchvision.transforms as transforms
import torchvision.models as models
import skimage.transform
import gradio as gr


def compute_spatial_similarity(conv1, conv2):
  """
  Takes in the last convolutional layer from two images, computes the pooled output
  feature, and then generates the spatial similarity map for both images.
  """
  conv1 = conv1.reshape(-1, 7*7).T
  conv2 = conv2.reshape(-1, 7*7).T

  pool1 = np.mean(conv1, axis=0)
  pool2 = np.mean(conv2, axis=0)
  out_sz = (int(np.sqrt(conv1.shape[0])),int(np.sqrt(conv1.shape[0])))
  conv1_normed = conv1 / np.linalg.norm(pool1) / conv1.shape[0]
  conv2_normed = conv2 / np.linalg.norm(pool2) / conv2.shape[0]
  im_similarity = np.zeros((conv1_normed.shape[0], conv1_normed.shape[0]))

  for zz in range(conv1_normed.shape[0]):
    repPx = mb.repmat(conv1_normed[zz,:],conv1_normed.shape[0],1)
    im_similarity[zz,:] = np.multiply(repPx,conv2_normed).sum(axis=1)
  similarity1 = np.reshape(np.sum(im_similarity,axis=1),out_sz)
  similarity2 = np.reshape(np.sum(im_similarity,axis=0),out_sz)
  return similarity1, similarity2

# Get Layer 4

display_transform = transforms.Compose([transforms.Resize(256), transforms.CenterCrop((224, 224))])

imagenet_transform = transforms.Compose(
  [transforms.Resize(256),
   transforms.CenterCrop((224, 224)), 
   transforms.ToTensor(), 
   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])


class Wrapper(torch.nn.Module):
  def __init__(self, model):
    super(Wrapper, self).__init__()
    self.model = model
    self.layer4_ouputs = None
    def fw_hook(module, input, output):
      self.layer4_ouputs = output
    self.model.layer4.register_forward_hook(fw_hook)

  def forward(self, input):
    _ = self.model(input)
    return self.layer4_ouputs
    
  def __repr__(self):
    return "Wrapper"
  
def get_layer4(input_image):
  l4_model = models.resnet50(pretrained=True)
  l4_model = l4_model.cuda()
  l4_model.eval();
  wrapped_model = Wrapper(l4_model)

  with torch.no_grad():
    data = imagenet_transform(input_image).unsqueeze(0)
    data = data.cuda()
    reference_layer4 = wrapped_model(data)

  return reference_layer4.data.to('cpu').numpy()

# Visualization
def visualize_similarities(image1, image2):
  a = get_layer4(image1).squeeze()
  b = get_layer4(image2).squeeze()
  sim1, sim2 = compute_spatial_similarity(a, b)

  fig, axes = plt.subplots(1, 2, figsize=(12, 5))
  axes[0].imshow(display_transform(image1))
  im1=axes[0].imshow(skimage.transform.resize(sim1, (224, 224)), alpha=0.6, cmap='jet')
  # axes[0].colorbar()

  axes[1].imshow(display_transform(image2))
  im2=axes[1].imshow(skimage.transform.resize(sim2, (224, 224)), alpha=0.6, cmap='jet')
  # axes[1].colorbar()

  fig.colorbar(im1, ax=axes[0])
  fig.colorbar(im2, ax=axes[1])
  plt.tight_layout()
  return fig

# GRADIO APP
iface = gr.Interface(fn=visualize_similarities, 
                     inputs=[gr.inputs.Image(shape=(300, 300), type='pil'), 
                             gr.inputs.Image(shape=(300, 300), type='pil')], outputs="plot")
iface.launch()