# import tensorflow as tf
from .basics import *
import numpy as np
import imageio
from .flowlib import flow_to_image
from .flowlib import read_flow, evaluate_flow

# from warp import tf_warp
import sys
# sys.path.append('/home/user321/tf_flownet2-master/FlowNet2_src/')
modelspath = '../flow_pretrain_np/'

# from flow_warp import flow_warp

def gather_nd(img, idx):
    """
    same as tf.gather_nd in pytorch
    """
    idx = idx.long()
    idx1, idx2, idx3 = idx.chunk(3, dim=3)
    output = img[idx1, idx2, idx3].squeeze(3)
    return output

def get_pixel_value(img, x, y):
    """
    Utility function to get pixel value for coordinate
    vectors x and y from a  4D tensor image.
    Input
    -----
    - img: tensor of shape (B, H, W, C)
    - x: flattened tensor of shape (B*H*W, )
    - y: flattened tensor of shape (B*H*W, )
    Returns
    -------
    - output: tensor of shape (B, H, W, C)
    """
    shape = x.size()
    batch_size = shape[0]
    height = shape[1]
    width = shape[2]

    batch_idx = torch.arange(0, batch_size).int().cuda()
    batch_idx = batch_idx.view(batch_size, 1, 1)
    b = batch_idx.repeat(1, height, width)

    indices = torch.stack([b, y, x], 3)
    # print(indices.size())

    return gather_nd(img, indices)


Backward_tensorGrid = [{} for i in range(8)]

def torch_warp(tensorInput, tensorFlow):
    device_id = tensorInput.device.index
    if str(tensorFlow.size()) not in Backward_tensorGrid[device_id]:
            tensorHorizontal = torch.linspace(-1.0, 1.0, tensorFlow.size(3)).view(1, 1, 1, tensorFlow.size(3)).expand(tensorFlow.size(0), -1, tensorFlow.size(2), -1)
            tensorVertical = torch.linspace(-1.0, 1.0, tensorFlow.size(2)).view(1, 1, tensorFlow.size(2), 1).expand(tensorFlow.size(0), -1, -1, tensorFlow.size(3))
            Backward_tensorGrid[device_id][str(tensorFlow.size())] = torch.cat([ tensorHorizontal, tensorVertical ], 1).cuda().to(device_id)
            # B, C, H, W = tensorInput.size()
            # xx = torch.arange(0, W).view(1,-1).repeat(H,1)
            # yy = torch.arange(0, H).view(-1,1).repeat(1,W)
            # xx = xx.view(1,1,H,W).repeat(B,1,1,1)
            # yy = yy.view(1,1,H,W).repeat(B,1,1,1)
            # Backward_tensorGrid[device_id][str(tensorFlow.size())] = Variable(torch.cat([xx, yy], 1).float().cuda()).to(device_id)

    tensorFlow = torch.cat([tensorFlow[:, 0:1, :, :] / ((tensorInput.size(3) - 1.0) / 2.0), tensorFlow[:, 1:2, :, :] / ((tensorInput.size(2) - 1.0) / 2.0) ], 1)

    return torch.nn.functional.grid_sample(input=tensorInput, grid=(Backward_tensorGrid[device_id][str(tensorFlow.size())] + tensorFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='border')
# # end
# Backward_tensorGrid = [{} for i in range(8)]
# vgrids = {}
# def torch_warp(x, flo):
#     """
#     warp an image/tensor (im2) back to im1, according to the optical flow
#     x: [B, C, H, W] (im2)
#     flo: [B, 2, H, W] flow
#     """
#     torch.backends.cudnn.enabled = True
#     B, C, H, W = x.size()
#     if str(x.size()) not in vgrids:
#         # mesh grid 
#         xx = torch.arange(0, W).view(1,-1).repeat(H,1)
#         yy = torch.arange(0, H).view(-1,1).repeat(1,W)
#         xx = xx.view(1,1,H,W).repeat(B,1,1,1)
#         yy = yy.view(1,1,H,W).repeat(B,1,1,1)
#         grid = torch.cat((xx,yy),1).float()
#         # if x.is_cuda:
#         grid = grid.cuda()
#         vgrids[str(x.size())] = Variable(grid)
#     vgrid = vgrids[str(x.size())] + flo

#     # scale grid to [-1,1]
#     vgrid[:,0,:,:] = 2.0*vgrid[:,0,:,:].clone() / max(W-1,1)-1.0
#     vgrid[:,1,:,:] = 2.0*vgrid[:,1,:,:].clone() / max(H-1,1)-1.0
#     vgrid = vgrid.permute(0,2,3,1)
    

#     output = torch.nn.functional.grid_sample(x, grid=vgrid, mode='bilinear', padding_mode='border')
#     mask = torch.autograd.Variable(torch.ones(x.size())).cuda()
#     mask = torch.nn.functional.grid_sample(mask, grid=vgrid, mode='bilinear', padding_mode='border')

#     # if W==128:
#         # np.save('mask.npy', mask.cpu().data.numpy())
#         # np.save('warp.npy', output.cpu().data.numpy())
    
#     mask[mask<0.9999] = 0
#     mask[mask>0] = 1

#     return output*mask

def log10(x):
    numerator = torch.log(x)
    denominator = torch.log(10)
    return numerator / denominator


def flow_warp(im, flow):
    warp = torch_warp(im, flow)

    return warp


def loadweightformnp(layername):
    index = layername.find('modelL')
    if index == -1:
        print('laod models error!!')
    else:
        name = layername[index:index + 11]
        modelweight = modelspath + name + '-weight.npy'
        modelbias = modelspath + name + '-bias.npy'
        weightnp = np.load(modelweight)
        # weightnp = np.transpose(weightnp, [2, 3, 1, 0])
        # print(weightnp)
        biasnp = np.load(modelbias)

        # init_weight = lambda shape, dtype: weightnp
        # init_bias   = lambda shape, dtype: biasnp
        # print('Done!')

        return torch.from_numpy(weightnp), torch.from_numpy(biasnp)
        # return init_weight, init_bias

class MEBasic(nn.Module):
    '''
    Get flow
    '''
    def __init__(self, layername):
        super(MEBasic, self).__init__()
        self.conv1 = nn.Conv2d(8, 32, 7, 1, padding=3)
        self.conv1.weight.data, self.conv1.bias.data = loadweightformnp(layername + '_F-1')
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(32, 64, 7, 1, padding=3)
        self.conv2.weight.data, self.conv2.bias.data = loadweightformnp(layername + '_F-2')
        self.relu2 = nn.ReLU()
        self.conv3 = nn.Conv2d(64, 32, 7, 1, padding=3)
        self.conv3.weight.data, self.conv3.bias.data = loadweightformnp(layername + '_F-3')
        self.relu3 = nn.ReLU()
        self.conv4 = nn.Conv2d(32, 16, 7, 1, padding=3)
        self.conv4.weight.data, self.conv4.bias.data = loadweightformnp(layername + '_F-4')
        self.relu4 = nn.ReLU()
        self.conv5 = nn.Conv2d(16, 2, 7, 1, padding=3)
        self.conv5.weight.data, self.conv5.bias.data = loadweightformnp(layername + '_F-5')

    def forward(self, x):
        x = self.relu1(self.conv1(x))
        x = self.relu2(self.conv2(x))
        x = self.relu3(self.conv3(x))
        x = self.relu4(self.conv4(x))
        x = self.conv5(x)
        return x



def bilinearupsacling(inputfeature):
    inputheight = inputfeature.size()[2]
    inputwidth = inputfeature.size()[3]
    # print(inputfeature.size())
    outfeature = F.interpolate(inputfeature, (inputheight * 2, inputwidth * 2), mode='bilinear')
    # print(outfeature.size())
    return outfeature
def bilinearupsacling2(inputfeature):
    inputheight = inputfeature.size()[2]
    inputwidth = inputfeature.size()[3]
    outfeature = F.interpolate(inputfeature, (inputheight * 2, inputwidth * 2), mode='bilinear', align_corners=True)
    return outfeature


# def bilinearupsacling2(inputfeature):
#     inputheight = inputfeature.get_shape().as_list()[1]
#     inputwidth = inputfeature.get_shape().as_list()[2]
#     outfeature = tf.image.resize_images(inputfeature, [inputheight * 2, inputwidth * 2], align_corners=True)
#     return outfeature


# def conv2Block(layername,
#                inputfeature,
#                kernelSize,
#                kernelNum,
#                paddingType='SAME',
#                stride=[1, 1, 1, 1],
#                actFunc=tf.nn.relu,
#                init=None):

#     # get the inputchannel
#     inputfeatureChannle = inputfeature.get_shape().as_list()[-1]

#     if init == None:
#         weights = tf.get_variable(
#             name=layername + '_weights',
#             shape=[kernelSize[0], kernelSize[1], inputfeatureChannle, kernelNum],
#             dtype=tf.float32,
#             initializer=tf.contrib.layers.xavier_initializer())
#         biases = tf.get_variable(
#             name=layername + '_biases', shape=[kernelNum], dtype=tf.float32, initializer=tf.constant_initializer(0.0))
#     else:
#         init_w, init_b = loadweightformnp(layername)
#         weights = tf.get_variable(name=layername + '_weights', dtype=tf.float32, initializer=init_w)
#         biases = tf.get_variable(name=layername + '_biases', dtype=tf.float32, trainable=True, initializer=init_b)

#     PreActivation = tf.nn.conv2d(inputfeature, weights, stride, padding=paddingType, name=layername + '_conv')
#     PreActivation = tf.nn.bias_add(PreActivation, biases, name=layername + '_addBiases')

#     if not actFunc == None:
#         activation = actFunc(PreActivation)
#         return activation
#     else:
#         return PreActivation

class ResBlock(nn.Module):
    def __init__(self, inputchannel, outputchannel, kernel_size, stride=1):
        super(ResBlock, self).__init__()
        self.relu1 = nn.ReLU()
        self.conv1 = nn.Conv2d(inputchannel, outputchannel, kernel_size, stride, padding=kernel_size//2)
        torch.nn.init.xavier_uniform_(self.conv1.weight.data)
        torch.nn.init.constant_(self.conv1.bias.data, 0.0)
        self.relu2 = nn.ReLU()
        self.conv2 = nn.Conv2d(outputchannel, outputchannel, kernel_size, stride, padding=kernel_size//2)
        torch.nn.init.xavier_uniform_(self.conv2.weight.data)
        torch.nn.init.constant_(self.conv2.bias.data, 0.0)
        # self.resblock = nn.Sequential(
        #     nn.ReLU(),
        #     nn.Conv2d(inputchannel, outputchannel, kernel_size, stride, padding=kernel_size//2),
        #     nn.ReLU(),
        #     nn.Conv2d(outputchannel, outputchannel, kernel_size, stride, padding=kernel_size//2),
        # )
        if inputchannel != outputchannel:
            self.adapt_conv = nn.Conv2d(inputchannel, outputchannel, 1)
            torch.nn.init.xavier_uniform_(self.adapt_conv.weight.data)
            torch.nn.init.constant_(self.adapt_conv.bias.data, 0.0)
        else:
            self.adapt_conv = None

    def forward(self, x):
        x_1 = self.relu1(x)
        firstlayer = self.conv1(x_1)
        firstlayer = self.relu2(firstlayer)
        seclayer = self.conv2(firstlayer)
        if self.adapt_conv is None:
            return x + seclayer
        else:
            return self.adapt_conv(x) + seclayer

class Warp_net(nn.Module):
    def __init__(self):
        super(Warp_net, self).__init__()
        channelnum = 64

        self.feature_ext = nn.Conv2d(6, channelnum, 3, padding=1)# feature_ext
        self.f_relu = nn.ReLU()
        torch.nn.init.xavier_uniform_(self.feature_ext.weight.data)
        torch.nn.init.constant_(self.feature_ext.bias.data, 0.0)
        self.conv0 = ResBlock(channelnum, channelnum, 3)#c0
        self.conv0_p = nn.AvgPool2d(2, 2)# c0p
        self.conv1 = ResBlock(channelnum, channelnum, 3)#c1
        self.conv1_p = nn.AvgPool2d(2, 2)# c1p
        self.conv2 = ResBlock(channelnum, channelnum, 3)# c2
        self.conv3 = ResBlock(channelnum, channelnum, 3)# c3
        self.conv4 = ResBlock(channelnum, channelnum, 3)# c4
        self.conv5 = ResBlock(channelnum, channelnum, 3)# c5
        self.conv6 = nn.Conv2d(channelnum, 3, 3, padding=1)
        torch.nn.init.xavier_uniform_(self.conv6.weight.data)
        torch.nn.init.constant_(self.conv6.bias.data, 0.0)

    def forward(self, x):
        feature_ext = self.f_relu(self.feature_ext(x))
        c0 = self.conv0(feature_ext)
        c0_p = self.conv0_p(c0)
        c1 = self.conv1(c0_p)
        c1_p = self.conv1_p(c1)
        c2 = self.conv2(c1_p)
        c3 = self.conv3(c2)
        c3_u = c1 + bilinearupsacling2(c3)#torch.nn.functional.interpolate(input=c3, scale_factor=2, mode='bilinear', align_corners=True)
        c4 = self.conv4(c3_u)
        c4_u = c0 + bilinearupsacling2(c4)# torch.nn.functional.interpolate(input=c4, scale_factor=2, mode='bilinear', align_corners=True)
        c5 = self.conv5(c4_u)
        res = self.conv6(c5)
        return res


# def Preprocessing(im):
#     r, g, b = tf.split(im, 3, axis=3)
#     b = b - 0.406
#     g = g - 0.456
#     r = r - 0.485
#     b_v = b / 0.225
#     g_v = g / 0.224
#     r_v = r / 0.229

#     rgb = tf.concat(values=[r_v, g_v, b_v], axis=3)
#     return rgb

flowfiledsSamples = [{} for i in range(8)]
class ME_Spynet(nn.Module):
    '''
    Get flow
    '''
    def __init__(self, layername='motion_estimation'):
        super(ME_Spynet, self).__init__()
        self.L = 4
        self.moduleBasic = torch.nn.ModuleList([ MEBasic(layername + 'modelL' + str(intLevel + 1)) for intLevel in range(4) ])
        # self.meBasic1 = MEBasic(layername + 'modelL1')
        # self.meBasic2 = MEBasic(layername + 'modelL2')
        # self.meBasic3 = MEBasic(layername + 'modelL3')
        # self.meBasic4 = MEBasic(layername + 'modelL4')
        # self.flow_warp = Resample2d()
        
        # self.meBasic = [self.meBasic1, self.meBasic2, self.meBasic3, self.meBasic4]

    # def Preprocessing(self, im):
    #     im[:, 0, :, :] -= 0.406
    #     im[:, 1, :, :] -= 0.456
    #     im[:, 2, :, :] -= 0.485
    #     im[:, 0, :, :] /= 0.225
    #     im[:, 1, :, :] /= 0.224
    #     im[:, 2, :, :] /= 0.229
    #     return im

    def forward(self, im1, im2):
        batchsize = im1.size()[0]
        im1_pre = im1
        im2_pre = im2

        im1list = [im1_pre]
        im2list = [im2_pre]
        for intLevel in range(self.L - 1):
            im1list.append(F.avg_pool2d(im1list[intLevel], kernel_size=2, stride=2))# , count_include_pad=False))
            im2list.append(F.avg_pool2d(im2list[intLevel], kernel_size=2, stride=2))#, count_include_pad=False))

        shape_fine = im2list[self.L - 1].size()
        zeroshape = [batchsize, 2, shape_fine[2] // 2, shape_fine[3] // 2]
        device_id = im1.device.index
        flowfileds = torch.zeros(zeroshape, dtype=torch.float32, device=device_id)
        for intLevel in range(self.L):
            flowfiledsUpsample = bilinearupsacling(flowfileds) * 2.0
            flowfileds = flowfiledsUpsample + self.moduleBasic[intLevel](torch.cat([im1list[self.L - 1 - intLevel], flow_warp(im2list[self.L - 1 - intLevel], flowfiledsUpsample), flowfiledsUpsample], 1))# residualflow

        return flowfileds


# def ME_Ours(im1, im2):
#     h_d = im1.get_shape().as_list()[1]
#     w_d = im1.get_shape().as_list()[2]

#     im1_d = tf.contrib.layers.avg_pool2d(im1, kernel_size=[4, 4], stride=[4, 4])
#     im2_d = tf.contrib.layers.avg_pool2d(im2, kernel_size=[4, 4], stride=[4, 4])
#     flow_d, _ = ME_Spynet(im1_d, im2_d)
#     # feature extract...
#     with tf.variable_scope('FeatureExtract', reuse=tf.AUTO_REUSE) as scope:
#         c11 = conv2Block(layername='conv1a', inputfeature=im1, kernelNum=16, kernelSize=[3, 3], stride=[1, 1, 1, 1])
#         c21 = conv2Block(layername='conv1a', inputfeature=im2, kernelNum=16, kernelSize=[3, 3], stride=[1, 1, 1, 1])

#     with tf.variable_scope('GenerateFlow') as scope:
#         flow_up = tf.image.resize_images(flow_d, [h_d, w_d], align_corners=True)
#         c21_w = flow_warp(c21, flow_up)

#         c11 = conv2Block(
#             layername='conv1a',
#             inputfeature=tf.concat(values=[c11, c21_w], axis=3),
#             kernelNum=16,
#             kernelSize=[3, 3],
#             stride=[1, 1, 1, 1],
#             actFunc=leaky_relu)

#         flow_res = conv2Block(
#             layername='conv2a', inputfeature=c11, kernelNum=2, kernelSize=[3, 3], stride=[1, 1, 1, 1], actFunc=None)
#         flow = flow_res + flow_up

#         warpframe = flow_warp(im2, flow)

#     return flow, warpframe


def build_model():

    net = ME_Spynet().cuda()
    # inputImage1 = tf.placeholder(shape=[1,3, 448, 832], dtype=tf.float32, name="inputimage1")
    # inputImage2 = tf.placeholder(shape=[1,3, 448, 832], dtype=tf.float32, name="inputimage2")

    # # inputImage, inputRes, groudtruth
    # flowfiled, warpframe = ME_Spynet(inputImage1, inputImage2)

    # read images
    im1 = imageio.imread('input.png')
    im1 = im1 / 255.0
    im1 = np.expand_dims(im1, axis=0)
    im2 = imageio.imread('ref.png')
    im2 = im2 / 255.0
    im2 = np.expand_dims(im2, axis=0)
    
    # means = np.array([0.485, 0.456, 0.406])
    # stds = np.array([0.229, 0.224, 0.225])
    # input_image = (input_image - means) / stds
    # ref_image = (ref_image - means) / stds
    im1 = np.transpose(im1, [0, 3, 1, 2])
    im2 = np.transpose(im2, [0, 3, 1, 2])
    im1 = torch.from_numpy(im1).float().cuda()
    im2 = torch.from_numpy(im2).float().cuda()
    net.eval()
    flow, warp_frame = net(im1, im2)
    flow = flow.detach().cpu().numpy()
    warp_frame = warp_frame.cpu().detach().numpy()
    rgb_my = flow_to_image(flow[0, :, :, :])
    imageio.imwrite('flow.png', rgb_my)

    # # test spynet
    # flowspy = read_flow('result_hevc.flo')
    # flowspy_rgb = flow_to_image(flowspy)
    # imageio.imwrite('test2.png', flowspy_rgb)
    # psnr = CalcuPSNR(rgb_my / 255, flowspy_rgb / 255)
    # # print(rgb_my)
    # error = evaluate_flow(flow[0, :, :, :], flowspy)
    # # print(flow[0,:,:,:])
    # # print(flowspy)
    # print(error)

    # print(psnr)
    imageio.imwrite('warp2.png', warp_frame[0, :, :, :].transpose(1,2,0))
    # warp_frame_rgb = flow_warp(im2, flow)
    psnrwapr = CalcuPSNR(im1[0].cpu().numpy().transpose(1,2,0), warp_frame[0, :, :, :].transpose(1,2,0))
    print(psnrwapr)



if __name__ == '__main__':
    build_model()
