from scipy.misc import imread
from scipy import signal
from skimage.color import rgb2gray
import numpy as np
import matplotlib.pyplot as plt
import os

"""
    a function which reads an image file and converts it into a given representation.
    
    parameters:
    filename - the filename of an image on disk (could be grayscale or RGB).
    representation - representation code, either 1 or 2 defining whether the
    output should be a grayscale
    image (1) or an RGB image (2). If the input image is grayscale, we won’t
    call it with representation = 2.
"""
def read_image(filename, representation):
    img = imread(filename)
    if representation == 1:
        img = rgb2gray(img)
        img = img.astype(np.float64)

    elif representation == 2:
        img = img.astype(np.float64)

    else:
        print("representation need to be '1' or '2'")
        return None

    return img / 255


""" ##########################################  3 - Image Pyramids  ######################################## """


def reduce(img, filter):

    conv_img = signal.convolve2d(img, filter, "same")
    conv_img = signal.convolve2d(conv_img, filter.T, "same")

    reduced_img = conv_img[::2, ::2]
    return reduced_img


"""
    3.1 - Gaussian pyramid 
    construct a Gaussian pyramid of a given image.
    
    parameters:
    im – a grayscale image with double values in [0, 1].
    max_levels – the maximal number of levels in the resulting pyramid.
    filter_size – the size of the Gaussian filter (an odd scalar that
    represents a squared filter) to be used in constructing the pyramid filter
    (e.g for filter_size = 3 you should get [0.25, 0.5, 0.25]).
"""
def build_gaussian_pyramid(im, max_levels, filter_size):
    pyr = []
    pyr.append(im)  # append G_0

    # built the filter using a consequent 1D convolutions of [1 1]
    base_filter = np.array([1, 1]).reshape(1, 2)
    filter_vec = base_filter

    for i in range(2, filter_size):
        filter_vec = signal.convolve2d(filter_vec, base_filter)

    # The filter_vec should be normalized
    filter_vec = filter_vec / np.sum(filter_vec)

    for i in range(1, max_levels):
        g_i = reduce(pyr[i-1], filter_vec)
        pyr.append(g_i)

    return pyr, filter_vec


# print('3.1 gaussian')
# im_orig = read_image('presubmit_externals/monkey.jpg', 1)
# gpyr, filter_vec = build_gaussian_pyramid(im_orig, 3, 3)


def expand(img, filter):
    rows = img.shape[0] * 2
    cols = img.shape[1] * 2
    expanded_img = np.zeros((rows, cols), np.float64)
    expanded_img[::2, ::2] = img

    filter = filter*2
    conv_img = signal.convolve2d(expanded_img, filter, "same")
    conv_img = signal.convolve2d(conv_img, filter.T, "same")

    return conv_img


"""
    3.1 - Laplacian pyramid
    construct a Laplacian pyramid of a given image
"""
def build_laplacian_pyramid(im, max_levels, filter_size):
    pyr = []
    gaussian_pyr, filter_vec = build_gaussian_pyramid(im, max_levels, filter_size)

    for i in range(max_levels - 1):
        l_i = gaussian_pyr[i] - expand(gaussian_pyr[i + 1], filter_vec)
        pyr.append(l_i)

    pyr.append(gaussian_pyr[max_levels - 1])
    return pyr, filter_vec


# print('3.1 laplacian')
# im_orig = read_image('presubmit_externals/monkey.jpg', 1)
# lpyr, filter_vec = build_laplacian_pyramid(im_orig, 3, 3)


"""
    3.2 - Laplacian pyramid reconstruction
    This function implement the reconstruction of an image from its Laplacian Pyramid
    
    parameter:
    lpyr and filter_vec are the Laplacian pyramid and the filter that are generated by
    the second function in 3.1.
    coeff is a python list. The list length is the same as the number of levels in the
    pyramid lpyr. Before reconstructing the image img you should multiply each level i
    of the laplacian pyramid by its corresponding coefficient coeff[i].
    only when this list is all ones we get the original image (up to a negligible
    floating error, e.g. maximal absolute difference around 10−12). When some values
    are different than 1 we will get filtering effects.
"""
def laplacian_to_image(lpyr, filter_vec, coeff):

    # multiply each level i of the laplacian pyramid by its corresponding
    # coefficient coeff[i]
    n = len(lpyr)
    for i in range(n):
        lpyr[i] = np.multiply(lpyr[i], coeff[i])

    result = lpyr[n-1]
    for j in range(2, n+1):
        a = expand(result, filter_vec)
        result = a + lpyr[-j]

    return result


# print('3.2')
# im_orig = read_image('presubmit_externals/monkey.jpg', 1)
# lpyr, filter_vec = build_laplacian_pyramid(im_orig, 3, 3)
# im_r = laplacian_to_image(lpyr, filter_vec, [1, 1, 1])
# plt.imshow(im_r, cmap="gray")
# plt.show()


"""
    3.3 - Pyramid display
    
    parameters:
    pyr - a Gaussian or Laplacian pyramid
    level - the number of levels to present in the result ≤ max_levels
    
    return: res - s a single black image in which the pyramid levels of the given
    pyramid pyr are stacked horizontally 
"""
def  render_pyramid(pyr, levels):
    shape_im_0 = pyr[0].shape
    length = min(levels, len(pyr))
    zeros = np.zeros(length)

    for i in range(length):
        zeros[i] = pyr[i].shape[0]

    cum_shapes = np.cumsum(zeros)
    cum_shapes = cum_shapes.astype(np.int64)

    # add zero at the begining
    shapes = np.insert(cum_shapes, 0, 0)

    width = shape_im_0[1]
    heigth = shapes[length]

    res = np.zeros((width, heigth))
    for j in range(length):
        min_val = pyr[j].min()
        max_val = pyr[j].max()

        # insert img to res, in the left up corner
        img = pyr[j] - min_val
        norm_img = img / (max_val - min_val)

        res[0:pyr[j].shape[1], shapes[j]:shapes[j+1]] = norm_img

    return res


def display_pyramid(pyr, levels):
    res = render_pyramid(pyr, levels)
    plt.imshow(res, cmap="gray")
    plt.show()


print('3.3')
im_orig = read_image('presubmit_externals/monkey.jpg', 1)
gpyr, filter_vec = build_gaussian_pyramid(im_orig, 3, 3)
display_pyramid(gpyr, 2)


""" ##########################################  4 - Pyramid Blending  ######################################## """


"""
    4 - pyramid blending
    
    parameter:
    im1, im2 – are two input grayscale images to be blended.
    mask – is a boolean mask containing True and False representing which parts
    of im1 and im2 should appear in the resulting im_blend.
    max_levels – is the max_levels parameter you should use when generating the
    Gaussian and Laplacian pyramids.
    filter_size_im – is the size of the Gaussian filter (an odd scalar that represents
    a squared filter) which defining the filter used in the construction of the
    Laplacian pyramids of im1 and im2
    filter_size_mask – is the size of the Gaussian filter(an odd scalar that represents
    a squared filter) which defining the filter used in the construction of the Gaussian
    pyramid of mask.
"""
def pyramid_blending(im1, im2, mask, max_levels, filter_size_im, filter_size_mask):
    l1, filter_vec = build_laplacian_pyramid(im1, max_levels, filter_size_im)
    l2, filter_vec = build_laplacian_pyramid(im2, max_levels, filter_size_im)

    mask = mask.astype(np.float64)
    g_m, g_filter = build_gaussian_pyramid(mask, max_levels, filter_size_mask)

    l_out = []
    for i in range(max_levels):
        l_out_i = (g_m[i] * l1[i]) + ((1 - g_m[i]) * l2[i])
        l_out.append(l_out_i)

    coeff = [1] * max_levels
    im_blend = laplacian_to_image(l_out, filter_vec, coeff)
    im_blend = np.clip(im_blend, a_min=0, a_max=1)

    return im_blend


# print("4")
# im1 = read_image('externals/5_1024.jpg', 1)
# im2 = read_image('externals/2.jpg', 1)
# mask = read_image('externals/mask1.jpg', 1).astype(np.bool)
# mask[mask < 0.5] = 0
# mask[mask >= 0.5] = 1
# im_blend = pyramid_blending(im1, im2, mask, 3, 3, 5)
# plt.imshow(im_blend)
# plt.show()


"""
    4.1 - blending examples
"""
def relpath(filename):
    return os.path.join(os.path.dirname(__file__), filename)


def blending_example1():
    im1 = read_image(relpath("5_1024.jpg"), 2)
    im2 = read_image(relpath("2.jpg"), 2)
    mask = read_image(relpath("mask1.jpg"), 1).astype(np.bool)

    shape = im1.shape
    im_blend = np.zeros(shape).astype(np.float64)

    im_blend[:, :, 0] = pyramid_blending(im1[:, :, 0], im2[:, :, 0], mask, 3, 3, 5)
    im_blend[:, :, 1] = pyramid_blending(im1[:, :, 1], im2[:, :, 1], mask, 3, 3, 5)
    im_blend[:, :, 2] = pyramid_blending(im1[:, :, 2], im2[:, :, 2], mask, 3, 3, 5)

    plt.figure()
    fig, axs = plt.subplots(2, 2)
    axs[0, 0].imshow(im1, cmap=plt.get_cmap("gray"), interpolation='nearest')
    axs[0, 1].imshow(im2, cmap=plt.get_cmap("gray"), interpolation='nearest')
    axs[1, 0].imshow(mask, cmap=plt.get_cmap("gray"), interpolation='nearest')
    axs[1, 1].imshow(im_blend, cmap=plt.get_cmap("gray"), interpolation='nearest')

    plt.show()

    return im1, im2, mask, im_blend


# im1, im2, mask, im_blend = blending_example1()


def blending_example2():
    im1 = read_image(relpath("3_512.jpg"), 2)
    im2 = read_image(relpath("7.jpg"), 2)
    mask = read_image(relpath("mask2.jpg"), 1).astype(np.bool)

    shape = im1.shape
    im_blend = np.zeros(shape).astype(np.float64)

    im_blend[:, :, 0] = pyramid_blending(im1[:, :, 0], im2[:, :, 0], mask, 3, 3, 5)
    im_blend[:, :, 1] = pyramid_blending(im1[:, :, 1], im2[:, :, 1], mask, 3, 3, 5)
    im_blend[:, :, 2] = pyramid_blending(im1[:, :, 2], im2[:, :, 2], mask, 3, 3, 5)

    plt.figure()
    fig, axs = plt.subplots(2, 2)
    axs[0, 0].imshow(im1, cmap=plt.get_cmap("gray"), interpolation='nearest')
    axs[0, 1].imshow(im2, cmap=plt.get_cmap("gray"), interpolation='nearest')
    axs[1, 0].imshow(mask, cmap=plt.get_cmap("gray"), interpolation='nearest')
    axs[1, 1].imshow(im_blend, cmap=plt.get_cmap("gray"), interpolation='nearest')

    plt.show()

    return im1, im2, mask, im_blend


# im1, im2, mask, im_blend = blending_example2()
