# -*- coding: utf-8 -*-
# MIT License
#
# Copyright (c) 2020 <NAME>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, 
# publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, 
# subject to the following conditions:
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#
# This repo is licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import math
import collections

import megengine.module as nn
import megengine.functional as F

from model.nn_upsample import NeuralUpsampler, FlowMaskEstimator

from common.utils import flow_warp, upsample2d_flow_as


def conv(inp, out, k=3, s=1, d=1, isReLU=True):
    if isReLU:
        ret = nn.Sequential(nn.Conv2d(inp, out, k, s, padding=((k - 1) * d) // 2, dilation=d, bias=True), nn.LeakyReLU(0.1))
    else:
        ret = nn.Sequential(nn.Conv2d(inp, out, k, s, padding=((k - 1) * d) // 2, dilation=d, bias=True))
    return ret


class ContextNetwork(nn.Module):
    def __init__(self, ch_in):
        super(ContextNetwork, self).__init__()
        self.convs = nn.Sequential(conv(ch_in, 128, 3, 1, 1), conv(128, 128, 3, 1, 2), conv(128, 128, 3, 1, 4), conv(128, 96, 3, 1, 8),
                                   conv(96, 64, 3, 1, 16), conv(64, 32, 3, 1, 1), conv(32, 2, isReLU=False))

    def forward(self, x):
        return self.convs(x)


class FlowEstimator(nn.Module):
    def __init__(self, ch_in):
        super(FlowEstimator, self).__init__()
        self.conv1 = conv(ch_in, 128)
        self.conv2 = conv(128, 128)
        self.conv3 = conv(128 + 128, 96)
        self.conv4 = conv(96 + 128, 64)
        self.conv5 = conv(96 + 64, 32)
        # channels of the second last layer
        self.feat_dim = 32
        self.predict_flow = conv(64 + 32, 2, isReLU=False)

    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        x3 = self.conv3(F.concat([x1, x2], axis=1))
        x4 = self.conv4(F.concat([x2, x3], axis=1))
        x5 = self.conv5(F.concat([x3, x4], axis=1))
        flow = self.predict_flow(F.concat([x4, x5], axis=1))
        return x5, flow


class CostVolume(nn.Module):
    def __init__(self, d=4, *args, **kwargs):
        super(CostVolume, self).__init__()
        self.d = d
        self.out_dim = 2 * self.d + 1
        self.pad_size = self.d

    def forward(self, x1, x2):
        _, _, H, W = x1.shape

        x2 = F.nn.pad(x2, ((0, 0), (0, 0), (self.pad_size, self.pad_size), (self.pad_size, self.pad_size)))
        cv = []
        for i in range(self.out_dim):
            for j in range(self.out_dim):
                cost = x1 * x2[:, :, i:(i + H), j:(j + W)]
                cost = 
F.mean(cost, 1, keepdims=True)