#!/usr/bin/env python
# -*- coding: utf-8 -*-are not covered by the UCLB ACP-A Licence,

from __future__ import absolute_import, division, print_function
import tensorflow as tf

def bilinear_sampler_1d_h(input_images, x_offset, wrap_mode='border', name='bilinear_sampler', **kwargs):
    '''
    一维双线性采样: x_offset--输入X上偏移量的图
    重复函数 : 先将一维的x后面扩展一个维度, 然后在扩展的维度上复制相应的值, 随后将其转成一维的值, exsamples:[1,2,3] --> [1,1,2,2,3,3]
    '''
    def _repeat(x, n_repeats):
        with tf.variable_scope('_repeat'):
            rep = tf.tile(tf.expand_dims(x, 1), [1, n_repeats])
            return tf.reshape(rep, [-1])

    def _interpolate(im, x, y): #插值函数
        with tf.variable_scope('_interpolate'):

            # handle both texture border types
            _edge_size = 0
            # 如果包围方式是border, 那么边界长度是1, 在h和w维两侧加一排0
            if _wrap_mode == 'border':
                _edge_size = 1
                im = tf.pad(im, [[0, 0], [1, 1], [1, 1], [0, 0]], mode='CONSTANT')
                x = x + _edge_size
                y = y + _edge_size
            elif _wrap_mode == 'edge':
                _edge_size = 0
            else:
                return None

            # 修剪偏移量x, 让它在0到width-1+2*edge_size之间(因为偏移量不能太大,要小于等于padding之后).
            x = tf.clip_by_value(x, 0.0,  _width_f - 1 + 2 * _edge_size)

            # 向下取整x,y然后x加1向上取整x
            x0_f = tf.floor(x)
            y0_f = tf.floor(y)
            x1_f = x0_f + 1
            # 将向下取整的x y变成整数, 向上取整的x不能大于padding之后的宽度减1
            # cast: 类型转换
            x0 = tf.cast(x0_f, tf.int32)
            y0 = tf.cast(y0_f, tf.int32)
            x1 = tf.cast(tf.minimum(x1_f,  _width_f - 1 + 2 * _edge_size), tf.int32)
            # 第二维也就是宽度维的宽是padding之后的宽
            dim2 = (_width + 2 * _edge_size)
            # 第一维也就是图像维的宽是padding之后的分辨率
            dim1 = (_width + 2 * _edge_size) * (_height + 2 * _edge_size)
            # 计算偏移量索引的基,先得到[0,1,2,...,batch],再将它乘宽度,变成
            # [0,dim1,2*dim1,...,batch*dim1],然后重复原图分辨率,变成
            # [0,0,......,0,dim1,dim1,......,dim1,2*dim1,2*dim1,......,2*dim1 . . batch * dim, batch * dim, ......, batch * dim]
            # 这样就变成基底了,表达的是有batch个图的基
            base = _repeat(tf.range(_num_batch) * dim1, _height * _width)
            # 将y的偏移乘以dim2,也就是乘以宽度,这样就得到加上y之后的基
            # y0是[0,0,...,0,1,1,....,1, . . h + 2 * e, h + 2 * e, ..., h + 2 * e]
            # 乘了dim2之后变成
            # [0, 0, ..., 0, w+2*e, w+2*e, ..., w+2*e, . . (h + 2 * e) * (w + 2 * e), ..., (h + 2 * e) * (w + 2 * e)]
            # 加上base之后得到了考虑了batch,height之后的索引
            base_y0 = base + y0 * dim2
            # 这个索引加上向上下取整的x索引和向上取整的x索引就得到了现在点的左侧点和右侧点
            idx_l = base_y0 + x0
            idx_r = base_y0 + x1
            # 将图变成[batch*w*h,channel]的形状
            im_flat = tf.reshape(im, tf.stack([-1, _num_channels]))
            # 利用tf.gather根据左右侧点的索引重新排列图,得到重排之后的左右像素
            pix_l = tf.gather(im_flat, idx_l)
            pix_r = tf.gather(im_flat, idx_r)
            # 计算双线性差值的系数x1-1和x-x0
            weight_l = tf.expand_dims(x1_f - x, 1)
            weight_r = tf.expand_dims(x - x0_f, 1)
            # 利用双线性差值方法计算像素值
            return weight_l * pix_l + weight_r * pix_r

    # get_disp函数生成视差图后,调用插值函数获得更好的图.
    def _transform(input_images, x_offset):
        '''
        转换函数首先调用meshgrid生成关于X轴和Y轴的索引
        exsamples:
        假设_width=3，经过linspace(0.0,_width_f-1.0,_width)是[ 0., 1., 2.]。height同理
        >>> x = tf.linspace(0.0, 2.0, 3)
        >>> sess.run(x)
        array([0., 1., 2. ], dtype = float32)
        >>> x = tf.linspace(0.0, 2.0, 3)
        >>> y = tf.linspace(0.0, 4.0, 5)
        >>> x_t, y_t = tf.meshgrid(x, y)
        >>> sess.run(x_t)
        array([0., 1., 2.],
            [0., 1., 2.],
            [0., 1., 2.],
            [0., 1., 2.],
            [0., 1., 2.]], dtype=float32)
        >>> sess.run(y_t)
        array([0., 0., 0.],
            [1., 1., 1.],
            [2., 2., 2.],
            [3., 3., 3.],
            [4., 4., 4.]], dtype=float32)
        >>> x_t_flat = tf.reshape(x_t, (1, -1))
        >>> y_t_flat = tf.reshape(y_t, (1, -1))
        >>> sess.run(x_t_flat)
        array([[0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2.]], dtype=float32)
        >>> sess.run(y_t_flat)
        array([[0., 0., 0., 1., 1., 1., 2., 2., 2., 3., 3., 3., 4., 4., 4.]], dtype=float32)
        >>> x_t_flat = tf.tile(x_t_flat, tf.stack([2,1]))
        >>> sess.run(x_t_flat)
        arraay([[0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2.], [0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2.]], dtype=float32)
        >>> x_t_flat = tf.reshape(x_t_flat, (1, -1))
        >>> sess.run(x_t_flat)
        array([[0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2., 0., 1., 2.]], dtype=float32)
        '''
        with tf.variable_scope('transform'):
            # grid of (x_t, y_t, 1), eq (1) in ref [1]
            x_t, y_t = tf.meshgrid(tf.linspace(0.0,   _width_f - 1.0,  _width),
                                   tf.linspace(0.0 , _height_f - 1.0 , _height))

            x_t_flat = tf.reshape(x_t, (1, -1))
            y_t_flat = tf.reshape(y_t, (1, -1))

            x_t_flat = tf.tile(x_t_flat, tf.stack([_num_batch, 1]))
            y_t_flat = tf.tile(y_t_flat, tf.stack([_num_batch, 1]))

            x_t_flat = tf.reshape(x_t_flat, [-1])
            y_t_flat = tf.reshape(y_t_flat, [-1])

            x_t_flat = x_t_flat + tf.reshape(x_offset, [-1]) * _width_f

            input_transformed = _interpolate(input_images, x_t_flat, y_t_flat)

            output = tf.reshape(
                input_transformed, tf.stack([_num_batch, _height, _width, _num_channels]))

            return output


    with tf.variable_scope(name):
        '''
        [num_batch, height, width, num_channels]
        '''
        _num_batch    = tf.shape(input_images)[0]
        _height       = tf.shape(input_images)[1]
        _width        = tf.shape(input_images)[2]
        _num_channels = tf.shape(input_images)[3]

        _height_f = tf.cast(_height, tf.float32)
        _width_f  = tf.cast(_width,  tf.float32)

        _wrap_mode = wrap_mode

        output = _transform(input_images, x_offset)
        return output

[((135, 9, 135, 32), 'tensorflow.variable_scope', 'tf.variable_scope', ({(135, 27, 135, 31): 'name'}, {}), '(name)', True, 'import tensorflow as tf\n'), ((144, 20, 144, 48), 'tensorflow.cast', 'tf.cast', ({(144, 28, 144, 35): '_height', (144, 37, 144, 47): 'tf.float32'}, {}), '(_height, tf.float32)', True, 'import tensorflow as tf\n'), ((145, 20, 145, 48), 'tensorflow.cast', 'tf.cast', ({(145, 28, 145, 34): '_width', (145, 37, 145, 47): 'tf.float32'}, {}), '(_width, tf.float32)', True, 'import tensorflow as tf\n'), ((13, 13, 13, 41), 'tensorflow.variable_scope', 'tf.variable_scope', ({(13, 31, 13, 40): '"""_repeat"""'}, {}), "('_repeat')", True, 'import tensorflow as tf\n'), ((15, 19, 15, 40), 'tensorflow.reshape', 'tf.reshape', ({(15, 30, 15, 33): 'rep', (15, 35, 15, 39): '[-1]'}, {}), '(rep, [-1])', True, 'import tensorflow as tf\n'), ((18, 13, 18, 46), 'tensorflow.variable_scope', 'tf.variable_scope', ({(18, 31, 18, 45): '"""_interpolate"""'}, {}), "('_interpolate')", True, 'import tensorflow as tf\n'), ((34, 16, 34, 72), 'tensorflow.clip_by_value', 'tf.clip_by_value', ({(34, 33, 34, 34): 'x', (34, 36, 34, 39): '0.0', (34, 42, 34, 71): '_width_f - 1 + 2 * _edge_size'}, {}), '(x, 0.0, _width_f - 1 + 2 * _edge_size)', True, 'import tensorflow as tf\n'), ((37, 19, 37, 30), 'tensorflow.floor', 'tf.floor', ({(37, 28, 37, 29): 'x'}, {}), '(x)', True, 'import tensorflow as tf\n'), ((38, 19, 38, 30), 'tensorflow.floor', 'tf.floor', ({(38, 28, 38, 29): 'y'}, {}), '(y)', True, 'import tensorflow as tf\n'), ((42, 17, 42, 40), 'tensorflow.cast', 'tf.cast', ({(42, 25, 42, 29): 'x0_f', (42, 31, 42, 39): 'tf.int32'}, {}), '(x0_f, tf.int32)', True, 'import tensorflow as tf\n'), ((43, 17, 43, 40), 'tensorflow.cast', 'tf.cast', ({(43, 25, 43, 29): 'y0_f', (43, 31, 43, 39): 'tf.int32'}, {}), '(y0_f, tf.int32)', True, 'import tensorflow as tf\n'), ((66, 20, 66, 45), 'tensorflow.gather', 'tf.gather', ({(66, 30, 66, 37): 'im_flat', (66, 39, 66, 44): 'idx_l'}, {}), '(im_flat, idx_l)', True, 'import tensorflow as tf\n'), ((67, 20, 67, 45), 'tensorflow.gather', 'tf.gather', ({(67, 30, 67, 37): 'im_flat', (67, 39, 67, 44): 'idx_r'}, {}), '(im_flat, idx_r)', True, 'import tensorflow as tf\n'), ((69, 23, 69, 50), 'tensorflow.expand_dims', 'tf.expand_dims', ({(69, 38, 69, 46): 'x1_f - x', (69, 48, 69, 49): '1'}, {}), '(x1_f - x, 1)', True, 'import tensorflow as tf\n'), ((70, 23, 70, 50), 'tensorflow.expand_dims', 'tf.expand_dims', ({(70, 38, 70, 46): 'x - x0_f', (70, 48, 70, 49): '1'}, {}), '(x - x0_f, 1)', True, 'import tensorflow as tf\n'), ((111, 13, 111, 43), 'tensorflow.variable_scope', 'tf.variable_scope', ({(111, 31, 111, 42): '"""transform"""'}, {}), "('transform')", True, 'import tensorflow as tf\n'), ((116, 23, 116, 47), 'tensorflow.reshape', 'tf.reshape', ({(116, 34, 116, 37): 'x_t', (116, 39, 116, 46): '(1, -1)'}, {}), '(x_t, (1, -1))', True, 'import tensorflow as tf\n'), ((117, 23, 117, 47), 'tensorflow.reshape', 'tf.reshape', ({(117, 34, 117, 37): 'y_t', (117, 39, 117, 46): '(1, -1)'}, {}), '(y_t, (1, -1))', True, 'import tensorflow as tf\n'), ((122, 23, 122, 49), 'tensorflow.reshape', 'tf.reshape', ({(122, 34, 122, 42): 'x_t_flat', (122, 44, 122, 48): '[-1]'}, {}), '(x_t_flat, [-1])', True, 'import tensorflow as tf\n'), ((123, 23, 123, 49), 'tensorflow.reshape', 'tf.reshape', ({(123, 34, 123, 42): 'y_t_flat', (123, 44, 123, 48): '[-1]'}, {}), '(y_t_flat, [-1])', True, 'import tensorflow as tf\n'), ((139, 24, 139, 46), 'tensorflow.shape', 'tf.shape', ({(139, 33, 139, 45): 'input_images'}, {}), '(input_images)', True, 'import tensorflow as tf\n'), ((140, 24, 140, 46), 'tensorflow.shape', 'tf.shape', ({(140, 33, 140, 45): 'input_images'}, {}), '(input_images)', True, 'import tensorflow as tf\n'), ((141, 24, 141, 46), 'tensorflow.shape', 'tf.shape', ({(141, 33, 141, 45): 'input_images'}, {}), '(input_images)', True, 'import tensorflow as tf\n'), ((142, 24, 142, 46), 'tensorflow.shape', 'tf.shape', ({(142, 33, 142, 45): 'input_images'}, {}), '(input_images)', True, 'import tensorflow as tf\n'), ((14, 26, 14, 46), 'tensorflow.expand_dims', 'tf.expand_dims', ({(14, 41, 14, 42): 'x', (14, 44, 14, 45): '1'}, {}), '(x, 1)', True, 'import tensorflow as tf\n'), ((25, 21, 25, 82), 'tensorflow.pad', 'tf.pad', (), '', True, 'import tensorflow as tf\n'), ((44, 25, 44, 73), 'tensorflow.minimum', 'tf.minimum', ({(44, 36, 44, 40): 'x1_f', (44, 43, 44, 72): '_width_f - 1 + 2 * _edge_size'}, {}), '(x1_f, _width_f - 1 + 2 * _edge_size)', True, 'import tensorflow as tf\n'), ((64, 37, 64, 66), 'tensorflow.stack', 'tf.stack', ({(64, 46, 64, 65): '[-1, _num_channels]'}, {}), '([-1, _num_channels])', True, 'import tensorflow as tf\n'), ((113, 35, 113, 78), 'tensorflow.linspace', 'tf.linspace', ({(113, 47, 113, 50): '0.0', (113, 54, 113, 68): '_width_f - 1.0', (113, 71, 113, 77): '_width'}, {}), '(0.0, _width_f - 1.0, _width)', True, 'import tensorflow as tf\n'), ((114, 35, 114, 79), 'tensorflow.linspace', 'tf.linspace', ({(114, 47, 114, 50): '0.0', (114, 53, 114, 68): '_height_f - 1.0', (114, 71, 114, 78): '_height'}, {}), '(0.0, _height_f - 1.0, _height)', True, 'import tensorflow as tf\n'), ((119, 41, 119, 66), 'tensorflow.stack', 'tf.stack', ({(119, 50, 119, 65): '[_num_batch, 1]'}, {}), '([_num_batch, 1])', True, 'import tensorflow as tf\n'), ((120, 41, 120, 66), 'tensorflow.stack', 'tf.stack', ({(120, 50, 120, 65): '[_num_batch, 1]'}, {}), '([_num_batch, 1])', True, 'import tensorflow as tf\n'), ((130, 35, 130, 89), 'tensorflow.stack', 'tf.stack', ({(130, 44, 130, 88): '[_num_batch, _height, _width, _num_channels]'}, {}), '([_num_batch, _height, _width, _num_channels])', True, 'import tensorflow as tf\n'), ((53, 27, 53, 47), 'tensorflow.range', 'tf.range', ({(53, 36, 53, 46): '_num_batch'}, {}), '(_num_batch)', True, 'import tensorflow as tf\n'), ((125, 34, 125, 60), 'tensorflow.reshape', 'tf.reshape', ({(125, 45, 125, 53): 'x_offset', (125, 55, 125, 59): '[-1]'}, {}), '(x_offset, [-1])', True, 'import tensorflow as tf\n')]