import numpy as np
import gradio as gr
import imageio
import cv2
import kornia as K
import kornia.geometry as KG
from copy import deepcopy
from tqdm import tqdm
from base64 import b64encode
import torch
import torch.nn.functional as F

use_cuda: bool = torch.cuda.is_available()
device = torch.device('cuda' if use_cuda else 'cpu')
registrator = KG.ImageRegistrator('similarity', 
                                  loss_fn = F.mse_loss, 
                                  lr=8e-4, pyramid_levels=3, num_iterations=500).to(device)

models = []

def check_image_sizes(f_names):
  sizes = []
  for f_name in f_names:
    img = cv2.imread(f_name, cv2.IMREAD_COLOR)
    sizes.append(img.shape)

  return len(set(sizes))==1


def convert_img(f_name):
    img = cv2.imread(f_name, cv2.IMREAD_COLOR)
    # convert image to torch tensor                                                                  
    tensor = K.image_to_tensor(img, None).float() / 255.
    return K.color.bgr_to_rgb(tensor)

def merge_sharp1_into2(timg1, timg2, trans1to2, verbose=False):
    curr_img = timg2.clone()
    warped = KG.homography_warp(timg1, torch.inverse(trans1to2), timg1.shape[-2:])
    mask1 = K.filters.laplacian(K.color.rgb_to_grayscale(timg1), 7).abs()
    mask1_norm = (mask1-mask1.min()) / (mask1.max() - mask1.min())
    mask1_blur = K.filters.gaussian_blur2d(mask1_norm, (9,9), (1.6, 1.6))
    mask1_blur = mask1_blur / mask1_blur.max()
    warped_mask = KG.homography_warp(mask1_blur.float(), torch.inverse(trans1to2), timg1.shape[-2:])
    curr_img = warped_mask * warped + (1-warped_mask) * curr_img
    return curr_img

def img_registration(images):
  f_names = [f.name for f in images]
  # print(check_image_sizes(f_names))
  # if not check_image_sizes(f_names):
  #   return "All images should be same size!"

  for i, f_name in tqdm(enumerate(f_names)):
      if i == 0:
          continue
      prev_img = convert_img(f_names[i-1]).to(device)
      curr_img = convert_img(f_name).to(device)
      model = registrator.register(prev_img, curr_img)
      models.append(deepcopy(model.detach()))

  models_to_final = [torch.eye(3, device=device)[None]]
  for m in models[::-1]:
      models_to_final.append(m @ models_to_final[-1])
  models_to_final = models_to_final[::-1]

  base_img = convert_img(f_names[-1])
  curr_img = deepcopy(base_img)
  _, layers, height, width  = curr_img.shape
  video_file = 'video.avi'
  video = cv2.VideoWriter(video_file, 0, 1, (width,height))

  with torch.no_grad():
      for i, image in tqdm(enumerate(f_names)):
          timg = convert_img(image)
          curr_img = merge_sharp1_into2(timg.to(device), curr_img.to(device), models_to_final[i].to(device))
          video.write(cv2.cvtColor(K.tensor_to_image(curr_img.float()*255).astype(np.uint8), cv2.COLOR_BGR2RGB))
  video.release()

  return K.tensor_to_image(curr_img.float()), video_file

title = 'Image Registration with Kornia!'
description = '''Image registration is the process of transforming different sets of data into one coordinate system. Data may be multiple photographs, data from different sensors, times, depths, or viewpoints. It is used in computer vision, medical imaging, and compiling and analyzing images and data from satellites. Registration is necessary in order to be able to compare or integrate the data obtained from these different measurements.

*Note that you can upload only image files, e.g. jpg, png etc and all images should have same width and height!* 

Learn more about [image registraion and Kornia](https://kornia.readthedocs.io/en/latest/applications/image_registration.html)'''

examples = [["IMG_3020.JPG", "IMG_3027.JPG", "IMG_3034.JPG", "IMG_3040.JPG", "IMG_3058.JPG", "IMG_3070.JPG", "IMG_3083.JPG", "IMG_3100.JPG", "IMG_3106.JPG", "IMG_3112.JPG"]]

iface = gr.Interface(
    img_registration, 
    inputs='files', 
    outputs=["image", gr.Video()], 
    allow_flagging="never",
    title=title,
    description=description
    )

if __name__ == "__main__":
  iface.launch(show_error=True)