from __future__ import absolute_import, division, print_function, unicode_literals

import scipy.ndimage
from scipy.sparse.linalg import spsolve
from scipy import sparse
import scipy.io as sio
import numpy as np
from PIL import Image
import copy
import cv2
import os
import argparse


def sub2ind(pi, pj, imgH, imgW):
    return pj + pi * imgW


def Poisson_blend_img(imgTrg, imgSrc_gx, imgSrc_gy, holeMask, gradientMask=None, edge=None):

    imgH, imgW, nCh = imgTrg.shape

    if not isinstance(gradientMask, np.ndarray):
        gradientMask = np.zeros((imgH, imgW), dtype=np.float32)

    if not isinstance(edge, np.ndarray):
        edge = np.zeros((imgH, imgW), dtype=np.float32)

    # Initialize the reconstructed image
    imgRecon = np.zeros((imgH, imgW, nCh), dtype=np.float32)

    # prepare discrete Poisson equation
    A, b, UnfilledMask = solvePoisson(holeMask, imgSrc_gx, imgSrc_gy, imgTrg,
                                                  gradientMask, edge)

    # Independently process each channel
    for ch in range(nCh):

        # solve Poisson equation
        x = scipy.sparse.linalg.lsqr(A, b[:, ch])[0]

        imgRecon[:, :, ch] = x.reshape(imgH, imgW)

    # Combined with the known region in the target
    holeMaskC = np.tile(np.expand_dims(holeMask, axis=2), (1, 1, nCh))
    imgBlend = holeMaskC * imgRecon + (1 - holeMaskC) * imgTrg


    # while((UnfilledMask * edge).sum() != 0):
    #     # Fill in edge pixel
    #     pi = np.expand_dims(np.where((UnfilledMask * edge) == 1)[0], axis=1) # y, i
    #     pj = np.expand_dims(np.where((UnfilledMask * edge) == 1)[1], axis=1) # x, j
    #
    #     for k in range(len(pi)):
    #         if pi[k, 0] - 1 >= 0:
    #             if (UnfilledMask * edge)[pi[k, 0] - 1, pj[k, 0]] == 0:
    #                 imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0] - 1, pj[k, 0], :]
    #                 UnfilledMask[pi[k, 0], pj[k, 0]] = 0
    #                 continue
    #         if pi[k, 0] + 1 <= imgH - 1:
    #             if (UnfilledMask * edge)[pi[k, 0] + 1, pj[k, 0]] == 0:
    #                 imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0] + 1, pj[k, 0], :]
    #                 UnfilledMask[pi[k, 0], pj[k, 0]] = 0
    #                 continue
    #         if pj[k, 0] - 1 >= 0:
    #             if (UnfilledMask * edge)[pi[k, 0], pj[k, 0] - 1] == 0:
    #                 imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0], pj[k, 0] - 1, :]
    #                 UnfilledMask[pi[k, 0], pj[k, 0]] = 0
    #                 continue
    #         if pj[k, 0] + 1 <= imgW - 1:
    #             if (UnfilledMask * edge)[pi[k, 0], pj[k, 0] + 1] == 0:
    #                 imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0], pj[k, 0] + 1, :]
    #                 UnfilledMask[pi[k, 0], pj[k, 0]] = 0

    return imgBlend, UnfilledMask

def solvePoisson(holeMask, imgSrc_gx, imgSrc_gy, imgTrg,
                           gradientMask, edge):

    # UnfilledMask indicates the region that is not completed
    UnfilledMask_topleft = copy.deepcopy(holeMask)
    UnfilledMask_bottomright = copy.deepcopy(holeMask)

    # Prepare the linear system of equations for Poisson blending
    imgH, imgW = holeMask.shape
    N = imgH * imgW

    # Number of unknown variables
    numUnknownPix = holeMask.sum()

    # 4-neighbors: dx and dy
    dx = [1, 0, -1,  0]
    dy = [0, 1,  0, -1]

    #      3
    #      |
    # 2 -- * -- 0
    #      |
    #      1
    #

    # Initialize (I, J, S), for sparse matrix A where A(I(k), J(k)) = S(k)
    I = np.empty((0, 1), dtype=np.float32)
    J = np.empty((0, 1), dtype=np.float32)
    S = np.empty((0, 1), dtype=np.float32)

    # Initialize b
    b = np.empty((0, 3), dtype=np.float32)

    # Precompute unkonwn pixel position
    pi = np.expand_dims(np.where(holeMask == 1)[0], axis=1) # y, i
    pj = np.expand_dims(np.where(holeMask == 1)[1], axis=1) # x, j
    pind = sub2ind(pi, pj, imgH, imgW)

    # |--------------------|
    # |        y (i)       |
    # |   x (j)  *         |
    # |                    |
    # |--------------------|
    # p[y, x]

    qi = np.concatenate((pi + dy[0],
                         pi + dy[1],
                         pi + dy[2],
                         pi + dy[3]), axis=1)

    qj = np.concatenate((pj + dx[0],
                         pj + dx[1],
                         pj + dx[2],
                         pj + dx[3]), axis=1)

    # Handling cases at image borders
    validN = (qi >= 0) & (qi <= imgH - 1) & (qj >= 0) & (qj <= imgW - 1)
    qind = np.zeros((validN.shape), dtype=np.float32)
    qind[validN] = sub2ind(qi[validN], qj[validN], imgH, imgW)

    e_start = 0  # equation counter start
    e_stop  = 0  # equation stop

    # 4 neighbors
    I, J, S, b, e_start, e_stop = constructEquation(0, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop)
    I, J, S, b, e_start, e_stop = constructEquation(1, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop)
    I, J, S, b, e_start, e_stop = constructEquation(2, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop)
    I, J, S, b, e_start, e_stop = constructEquation(3, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop)

    nEqn = len(b)
    # Construct the sparse matrix A
    A = sparse.csr_matrix((S[:, 0], (I[:, 0], J[:, 0])), shape=(nEqn, N))

    # Check connected pixels
    for ind in range(0, len(pi), 1):
        ii = pi[ind, 0]
        jj = pj[ind, 0]

        # check up (3)
        if ii - 1 >= 0:
            if UnfilledMask_topleft[ii - 1, jj] == 0 and gradientMask[ii - 1, jj] == 0:
                UnfilledMask_topleft[ii, jj] = 0
        # check left (2)
        if jj - 1 >= 0:
            if UnfilledMask_topleft[ii, jj - 1] == 0 and gradientMask[ii, jj - 1] == 0:
                UnfilledMask_topleft[ii, jj] = 0


    for ind in range(len(pi) - 1, -1, -1):
        ii = pi[ind, 0]
        jj = pj[ind, 0]

        # check bottom (1)
        if ii + 1 <= imgH - 1:
            if UnfilledMask_bottomright[ii + 1, jj] == 0 and gradientMask[ii, jj] == 0:
                UnfilledMask_bottomright[ii, jj] = 0
        # check right (0)
        if jj + 1 <= imgW - 1:
            if UnfilledMask_bottomright[ii, jj + 1] == 0 and gradientMask[ii, jj] == 0:
                UnfilledMask_bottomright[ii, jj] = 0

    UnfilledMask = UnfilledMask_topleft * UnfilledMask_bottomright

    return A, b, UnfilledMask


def constructEquation(n, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop):

    # Pixel that has valid neighbors
    validNeighbor = validN[:, n]

    # Change the out-of-boundary value to 0, in order to run edge[y,x]
    # in the next line. It won't affect anything as validNeighbor is saved already

    qi_tmp = copy.deepcopy(qi)
    qj_tmp = copy.deepcopy(qj)
    qi_tmp[np.invert(validNeighbor), n] = 0
    qj_tmp[np.invert(validNeighbor), n] = 0

    NotEdge = (edge[pi[:, 0], pj[:, 0]] == 0) * (edge[qi_tmp[:, n], qj_tmp[:, n]] == 0)

    # Have gradient
    if n == 0:
        HaveGrad = gradientMask[pi[:, 0], pj[:, 0]] == 0
    elif n == 2:
        HaveGrad = gradientMask[pi[:, 0], pj[:, 0] - 1] == 0
    elif n == 1:
        HaveGrad = gradientMask[pi[:, 0], pj[:, 0]] == 0
    elif n == 3:
        HaveGrad = gradientMask[pi[:, 0] - 1, pj[:, 0]] == 0

    # Boundary constraint
    Boundary = holeMask[qi_tmp[:, n], qj_tmp[:, n]] == 0

    valid = validNeighbor * NotEdge * HaveGrad * Boundary

    J_tmp = pind[valid, :]

    # num of equations: len(J_tmp)
    e_stop = e_start + len(J_tmp)
    I_tmp = np.arange(e_start, e_stop, dtype=np.float32).reshape(-1, 1)
    e_start = e_stop

    S_tmp = np.ones(J_tmp.shape, dtype=np.float32)

    if n == 0:
        b_tmp = - imgSrc_gx[pi[valid, 0], pj[valid, 0], :] + imgTrg[qi[valid, n], qj[valid, n], :]
    elif n == 2:
        b_tmp = imgSrc_gx[pi[valid, 0], pj[valid, 0] - 1, :] + imgTrg[qi[valid, n], qj[valid, n], :]
    elif n == 1:
        b_tmp = - imgSrc_gy[pi[valid, 0], pj[valid, 0], :] + imgTrg[qi[valid, n], qj[valid, n], :]
    elif n == 3:
        b_tmp = imgSrc_gy[pi[valid, 0] - 1, pj[valid, 0], :] + imgTrg[qi[valid, n], qj[valid, n], :]

    I = np.concatenate((I, I_tmp))
    J = np.concatenate((J, J_tmp))
    S = np.concatenate((S, S_tmp))
    b = np.concatenate((b, b_tmp))

    # Non-boundary constraint
    NonBoundary = holeMask[qi_tmp[:, n], qj_tmp[:, n]] == 1
    valid = validNeighbor * NotEdge * HaveGrad * NonBoundary

    J_tmp = pind[valid, :]

    # num of equations: len(J_tmp)
    e_stop = e_start + len(J_tmp)
    I_tmp = np.arange(e_start, e_stop, dtype=np.float32).reshape(-1, 1)
    e_start = e_stop

    S_tmp = np.ones(J_tmp.shape, dtype=np.float32)

    if n == 0:
        b_tmp = - imgSrc_gx[pi[valid, 0], pj[valid, 0], :]
    elif n == 2:
        b_tmp = imgSrc_gx[pi[valid, 0], pj[valid, 0] - 1, :]
    elif n == 1:
        b_tmp = - imgSrc_gy[pi[valid, 0], pj[valid, 0], :]
    elif n == 3:
        b_tmp = imgSrc_gy[pi[valid, 0] - 1, pj[valid, 0], :]

    I = np.concatenate((I, I_tmp))
    J = np.concatenate((J, J_tmp))
    S = np.concatenate((S, S_tmp))
    b = np.concatenate((b, b_tmp))

    S_tmp = - np.ones(J_tmp.shape, dtype=np.float32)
    J_tmp = qind[valid, n, None]

    I = np.concatenate((I, I_tmp))
    J = np.concatenate((J, J_tmp))
    S = np.concatenate((S, S_tmp))

    return I, J, S, b, e_start, e_stop


def getUnfilledMask(holeMask, gradientMask):
    # UnfilledMask indicates the region that is not completed
    UnfilledMask_topleft = copy.deepcopy(holeMask)
    UnfilledMask_bottomright = copy.deepcopy(holeMask)

    # Get the shape information of the mask
    imgH, imgW = holeMask.shape

    # Precompute the unknown pixel position
    pi = np.expand_dims(np.where(holeMask == 1)[0], axis=1)
    pj = np.expand_dims(np.where(holeMask == 1)[1], axis=1)

    # Check connected pixels
    for ind in range(0, len(pi), 1):
        ii = pi[ind, 0]
        jj = pj[ind, 0]

        # check up (3)
        if ii - 1 >= 0:
            if UnfilledMask_topleft[ii - 1, jj] == 0 and gradientMask[ii - 1, jj] == 0:
                UnfilledMask_topleft[ii, jj] = 0
        # check left (2)
        if jj - 1 >= 0:
            if UnfilledMask_topleft[ii, jj - 1] == 0 and gradientMask[ii, jj - 1] == 0:
                UnfilledMask_topleft[ii, jj] = 0

    for ind in range(len(pi) - 1, -1, -1):
        ii = pi[ind, 0]
        jj = pj[ind, 0]

        # check bottom (1)
        if ii + 1 <= imgH - 1:
            if UnfilledMask_bottomright[ii + 1, jj] == 0 and gradientMask[ii, jj] == 0:
                UnfilledMask_bottomright[ii, jj] = 0
        # check right (0)
        if jj + 1 <= imgW - 1:
            if UnfilledMask_bottomright[ii, jj + 1] == 0 and gradientMask[ii, jj] == 0:
                UnfilledMask_bottomright[ii, jj] = 0

    UnfilledMask = UnfilledMask_topleft * UnfilledMask_bottomright

    return UnfilledMask