import scipy
import scipy.misc

def split_rgb(aos):
    return array(split(aos, 3, axis=2)).squeeze()

def join_rgb(soa):
    return concatenate(expand_dims(soa, 3), axis=2)

class PoissonSolver:
    pass

def step1(self, ref_file, target_file, ofs, target_scale, common_scale):
    ref           = scipy.misc.imresize(scipy.misc.imread(ref_file), common_scale, interp="bicubic")
    target_scaled = scipy.misc.imresize(scipy.misc.imread(target_file), common_scale * target_scale, interp="bicubic")

    rn            = ref.shape[:2]
    tn            = target_scaled.shape[:2]
    
    Xg, Yg        = ogrid[:rn[0], :rn[1]]
    mask          = (Xg >= ofs[0]) & (Xg < ofs[0]+tn[0]) & (Yg >= ofs[1]) & (Yg < ofs[1]+tn[1])

    rgb_view_mode                                    = [("r", uint8), ("g", uint8), ("b", uint8)]

    target                                           = zeros_like(ref).view(dtype=rgb_view_mode).squeeze().copy()
    target[ofs[0]:ofs[0]+tn[0], ofs[1]:ofs[1]+tn[1]] = target_scaled.view(dtype=rgb_view_mode).squeeze()
    target                                           = reshape(target.view(dtype=uint8), ref.shape)
    
    naive_clone                                      = ref.copy()
    naive_clone[mask == 1]                           = target[mask == 1]
    
    figsize(20, 15)
    for i, image in enumerate(["ref", "target", "mask", "naive_clone"]):
        matplotlib.pyplot.subplot(3, 4, 1 + i)
        matplotlib.pyplot.imshow(eval(image))
        matplotlib.pyplot.title(image)

    ref_soa = split_rgb(ref)
    target_soa = split_rgb(target)

    color_mode = 'RGB'
    for i in range(3):
        matplotlib.pyplot.subplot(3, 4, 5 + i)
        matplotlib.pyplot.imshow(eval("ref_soa[%d]" % i))
        matplotlib.pyplot.title("ref %s" % color_mode[i])
        matplotlib.pyplot.subplot(3, 4, 9 + i)
        matplotlib.pyplot.imshow(eval("target_soa[%d]" % i))
        matplotlib.pyplot.title("target %s" % color_mode[i])

    self.ref_soa = ref_soa
    self.target_soa = target_soa
    self.mask = mask
    self.naive_clone = naive_clone
    self.rn = rn
    self.tn = tn

s = PoissonSolver()
step1(s, "sky.png", "bird.png", [10, 50], 0.5, 0.25)

import skimage
import skimage.morphology

def step2(self):
    self.strict_interior = skimage.morphology.erosion(self.mask, ones((3,3), dtype=bool))
    self.border          = self.mask & ~self.strict_interior
    
    self.strict_interior_indices                          = self.strict_interior.nonzero()
    self.num_strict_interior_pixels                       = self.strict_interior_indices[0].shape[0]
    self.variable_index_map                               = -1 * ones_like(self.mask, dtype=int32)
    self.variable_index_map[self.strict_interior_indices] = arange(self.num_strict_interior_pixels)
    
    figsize(13, 4)
    for i, image in enumerate(["strict_interior", "border", "variable_index_map"]):
        matplotlib.pyplot.subplot(1, 3, 1 + i)
        matplotlib.pyplot.imshow(eval("self.%s" % image), interpolation="nearest")
        matplotlib.pyplot.title(image)

step2(s)

import scipy.sparse
import scipy.sparse.linalg

def gen_poisson_coeffs(self, ref, target):
    A = scipy.sparse.dok_matrix((self.num_strict_interior_pixels, ) * 2, dtype=float32)
    B = zeros(self.num_strict_interior_pixels, dtype=float32)
    
    for i in range(self.num_strict_interior_pixels):
        
        y = self.strict_interior_indices[0][i]
        x = self.strict_interior_indices[1][i]
        
        x_left, x_right = x-1, x+1
        y_up,   y_down  = y-1, y+1
    
        assert(i == self.variable_index_map[y,x])
        
        x_neighbors = []
        y_neighbors = []
        neighbors   = []
        
        if x_right < self.rn[1]:
            y_neighbors.append(y)
            x_neighbors.append(x_right)
            neighbors.append(self.variable_index_map[y, x_right])
            
        if y_up >= 0:
            y_neighbors.append(y_up)
            x_neighbors.append(x)
            neighbors.append(self.variable_index_map[y_up, x])
    
        if x_left >= 0:
            y_neighbors.append(y)
            x_neighbors.append(x_left)
            neighbors.append(self.variable_index_map[y, x_left])
    
        if y_down < self.rn[0]:
            y_neighbors.append(y_down)
            x_neighbors.append(x)
            neighbors.append(self.variable_index_map[y_down, x])
    
        y_neighbors               = array(y_neighbors)
        x_neighbors               = array(x_neighbors)
        neighbors                 = array(neighbors)
        neighbors_strict_interior = (neighbors != -1).nonzero()
        neighbors_border          = (neighbors == -1).nonzero()
        num_neighbors             = y_neighbors.shape[0]
    
        sum_vq       = (num_neighbors * target[y, x]) - sum(target[(y_neighbors, x_neighbors)])
        sum_border_f = sum(ref[(y_neighbors[neighbors_border],x_neighbors[neighbors_border])])
        
        A[i,i]                                        = num_neighbors
        A[i, neighbors[neighbors_strict_interior[0]]] = -1
        B[i]                                          = sum_vq + sum_border_f
    
    return A, B

coeffs = [gen_poisson_coeffs(s, s.ref_soa[i], s.target_soa[i]) for i in range(3)]

X_rgb = [scipy.sparse.linalg.spsolve(A.tocsr(), B) for A, B in coeffs]

seamless_clone_soa = s.ref_soa.copy()

for i in range(3):
    seamless_clone_soa[i][s.strict_interior_indices] = clip(X_rgb[i], 0, 255)

ref = join_rgb(s.ref_soa)
target = join_rgb(s.target_soa)
seamless_clone = join_rgb(seamless_clone_soa)
naive_clone = s.naive_clone

figsize(18, 4)
for i, image in enumerate(["ref", "target", "naive_clone", "seamless_clone"]):
    matplotlib.pyplot.subplot(1, 4, 1 + i)
    matplotlib.pyplot.imshow(locals()[image])
    matplotlib.pyplot.title(image)

arranged_coeffs = array(split(array(coeffs), 2, axis=1)).squeeze()
A = scipy.sparse.block_diag(arranged_coeffs[0]).tocsr()
B = concatenate(arranged_coeffs[1])

X_rgb = reshape(scipy.sparse.linalg.spsolve(A, B), (3, s.num_strict_interior_pixels))

seamless_clone_soa = s.ref_soa.copy()

for i in range(3):
    seamless_clone_soa[i][s.strict_interior_indices] = clip(X_rgb[i], 0, 255)

ref = join_rgb(s.ref_soa)
target = join_rgb(s.target_soa)
naive_clone = s.naive_clone
seamless_clone_full = join_rgb(seamless_clone_soa)

figsize(18, 4)
for i, image in enumerate(["ref", "target", "naive_clone", "seamless_clone_full"]):
    matplotlib.pyplot.subplot(1, 4, 1 + i)
    matplotlib.pyplot.imshow(locals()[image]);
    matplotlib.pyplot.title(image);

s = PoissonSolver()
step1(s, "sky.png", "rings.png", [10, 50], 1.0, 0.25)
step2(s)
coeffs = [gen_poisson_coeffs(s, s.ref_soa[i], s.target_soa[i]) for i in range(3)]

X_rgb = [scipy.sparse.linalg.spsolve(A.tocsr(), B) for A, B in coeffs]

seamless_clone_soa = s.ref_soa.copy()

for i in range(3):
    seamless_clone_soa[i][s.strict_interior_indices] = clip(X_rgb[i], 0, 255)

ref = join_rgb(s.ref_soa)
target = join_rgb(s.target_soa)
seamless_clone = join_rgb(seamless_clone_soa)
naive_clone = s.naive_clone

figsize(18, 4)
for i, image in enumerate(["ref", "target", "naive_clone", "seamless_clone"]):
    matplotlib.pyplot.subplot(1, 4, 1 + i)
    matplotlib.pyplot.imshow(locals()[image]);
    matplotlib.pyplot.title(image);

arranged_coeffs = array(split(array(coeffs), 2, axis=1)).squeeze()
A = scipy.sparse.block_diag(arranged_coeffs[0]).tocsr()
B = concatenate(arranged_coeffs[1])

X_rgb = reshape(scipy.sparse.linalg.spsolve(A, B), (3, s.num_strict_interior_pixels))

seamless_clone_soa = s.ref_soa.copy()

for i in range(3):
    seamless_clone_soa[i][s.strict_interior_indices] = clip(X_rgb[i], 0, 255)

ref = join_rgb(s.ref_soa)
target = join_rgb(s.target_soa)
naive_clone = s.naive_clone
seamless_clone_full = join_rgb(seamless_clone_soa)

figsize(18, 4)
for i, image in enumerate(["ref", "target", "naive_clone", "seamless_clone_full"]):
    matplotlib.pyplot.subplot(1, 4, 1 + i)
    matplotlib.pyplot.imshow(locals()[image])
    matplotlib.pyplot.title(image)