import numpy as np
import scipy.interpolate
import scipy.ndimage
from sklearn.feature_extraction.image import extract_patches_2d, reconstruct_from_patches_2d


def _calc_patch_grid_dims(shape, patch_size, patch_stride):
    x_w, x_h, x_c = shape
    num_rows = 1 + (x_h - patch_size) // patch_stride
    num_cols = 1 + (x_w - patch_size) // patch_stride
    return num_rows, num_cols


def make_patch_grid(x, patch_size, patch_stride=1):
    '''x shape: (num_channels, rows, cols)'''
    x = x.transpose(2, 1, 0)
    patches = extract_patches_2d(x, (patch_size, patch_size))
    x_w, x_h, x_c  = x.shape
    num_rows, num_cols = _calc_patch_grid_dims(x.shape, patch_size, patch_stride)
    patches = patches.reshape((num_rows, num_cols, patch_size, patch_size, x_c))
    patches = patches.transpose((0, 1, 4, 2, 3))
    #patches = np.rollaxis(patches, -1, 2)
    return patches


def combine_patches_grid(in_patches, out_shape):
    '''Reconstruct an image from these `patches`

    input shape: (rows, cols, channels, patch_row, patch_col)
    '''
    num_rows, num_cols = in_patches.shape[:2]
    num_channels = in_patches.shape[-3]
    patch_size = in_patches.shape[-1]
    num_patches = num_rows * num_cols
    in_patches = np.reshape(in_patches, (num_patches, num_channels, patch_size, patch_size))  # (patches, channels, pr, pc)
    in_patches = np.transpose(in_patches, (0, 2, 3, 1)) # (patches, p, p, channels)
    recon = reconstruct_from_patches_2d(in_patches, out_shape)
    return recon.transpose(2, 1, 0).astype(np.float32)


class PatchMatcher(object):
    '''A matcher of image patches inspired by the PatchMatch algorithm.

    image shape: (width, height, channels)
    '''
    def __init__(self, input_shape, target_img, patch_size=1, patch_stride=1, jump_size=0.5,
            num_propagation_steps=5, num_random_steps=5, random_max_radius=1.0, random_scale=0.5):
        self.input_shape = input_shape
        self.patch_size = patch_size
        self.patch_stride = patch_stride
        self.jump_size = jump_size
        self.num_propagation_steps = num_propagation_steps
        self.num_random_steps = num_random_steps
        self.random_max_radius = random_max_radius
        self.random_scale = random_scale
        self.num_input_rows, self.num_input_cols = _calc_patch_grid_dims(input_shape, patch_size, patch_stride)
        self.target_patches = make_patch_grid(target_img, patch_size)
        self.target_patches_normed = self.normalize_patches(self.target_patches)
        self.coords = np.random.uniform(0.0, 1.0,  # TODO: switch to pixels
            (2, self.num_input_rows, self.num_input_cols))# * [[[self.num_input_rows]],[[self.num_input_cols]]]
        self.similarity = np.zeros(input_shape[:2:-1], dtype=np.float32)
        self.min_propagration_row = 1.0 / self.num_input_rows
        self.min_propagration_col = 1.0 / self.num_input_cols
        self.delta_row = np.array([[[self.min_propagration_row]], [[0.0]]])
        self.delta_col = np.array([[[0.0]], [[self.min_propagration_col]]])

    def update(self, input_img, reverse_propagation=False):
        input_patches = self.get_patches_for(input_img)
        self.update_with_patches(self.normalize_patches(input_patches), reverse_propagation=reverse_propagation)

    def update_with_patches(self, input_patches, reverse_propagation=False):
        self._propagate(input_patches, reverse_propagation=reverse_propagation)
        self._random_update(input_patches)

    def get_patches_for(self, img):
        return make_patch_grid(img, self.patch_size);

    def normalize_patches(self, patches):
        norm = np.sqrt(np.sum(np.square(patches), axis=(2, 3, 4), keepdims=True))
        return patches / norm

    def _propagate(self, input_patches, reverse_propagation=False):
        if reverse_propagation:
            roll_direction = 1
        else:
            roll_direction = -1
        sign = float(roll_direction)
        for step_i in range(self.num_propagation_steps):
            new_coords = self.clip_coords(np.roll(self.coords, roll_direction, 1) + self.delta_row * sign)
            coords_row, similarity_row = self.eval_state(new_coords, input_patches)
            new_coords = self.clip_coords(np.roll(self.coords, roll_direction, 2) + self.delta_col * sign)
            coords_col, similarity_col = self.eval_state(new_coords, input_patches)
            self.coords, self.similarity = self.take_best(coords_row, similarity_row, coords_col, similarity_col)

    def _random_update(self, input_patches):
        for alpha in range(1, self.num_random_steps + 1):  # NOTE this should actually stop when the move is < 1
            new_coords = self.clip_coords(self.coords + np.random.uniform(-self.random_max_radius, self.random_max_radius, self.coords.shape) * self.random_scale ** alpha)
            self.coords, self.similarity = self.eval_state(new_coords, input_patches)

    def eval_state(self, new_coords, input_patches):
        new_similarity = self.patch_similarity(input_patches, new_coords)
        delta_similarity = new_similarity - self.similarity
        coords = np.where(delta_similarity > 0, new_coords, self.coords)
        best_similarity = np.where(delta_similarity > 0, new_similarity, self.similarity)
        return coords, best_similarity

    def take_best(self, coords_a, similarity_a, coords_b, similarity_b):
        delta_similarity = similarity_a - similarity_b
        best_coords = np.where(delta_similarity > 0, coords_a, coords_b)
        best_similarity = np.where(delta_similarity > 0, similarity_a, similarity_b)
        return best_coords, best_similarity

    def patch_similarity(self, source, coords):
        '''Check the similarity of the patches specified in coords.'''
        target_vals = self.lookup_coords(self.target_patches_normed, coords)
        err = source * target_vals
        return np.sum(err, axis=(2, 3, 4))

    def clip_coords(self, coords):
        # TODO: should this all be in pixel space?
        coords = np.clip(coords, 0.0, 1.0)
        return coords

    def lookup_coords(self, x, coords):
        x_shape = np.expand_dims(np.expand_dims(x.shape, -1), -1)
        i_coords = np.round(coords * (x_shape[:2] - 1)).astype('int32')
        return x[i_coords[0], i_coords[1]]

    def get_reconstruction(self, patches=None, combined=None):
        if combined is not None:
            patches = make_patch_grid(combined, self.patch_size)
        if patches is None:
            patches = self.target_patches
        patches = self.lookup_coords(patches, self.coords)
        recon = combine_patches_grid(patches, self.input_shape)
        return recon

    def scale(self, new_shape, new_target_img):
        '''Create a new matcher of the given shape and replace its
        state with a scaled up version of the current matcher's state.
        '''
        new_matcher = PatchMatcher(new_shape, new_target_img, patch_size=self.patch_size,
                patch_stride=self.patch_stride, jump_size=self.jump_size,
                num_propagation_steps=self.num_propagation_steps,
                num_random_steps=self.num_random_steps,
                random_max_radius=self.random_max_radius,
                random_scale=self.random_scale)
        new_matcher.coords = congrid(self.coords, new_matcher.coords.shape, method='neighbour')
        new_matcher.similarity = congrid(self.similarity, new_matcher.coords.shape, method='neighbour')
        return new_matcher


def congrid(a, newdims, method='linear', centre=False, minusone=False):
    '''Arbitrary resampling of source array to new dimension sizes.
    Currently only supports maintaining the same number of dimensions.
    To use 1-D arrays, first promote them to shape (x,1).

    Uses the same parameters and creates the same co-ordinate lookup points
    as IDL''s congrid routine, which apparently originally came from a VAX/VMS
    routine of the same name.

    method:
    neighbour - closest value from original data
    nearest and linear - uses n x 1-D interpolations using
                         scipy.interpolate.interp1d
    (see Numerical Recipes for validity of use of n 1-D interpolations)
    spline - uses ndimage.map_coordinates

    centre:
    True - interpolation points are at the centres of the bins
    False - points are at the front edge of the bin

    minusone:
    For example- inarray.shape = (i,j) & new dimensions = (x,y)
    False - inarray is resampled by factors of (i/x) * (j/y)
    True - inarray is resampled by(i-1)/(x-1) * (j-1)/(y-1)
    This prevents extrapolation one element beyond bounds of input array.
    '''
    if not a.dtype in [np.float64, np.float32]:
        a = np.cast[float](a)

    m1 = np.cast[int](minusone)
    ofs = np.cast[int](centre) * 0.5
    old = np.array( a.shape )
    ndims = len( a.shape )
    if len( newdims ) != ndims:
        print("[congrid] dimensions error. " \
              "This routine currently only support " \
              "rebinning to the same number of dimensions.")
        return None
    newdims = np.asarray( newdims, dtype=float )
    dimlist = []

    if method == 'neighbour':
        for i in range( ndims ):
            base = np.indices(newdims)[i]
            dimlist.append( (old[i] - m1) / (newdims[i] - m1) \
                            * (base + ofs) - ofs )
        cd = np.array( dimlist ).round().astype(int)
        newa = a[list( cd )]
        return newa

    elif method in ['nearest','linear']:
        # calculate new dims
        for i in range( ndims ):
            base = np.arange( newdims[i] )
            dimlist.append( (old[i] - m1) / (newdims[i] - m1) \
                            * (base + ofs) - ofs )
        # specify old dims
        olddims = [np.arange(i, dtype = np.float) for i in list( a.shape )]

        # first interpolation - for ndims = any
        mint = scipy.interpolate.interp1d( olddims[-1], a, kind=method )
        newa = mint( dimlist[-1] )

        trorder = [ndims - 1] + range( ndims - 1 )
        for i in range( ndims - 2, -1, -1 ):
            newa = newa.transpose( trorder )

            mint = scipy.interpolate.interp1d( olddims[i], newa, kind=method )
            newa = mint( dimlist[i] )

        if ndims > 1:
            # need one more transpose to return to original dimensions
            newa = newa.transpose( trorder )

        return newa
    elif method in ['spline']:
        oslices = [ slice(0,j) for j in old ]
        oldcoords = np.ogrid[oslices]
        nslices = [ slice(0,j) for j in list(newdims) ]
        newcoords = np.mgrid[nslices]

        newcoords_dims = range(np.rank(newcoords))
        #make first index last
        newcoords_dims.append(newcoords_dims.pop(0))
        newcoords_tr = newcoords.transpose(newcoords_dims)
        # makes a view that affects newcoords

        newcoords_tr += ofs

        deltas = (np.asarray(old) - m1) / (newdims - m1)
        newcoords_tr *= deltas

        newcoords_tr -= ofs

        newa = scipy.ndimage.map_coordinates(a, newcoords)
        return newa
    else:
        print("Congrid error: Unrecognized interpolation type.\n", \
              "Currently only \'neighbour\', \'nearest\',\'linear\',", \
              "and \'spline\' are supported.")
        return None


if __name__ == '__main__':
    import sys
    import time
    from scipy.misc import imsave

    from image_analogy.img_utils import load_image, preprocess_image, deprocess_image

    content_image_path, style_image_path, output_prefix = sys.argv[1:]
    jump_size = 1.0
    num_steps = 7
    patch_size = 1
    patch_stride = 1

    feat_chans = 512
    feat_style_shape = (feat_chans, 12, 18)
    feat_style = np.random.uniform(0.0, 1.0, feat_style_shape)
    feat_in_shape = (feat_chans, 17, 10)
    feat_in = np.random.uniform(0.0, 1.0, feat_in_shape)
    matcher = PatchMatcher(feat_in_shape[::-1], feat_style, patch_size=patch_size)
    feat_in_normed = matcher.normalize_patches(matcher.get_patches_for(feat_in))
    for i in range(num_steps):
        matcher.update_with_patches(feat_in_normed)
    r = matcher.get_reconstruction()

    content_img_img = load_image(content_image_path)
    content_n_channels, content_n_rows, content_n_cols = content_img_img.shape[::-1]
    content_img = preprocess_image(content_img_img, content_n_cols, content_n_rows)[0]#.transpose((2,1,0))
    style_img = load_image(style_image_path)
    style_n_channels, style_n_rows, style_n_cols = content_img_img.shape[::-1]
    style_img =  preprocess_image(
        load_image(style_image_path), style_n_cols, style_n_rows)[0]#.transpose((2,1,0))
    pg = make_patch_grid(content_img, patch_size)
    result = combine_patches_grid(pg, content_img.shape[::-1])
    outimg = deprocess_image(result, contrast_percent=0)
    imsave(output_prefix + '_bestre.png', outimg)

    # # #
    matcher = PatchMatcher((content_n_cols, content_n_rows, content_n_channels), style_img, patch_size=patch_size)
    for i in range(num_steps):
        start = time.time()
        matcher.update(content_img, reverse_propagation=bool(i % 2))
        print(matcher.similarity.min(), matcher.similarity.max(), matcher.similarity.mean())
        end = time.time()
        #print end-start
    start = time.time()
    result = matcher.get_reconstruction(patches=matcher.target_patches)
    print(result.shape)
    end = time.time()
    print(end-start)
    outimg = deprocess_image(result, contrast_percent=0)
    # # imsave takes (rows, cols, channels)
    imsave(output_prefix + '_best.png', outimg)

[((17, 14, 17, 61), 'sklearn.feature_extraction.image.extract_patches_2d', 'extract_patches_2d', ({(17, 33, 17, 34): 'x', (17, 36, 17, 60): '(patch_size, patch_size)'}, {}), '(x, (patch_size, patch_size))', False, 'from sklearn.feature_extraction.image import extract_patches_2d, reconstruct_from_patches_2d\n'), ((35, 17, 35, 92), 'numpy.reshape', 'np.reshape', ({(35, 28, 35, 38): 'in_patches', (35, 40, 35, 91): '(num_patches, num_channels, patch_size, patch_size)'}, {}), '(in_patches, (num_patches, num_channels, patch_size, patch_size))', True, 'import numpy as np\n'), ((36, 17, 36, 55), 'numpy.transpose', 'np.transpose', ({(36, 30, 36, 40): 'in_patches', (36, 42, 36, 54): '(0, 2, 3, 1)'}, {}), '(in_patches, (0, 2, 3, 1))', True, 'import numpy as np\n'), ((37, 12, 37, 62), 'sklearn.feature_extraction.image.reconstruct_from_patches_2d', 'reconstruct_from_patches_2d', ({(37, 40, 37, 50): 'in_patches', (37, 52, 37, 61): 'out_shape'}, {}), '(in_patches, out_shape)', False, 'from sklearn.feature_extraction.image import extract_patches_2d, reconstruct_from_patches_2d\n'), ((184, 10, 184, 29), 'numpy.array', 'np.array', ({(184, 20, 184, 27): 'a.shape'}, {}), '(a.shape)', True, 'import numpy as np\n'), ((191, 14, 191, 48), 'numpy.asarray', 'np.asarray', (), '', True, 'import numpy as np\n'), ((271, 17, 271, 62), 'numpy.random.uniform', 'np.random.uniform', ({(271, 35, 271, 38): '0.0', (271, 40, 271, 43): '1.0', (271, 45, 271, 61): 'feat_style_shape'}, {}), '(0.0, 1.0, feat_style_shape)', True, 'import numpy as np\n'), ((273, 14, 273, 56), 'numpy.random.uniform', 'np.random.uniform', ({(273, 32, 273, 35): '0.0', (273, 37, 273, 40): '1.0', (273, 42, 273, 55): 'feat_in_shape'}, {}), '(0.0, 1.0, feat_in_shape)', True, 'import numpy as np\n'), ((280, 22, 280, 52), 'image_analogy.img_utils.load_image', 'load_image', ({(280, 33, 280, 51): 'content_image_path'}, {}), '(content_image_path)', False, 'from image_analogy.img_utils import load_image, preprocess_image, deprocess_image\n'), ((283, 16, 283, 44), 'image_analogy.img_utils.load_image', 'load_image', ({(283, 27, 283, 43): 'style_image_path'}, {}), '(style_image_path)', False, 'from image_analogy.img_utils import load_image, preprocess_image, deprocess_image\n'), ((289, 13, 289, 56), 'image_analogy.img_utils.deprocess_image', 'deprocess_image', (), '', False, 'from image_analogy.img_utils import load_image, preprocess_image, deprocess_image\n'), ((290, 4, 290, 49), 'scipy.misc.imsave', 'imsave', ({(290, 11, 290, 40): "(output_prefix + '_bestre.png')", (290, 42, 290, 48): 'outimg'}, {}), "(output_prefix + '_bestre.png', outimg)", False, 'from scipy.misc import imsave\n'), ((300, 12, 300, 23), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((303, 10, 303, 21), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((305, 13, 305, 56), 'image_analogy.img_utils.deprocess_image', 'deprocess_image', (), '', False, 'from image_analogy.img_utils import load_image, preprocess_image, deprocess_image\n'), ((307, 4, 307, 47), 'scipy.misc.imsave', 'imsave', ({(307, 11, 307, 38): "(output_prefix + '_best.png')", (307, 40, 307, 46): 'outimg'}, {}), "(output_prefix + '_best.png', outimg)", False, 'from scipy.misc import imsave\n'), ((59, 22, 60, 58), 'numpy.random.uniform', 'np.random.uniform', ({(59, 40, 59, 43): '0.0', (59, 45, 59, 48): '1.0', (60, 12, 60, 57): '(2, self.num_input_rows, self.num_input_cols)'}, {}), '(0.0, 1.0, (2, self.num_input_rows, self.num_input_cols))', True, 'import numpy as np\n'), ((61, 26, 61, 72), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((64, 25, 64, 75), 'numpy.array', 'np.array', ({(64, 34, 64, 74): '[[[self.min_propagration_row]], [[0.0]]]'}, {}), '([[[self.min_propagration_row]], [[0.0]]])', True, 'import numpy as np\n'), ((65, 25, 65, 75), 'numpy.array', 'np.array', ({(65, 34, 65, 74): '[[[0.0]], [[self.min_propagration_col]]]'}, {}), '([[[0.0]], [[self.min_propagration_col]]])', True, 'import numpy as np\n'), ((103, 17, 103, 72), 'numpy.where', 'np.where', ({(103, 26, 103, 46): 'delta_similarity > 0', (103, 48, 103, 58): 'new_coords', (103, 60, 103, 71): 'self.coords'}, {}), '(delta_similarity > 0, new_coords, self.coords)', True, 'import numpy as np\n'), ((104, 26, 104, 89), 'numpy.where', 'np.where', ({(104, 35, 104, 55): 'delta_similarity > 0', (104, 57, 104, 71): 'new_similarity', (104, 73, 104, 88): 'self.similarity'}, {}), '(delta_similarity > 0, new_similarity, self.similarity)', True, 'import numpy as np\n'), ((109, 22, 109, 72), 'numpy.where', 'np.where', ({(109, 31, 109, 51): 'delta_similarity > 0', (109, 53, 109, 61): 'coords_a', (109, 63, 109, 71): 'coords_b'}, {}), '(delta_similarity > 0, coords_a, coords_b)', True, 'import numpy as np\n'), ((110, 26, 110, 84), 'numpy.where', 'np.where', ({(110, 35, 110, 55): 'delta_similarity > 0', (110, 57, 110, 69): 'similarity_a', (110, 71, 110, 83): 'similarity_b'}, {}), '(delta_similarity > 0, similarity_a, similarity_b)', True, 'import numpy as np\n'), ((117, 15, 117, 42), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((121, 17, 121, 42), 'numpy.clip', 'np.clip', ({(121, 25, 121, 31): 'coords', (121, 33, 121, 36): '0.0', (121, 38, 121, 41): '1.0'}, {}), '(coords, 0.0, 1.0)', True, 'import numpy as np\n'), ((282, 18, 282, 83), 'image_analogy.img_utils.preprocess_image', 'preprocess_image', ({(282, 35, 282, 50): 'content_img_img', (282, 52, 282, 66): 'content_n_cols', (282, 68, 282, 82): 'content_n_rows'}, {}), '(content_img_img, content_n_cols, content_n_rows)', False, 'from image_analogy.img_utils import load_image, preprocess_image, deprocess_image\n'), ((295, 16, 295, 27), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((298, 14, 298, 25), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((125, 33, 125, 60), 'numpy.expand_dims', 'np.expand_dims', ({(125, 48, 125, 55): 'x.shape', (125, 57, 125, 59): '-1'}, {}), '(x.shape, -1)', True, 'import numpy as np\n'), ((286, 8, 286, 36), 'image_analogy.img_utils.load_image', 'load_image', ({(286, 19, 286, 35): 'style_image_path'}, {}), '(style_image_path)', False, 'from image_analogy.img_utils import load_image, preprocess_image, deprocess_image\n'), ((79, 30, 79, 48), 'numpy.square', 'np.square', ({(79, 40, 79, 47): 'patches'}, {}), '(patches)', True, 'import numpy as np\n'), ((126, 19, 126, 55), 'numpy.round', 'np.round', ({(126, 28, 126, 54): 'coords * (x_shape[:2] - 1)'}, {}), '(coords * (x_shape[:2] - 1))', True, 'import numpy as np\n'), ((196, 19, 196, 38), 'numpy.indices', 'np.indices', ({(196, 30, 196, 37): 'newdims'}, {}), '(newdims)', True, 'import numpy as np\n'), ((206, 19, 206, 42), 'numpy.arange', 'np.arange', ({(206, 30, 206, 40): 'newdims[i]'}, {}), '(newdims[i])', True, 'import numpy as np\n'), ((210, 19, 210, 49), 'numpy.arange', 'np.arange', (), '', True, 'import numpy as np\n'), ((89, 42, 89, 81), 'numpy.roll', 'np.roll', ({(89, 50, 89, 61): 'self.coords', (89, 63, 89, 77): 'roll_direction', (89, 79, 89, 80): '1'}, {}), '(self.coords, roll_direction, 1)', True, 'import numpy as np\n'), ((91, 42, 91, 81), 'numpy.roll', 'np.roll', ({(91, 50, 91, 61): 'self.coords', (91, 63, 91, 77): 'roll_direction', (91, 79, 91, 80): '2'}, {}), '(self.coords, roll_direction, 2)', True, 'import numpy as np\n'), ((234, 31, 234, 49), 'numpy.rank', 'np.rank', ({(234, 39, 234, 48): 'newcoords'}, {}), '(newcoords)', True, 'import numpy as np\n'), ((97, 56, 97, 141), 'numpy.random.uniform', 'np.random.uniform', ({(97, 74, 97, 97): '-self.random_max_radius', (97, 99, 97, 121): 'self.random_max_radius', (97, 123, 97, 140): 'self.coords.shape'}, {}), '(-self.random_max_radius, self.random_max_radius, self.\n    coords.shape)', True, 'import numpy as np\n'), ((199, 13, 199, 32), 'numpy.array', 'np.array', ({(199, 23, 199, 30): 'dimlist'}, {}), '(dimlist)', True, 'import numpy as np\n'), ((242, 18, 242, 33), 'numpy.asarray', 'np.asarray', ({(242, 29, 242, 32): 'old'}, {}), '(old)', True, 'import numpy as np\n')]