# CSC320 Winter 2018
# Assignment 3
# (c) Olga (Ge Ya) Xu, Kyros Kutulakos
#
# DISTRIBUTION OF THIS CODE ANY FORM (ELECTRONIC OR OTHERWISE,
# AS-IS, MODIFIED OR IN PART), WITHOUT PRIOR WRITTEN AUTHORIZATION
# BY KYROS KUTULAKOS IS STRICTLY PROHIBITED. VIOLATION OF THIS
# POLICY WILL BE CONSIDERED AN ACT OF ACADEMIC DISHONESTY

#
# DO NOT MODIFY THIS FILE ANYWHERE EXCEPT WHERE INDICATED
#

# import basic packages
import numpy as np
import math

# basic numpy configuration

# set random seed
np.random.seed(seed=131)
# ignore division by zero warning
np.seterr(divide='ignore', invalid='ignore')


# This function implements the basic loop of the PatchMatch
# algorithm, as explained in Section 3.2 of the paper.
# The function takes an NNF f as input, performs propagation and random search,
# and returns an updated NNF.
#
# The function takes several input arguments:
#     - source_patches:      The matrix holding the patches of the source image,
#                            as computed by the make_patch_matrix() function. For an
#                            NxM source image and patches of width P, the matrix has
#                            dimensions NxMxCx(P^2) where C is the number of color channels
#                            and P^2 is the total number of pixels in the patch. The
#                            make_patch_matrix() is defined below and is called by the
#                            initialize_algorithm() method of the PatchMatch class. For
#                            your purposes, you may assume that source_patches[i,j,c,:]
#                            gives you the list of intensities for color channel c of
#                            all pixels in the patch centered at pixel [i,j]. Note that patches
#                            that go beyond the image border will contain NaN values for
#                            all patch pixels that fall outside the source image.
#     - target_patches:      The matrix holding the patches of the target image.
#     - f:                   The current nearest-neighbour field
#     - alpha, w:            Algorithm parameters, as explained in Section 3 and Eq.(1)
#     - propagation_enabled: If true, propagation should be performed.
#                            Use this flag for debugging purposes, to see how your
#                            algorithm performs with (or without) this step
#     - random_enabled:      If true, random search should be performed.
#                            Use this flag for debugging purposes, to see how your
#                            algorithm performs with (or without) this step.
#     - odd_iteration:       True if and only if this is an odd-numbered iteration.
#                            As explained in Section 3.2 of the paper, the algorithm
#                            behaves differently in odd and even iterations and this
#                            parameter controls this behavior.
#     - best_D:              And NxM matrix whose element [i,j] is the similarity score between
#                            patch [i,j] in the source and its best-matching patch in the
#                            target. Use this matrix to check if you have found a better
#                            match to [i,j] in the current PatchMatch iteration
#     - global_vars:         (optional) if you want your function to use any global variables,
#                            you can pass them to/from your function using this argument

# Return arguments:
#     - new_f:               The updated NNF
#     - best_D:              The updated similarity scores for the best-matching patches in the
#                            target
#     - global_vars:         (optional) if you want your function to use any global variables,
#                            return them in this argument and they will be stored in the
#                            PatchMatch data structure


def propagation_and_random_search(source_patches, target_patches,
                                  f, alpha, w,
                                  propagation_enabled, random_enabled,
                                  odd_iteration, best_D=None,
                                  global_vars=None
                                  ):
    new_f = f.copy()

    #############################################
    ###  PLACE YOUR CODE BETWEEN THESE LINES  ###
    #############################################
    src_shape = source_patches.shape

    if best_D is None:
        # initialize best_D
        best_D = np.ones((src_shape[0], src_shape[1], 1)) * np.inf
    # propagation
    if propagation_enabled == False:
        if odd_iteration:
            for i in range(src_shape[0]):
                for j in range(src_shape[1]):
                    # middle patch
                    middle_patch = source_patches[i, j]
                    d_1 = np.inf
                    d_2 = np.inf
                    D_list = []
                    middle = target_patches[i + new_f[i, j, 0], j + new_f[i, j, 1]]
                    delta_0 = np.ndarray.flatten(middle_patch - middle)
                    w0 = np.count_nonzero(~np.isnan(delta_0))
                    delta_0[delta_0 == np.nan] = 0
                    d_0 = np.true_divide(np.sum(np.dot(delta_0, delta_0)), w0)
                    D_list.append(d_0)

                    # left patch
                    if (0 <= i - 1 < src_shape[0]):
                        if (0 <= i + new_f[i - 1, j, 0] < src_shape[0] and 0 <= j + new_f[i - 1, j, 1] <
                                src_shape[1]):
                            left = target_patches[i + new_f[i - 1, j, 0], j + new_f[i - 1, j, 1]]
                            delta_left = np.ndarray.flatten(middle_patch - left)
                            w1 = np.count_nonzero(~np.isnan(delta_left))
                            delta_left[delta_left == np.nan] = 0
                            d_1 = np.true_divide(np.sum(np.dot(delta_left, delta_left)), w1)
                            D_list.append(d_1)
                    # up patch
                    if (0 <= j - 1 < src_shape[1]):
                        if (0 <= i + new_f[i, j - 1, 0] < src_shape[0] and 0 <= j + new_f[i, j - 1, 1] <
                                src_shape[1]):
                            up = target_patches[i + new_f[i, j - 1, 0], j + new_f[i, j - 1, 1]]
                            delta_up = np.ndarray.flatten(middle_patch - up)
                            w2 = np.count_nonzero(~np.isnan(delta_up))
                            delta_up[delta_up == np.nan] = 0
                            d_2 = np.true_divide(np.sum(np.dot(delta_up, delta_up)), w2)
                            D_list.append(d_2)
                    # select the best d_value and update the new_f
                    temp_best = np.min(D_list)
                    if (temp_best == d_0):
                        new_f[i, j] = np.array([new_f[i, j, 0], new_f[i, j, 1]])
                        best_D[i, j] = temp_best

                    if (temp_best == d_1):
                        new_f[i, j] = np.array([new_f[i - 1, j, 0], new_f[i - 1, j, 1]])
                        best_D[i, j] = temp_best

                    if (temp_best == d_2):
                        new_f[i, j] = np.array([new_f[i, j - 1, 0], new_f[i, j - 1, 1]])
                        best_D[i, j] = temp_best

                    if random_enabled == False:
                        i_exp = 0

                        temp = (alpha ** i_exp) * w
                        src_temp = (source_patches[i, j])
                        while temp >= 1:
                            R = np.array([np.random.uniform(-1, 1, 1), np.random.uniform(-1, 1, 1)])
                            u_f = np.array([i + new_f[i, j, 0] + np.multiply(temp, R)[0].astype(int),
                                            j + new_f[i, j, 1] + np.multiply(temp, R)[1].astype(int)])
                            if (0 < u_f[0] < src_shape[0] and 0 < u_f[1] < src_shape[1]):
                                dest_temp = (target_patches[u_f[0], u_f[1]])
                                delta = np.ndarray.flatten(src_temp - dest_temp)
                                w4 = np.count_nonzero(~np.isnan(delta))
                                delta[delta == np.nan] = 0
                                abs_dist = np.true_divide(np.sum(np.dot(delta, delta)), w4)

                                if (abs_dist < best_D[i, j]):
                                    new_f[i, j, 0] = new_f[i, j, 0] + np.multiply(temp, R)[0].astype(int)
                                    new_f[i, j, 1] = new_f[i, j, 1] + np.multiply(temp, R)[1].astype(int)
                                    best_D[i, j] = abs_dist

                            temp = temp * (alpha ** i_exp)
                            i_exp += 1
        # even iteration
        else:
            for i in range(src_shape[0]-1,-1,-1):
                for j in range(src_shape[1]-1,-1,-1):
                    # middle patch
                    middle_patch = source_patches[i, j]
                    d_1 = np.inf
                    d_2 = np.inf
                    D_list = []
                    middle = target_patches[i + new_f[i, j, 0], j + new_f[i, j, 1]]
                    delta_0 = np.ndarray.flatten(middle_patch - middle)
                    w0 = np.count_nonzero(~np.isnan(delta_0))
                    delta_0[delta_0 == np.nan] = 0
                    d_0 = np.true_divide(np.sum(np.dot(delta_0, delta_0)), w0)
                    D_list.append(d_0)
                    # right patch
                    if (0 <= i + 1 < src_shape[0]):
                        if (0 <= i + new_f[i + 1, j, 0] < src_shape[0] and 0 <= j + new_f[i + 1, j, 1] < src_shape[
                                1]):
                            right = target_patches[i + new_f[i + 1, j, 0], j + new_f[i + 1, j, 1]]
                            delta_right = np.ndarray.flatten(middle_patch - right)
                            w1 = np.count_nonzero(~np.isnan(delta_right))
                            delta_right[delta_right == np.nan] = 0
                            d_1 = np.true_divide(np.sum(np.dot(delta_right, delta_right)), w1)
                            D_list.append(d_1)
                    # down patch
                    if (0 <= j + 1 < src_shape[1]):
                        if (0 <= i + new_f[i, j + 1, 0] < src_shape[0] and 0 <= j + new_f[i, j + 1, 1] < src_shape[
                            1]):
                            down = target_patches[i + new_f[i, j + 1, 0], j + new_f[i, j + 1, 1]]
                            delta_down = np.ndarray.flatten(middle_patch - down)
                            w2 = np.count_nonzero(~np.isnan(delta_down))
                            delta_down[delta_down == np.nan] = 0
                            d_2 = np.true_divide(np.sum(np.dot(delta_down, delta_down)), w2)
                            D_list.append(d_2)
                    # select the best d_value and update the new_f
                    temp_best = np.min(D_list)
                    if (temp_best == d_0):
                        new_f[i, j] = np.array([new_f[i, j, 0], new_f[i, j, 1]])
                        best_D[i, j] = temp_best

                    if (temp_best == d_1):
                        new_f[i, j] = np.array([new_f[i + 1, j, 0], new_f[i + 1, j, 1]])
                        best_D[i, j] = temp_best

                    if (temp_best == d_2):
                        new_f[i, j] = np.array([new_f[i, j + 1, 0], new_f[i, j + 1, 1]])
                        best_D[i, j] = temp_best
                # random
                if random_enabled == False:

                    i_exp = 0

                    temp = (alpha ** i_exp) * w
                    src_temp = (source_patches[i, j])
                    while temp >= 1:
                        R = np.array([np.random.uniform(-1, 1, 1), np.random.uniform(-1, 1, 1)])
                        u_f = np.array([i + new_f[i, j, 0] + np.multiply(temp, R)[0].astype(int),
                                        j + new_f[i, j, 1] + np.multiply(temp, R)[1].astype(int)])

                        if (0 < u_f[0] < src_shape[0] and 0 < u_f[1] < src_shape[1]):
                            dest_temp = (target_patches[u_f[0], u_f[1]])
                            delta = np.ndarray.flatten(src_temp - dest_temp)
                            w4 = np.count_nonzero(~np.isnan(delta))
                            delta[delta == np.nan] = 0
                            abs_dist = np.true_divide(np.sum(np.dot(delta, delta)), w4)

                            if (abs_dist < best_D[i, j]):
                                new_f[i, j, 0] = new_f[i, j, 0] + np.multiply(temp, R)[0].astype(int)
                                new_f[i, j, 1] = new_f[i, j, 1] + np.multiply(temp, R)[1].astype(int)
                                best_D[i, j] = abs_dist

                        temp = temp * (alpha ** i_exp)
                        i_exp += 1


    #############################################

    return new_f, best_D, global_vars


# This function uses a computed NNF to reconstruct the source image
# using pixels from the target image. The function takes two input
# arguments
#     - target: the target image that was used as input to PatchMatch
#     - f:      the nearest-neighbor field the algorithm computed
# and should return a reconstruction of the source image:
#     - rec_source: an openCV image that has the same shape as the source image
#
# To reconstruct the source, the function copies to pixel (x,y) of the source
# the color of pixel (x,y)+f(x,y) of the target.
#
# The goal of this routine is to demonstrate the quality of the computed NNF f.
# Specifically, if patch (x,y)+f(x,y) in the target image is indeed very similar
# to patch (x,y) in the source, then copying the color of target pixel (x,y)+f(x,y)
# to the source pixel (x,y) should not change the source image appreciably.
# If the NNF is not very high quality, however, the reconstruction of source image
# will not be very good.
#
# You should use matrix/vector operations to avoid looping over pixels,
# as this would be very inefficient

def reconstruct_source_from_target(target, f):
    rec_source = None

    #############################################
    ###  PLACE YOUR CODE BETWEEN THESE LINES  ###
    #############################################
    N = target.shape[0]
    M = target.shape[1]
    origin= make_coordinates_matrix(target.shape)
    target_map = origin + f
    mapping = target.copy() * 0
    for i in range(N):
        for j in range(M):
            mapping[i, j] = target[target_map[i, j, 0], target_map[i, j, 1]]
    rec_source = mapping.reshape(-1, f.shape[1], 3)

    #############################################
    return rec_source


# This function takes an NxM image with C color channels and a patch size P
# and returns a matrix of size NxMxCxP^2 that contains, for each pixel [i,j] in
# in the image, the pixels in the patch centered at [i,j].
#
# You should study this function very carefully to understand precisely
# how pixel data are organized, and how patches that extend beyond
# the image border are handled.


def make_patch_matrix(im, patch_size):
    phalf = patch_size // 2
    # create an image that is padded with patch_size/2 pixels on all sides
    # whose values are NaN outside the original image
    padded_shape = im.shape[0] + patch_size - 1, im.shape[1] + patch_size - 1, im.shape[2]
    padded_im = np.zeros(padded_shape) * np.NaN
    padded_im[phalf:(im.shape[0] + phalf), phalf:(im.shape[1] + phalf), :] = im

    # Now create the matrix that will hold the vectorized patch of each pixel. If the
    # original image had NxM pixels, this matrix will have NxMx(patch_size*patch_size)
    # pixels
    patch_matrix_shape = im.shape[0], im.shape[1], im.shape[2], patch_size ** 2
    patch_matrix = np.zeros(patch_matrix_shape) * np.NaN
    for i in range(patch_size):
        for j in range(patch_size):
            patch_matrix[:, :, :, i * patch_size + j] = padded_im[i:(i + im.shape[0]), j:(j + im.shape[1]), :]

    return patch_matrix


# Generate a matrix g of size (im_shape[0] x im_shape[1] x 2)
# such that g(y,x) = [y,x]
#
# Step is an optional argument used to create a matrix that is step times
# smaller than the full image in each dimension
#
# Pay attention to this function as it shows how to perform these types
# of operations in a vectorized manner, without resorting to loops


def make_coordinates_matrix(im_shape, step=1):
    """
    Return a matrix of size (im_shape[0] x im_shape[1] x 2) such that g(x,y)=[y,x]
    """
    range_x = np.arange(0, im_shape[1], step)
    range_y = np.arange(0, im_shape[0], step)
    axis_x = np.repeat(range_x[np.newaxis, ...], len(range_y), axis=0)
    axis_y = np.repeat(range_y[..., np.newaxis], len(range_x), axis=1)

    return np.dstack((axis_y, axis_x))
