# Copyright 2020 Adobe. All rights reserved.
# This file is licensed to you under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You may obtain a copy
# of the License at http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software distributed under
# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
# OF ANY KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
import torch
import torch.nn.functional as F
from torch.nn import init
import torchvision.transforms as transforms

import cv2
import imageio
import os
import sys
import numpy as np
import matplotlib.pyplot as plt


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def rescale_img(img, output_size=256):
    """
    Returns version of image `img` resized to a square image of specified `output_size`.
    """
    return cv2.resize(img, dsize=(output_size, output_size))

def crop_img(img, crop_size=700):
    """
    Returns version of image `img` cropped to specified output square size `crop_size` x `crop_size`.
    """
    H, W, _ = img.shape
    y_off, x_off = (H - crop_size) // 2, (W - crop_size) // 2
    return img[y_off:-y_off, x_off:-x_off]

def flow_warp(x, flow, padding_mode='border'):
    """
    Warps an image or feature map with optical flow.
    Arguments:
        `x` (Tensor): size (n, c, h, w)
        `flow` (Tensor): size (n, 2, h, w), values range from -1 to 1 (relevant to image width or height)
        `padding_mode` (str): 'zeros' or 'border'

    Returns:
        Tensor: warped image or feature map according to `flow`

    Code borrowed from https://github.com/hellock/cvbase/issues/4.
    """
    assert x.size()[-2:] == flow.size()[-2:]
    n, _, h, w = x.size()
    x_ = torch.arange(w).view(1, -1).expand(h, -1)
    y_ = torch.arange(h).view(-1, 1).expand(-1, w)
    grid = torch.stack([x_, y_], dim=0).float()
    grid = grid.unsqueeze(0).expand(n, -1, -1, -1)
    grid[:, 0, :, :] = 2 * grid[:, 0, :, :] / (w - 1) - 1
    grid[:, 1, :, :] = 2 * grid[:, 1, :, :] / (h - 1) - 1
    grid = grid + flow
    grid = grid.permute(0, 2, 3, 1)
    return F.grid_sample(x, grid, padding_mode=padding_mode)

def plot_quiver(dense_flow, path, save=True, overlaid=False):
    """
    Plots quiver field (arrows with direction, magnitude) for generated warping field.

    `dense_flow` is the upsampled warping field generated by the model.
    `path` is the path to which the quiver plot should be saved.
    `save` flag determines whether to save (or simply just to plot).
    `overlaid` determines whether the plot is being overlaid on another figure.
    """
    X, Y = np.meshgrid(np.arange(0, 256, 6), np.arange(256, 0, -6))
    if overlaid:
        X, Y = np.meshgrid(np.arange(0, 256, 6), np.arange(0, 256, 6))
    U = -dense_flow[0, 0, :, :][::6, ::6].detach().cpu()
    V = dense_flow[0, 1, :, :][::6, ::6].detach().cpu()

    fig, ax = plt.subplots()
    fig.set_size_inches(10, 10)
    plt.axis('off')
    # ax.set_title('Warp Flow - AutoToon')  # uncomment this line to set title of quiver plot
    ax.quiver(X, Y, U, V, units='width', width=0.007, color='white')
    Q = ax.quiver(X, Y, U, V, units='width', width=0.007)
    
    if save:
        plt.savefig(path, bbox_inches='tight', pad_inches=0)

def save_images(model, save_dir, img=None, img_path=None, img_name=None, scale=1, visualize=True):
    """
    Saves six images depicting results of the model.

    The following model outputs are produced and written to the output directory `save_dir`:
    - IMG_NAME_orig.jpg, the original cropped and centered input image,
    - IMG_NAME_out.jpg, the output of the AutoToon model for the input image,
    - IMG_NAME_quiver.jpg, the warping field quiver plot generated by the model,
    - IMG_NAME_overlaid.jpg, the overlaid warping field quiver plot on top of the output image,
    - IMG_NAME_xflow.jpg, the visualized heatmap for the x-direction warping field with the overlaid quiver plot,
    - IMG_NAME_yflow.jpg, the visualized heatmap for the y-direction warping field with the overlaid quiver plot.
    """
    assert img != None or img_path != None  # must provide either an image or an image path to load
    if img:
        assert img_name  # must provide image name if image object is provided

    # normalize tensor values with stats from AutoToon training dataset
    transformation = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5072122 , 0.4338291 , 0.38466746], 
                                                                [0.08460502, 0.07533269, 0.07269038])])
    # pre-process image
    in_img = None
    if img:  # image takes precedent; otherwise load image from path
        in_img = np.array(img)
        img_name, img_ext = (img_name.split('.')[0], '.' + img_name.split('.')[1])
    elif img_path:
        in_img = imageio.imread(img_path)
        img_name = os.path.basename(img_path)
        img_name, img_ext = (img_name.split('.')[0], '.' + img_name.split('.')[1])
    in_img = rescale_img(crop_img(in_img))
    in_img = in_img[:, :, :3]
    p_img = transformation(in_img).unsqueeze(0)

    # run model and get output image, upsampled warping field (flow), and 32 x 32 warping field (small_flow)
    output, flow, small_flow = model(p_img.float().to(device))
    output = output.detach().squeeze(0).permute(1,2,0).cpu()

    # save cropped input image
    imageio.imwrite(os.path.join(save_dir, img_name + '_orig' + img_ext), np.array(in_img).astype('uint8'))
    if visualize:
        plt.figure()
        plt.axis('off')
        plt.imshow(in_img)
        plt.show()
        plt.close()

    # save output warped image
    manual_warp = flow_warp(torch.from_numpy(in_img).permute(2,0,1).unsqueeze(0).float(), flow.cpu().detach() * scale)
    manual_warp = manual_warp.detach().squeeze(0).permute(1,2,0)
    imageio.imwrite(os.path.join(save_dir, img_name + '_out' + img_ext), np.array(manual_warp).astype('uint8'))
    if visualize:
        plt.figure()
        plt.axis('off')
        plt.imshow(manual_warp / 255., alpha=1)
        plt.show()
        plt.close()

    # save quiver plot of warping field
    plot_quiver(flow.cpu().detach() * scale, os.path.join(save_dir, img_name + '_quiver' + img_ext), save=True)
    plt.close()

    # save quiver plot overlaid on top of warped cartoon
    plot_quiver(flow.cpu().detach() * scale, '', save=False, overlaid=True)
    manual_warp = flow_warp(torch.from_numpy(in_img).permute(2,0,1).unsqueeze(0).float(), flow.cpu().detach() * scale)
    manual_warp = manual_warp.detach().squeeze(0).permute(1,2,0)
    plt.imshow(manual_warp / 255., alpha=0.7)
    plt.savefig(os.path.join(save_dir, img_name + '_overlaid' + img_ext), bbox_inches='tight', pad_inches=0)
    plt.close()

    # save x flow
    plot_quiver(flow.cpu().detach() * scale, '', save=False, overlaid=True)
    plt.imshow(flow.detach().cpu().squeeze(0).permute(1,2,0)[:, :, 0] * scale)
    plt.colorbar()
    plt.savefig(os.path.join(save_dir, img_name + '_xflow' + img_ext), bbox_inches='tight', pad_inches=0)
    plt.clf()
    plt.close()

    # save y flow
    plot_quiver(flow.cpu().detach() * scale, '', save=False, overlaid=True)
    plt.imshow(flow.detach().cpu().squeeze(0).permute(1,2,0)[:, :, 1] * scale)
    plt.colorbar()
    plt.savefig(os.path.join(save_dir, img_name + '_yflow' + img_ext), bbox_inches='tight', pad_inches=0)
    plt.close()
