# -*- coding: utf-8 -*-
# MIT License
#
# Copyright (c) 2021 coolbeam
#
# 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:
#
# 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 megengine.module as nn
import megengine.functional as F

from common.utils import flow_warp, upsample2d_flow_as


def conv(in_planes, out_planes, kernel_size=3, stride=1, dilation=1, isReLU=True, if_IN=False, IN_affine=False, if_BN=False):
    if isReLU:
        if if_IN:
            return nn.Sequential(
                nn.Conv2d(in_planes,
                          out_planes,
                          kernel_size=kernel_size,
                          stride=stride,
                          dilation=dilation,
                          padding=((kernel_size - 1) * dilation) // 2,
                          bias=True), nn.LeakyReLU(0.1), nn.InstanceNorm(out_planes, affine=IN_affine))
        elif if_BN:
            return nn.Sequential(
                nn.Conv2d(in_planes,
                          out_planes,
                          kernel_size=kernel_size,
                          stride=stride,
                          dilation=dilation,
                          padding=((kernel_size - 1) * dilation) // 2,
                          bias=True), nn.LeakyReLU(0.1), nn.BatchNorm2d(out_planes, affine=IN_affine))
        else:
            return nn.Sequential(
                nn.Conv2d(in_planes,
                          out_planes,
                          kernel_size=kernel_size,
                          stride=stride,
                          dilation=dilation,
                          padding=((kernel_size - 1) * dilation) // 2,
                          bias=True), nn.LeakyReLU(0.1))
    else:
        if if_IN:
            return nn.Sequential(
                nn.Conv2d(in_planes,
                          out_planes,
                          kernel_size=kernel_size,
                          stride=stride,
                          dilation=dilation,
                          padding=((kernel_size - 1) * dilation) // 2,
                          bias=True), nn.InstanceNorm(out_planes, affine=IN_affine))
        elif if_BN:
            return nn.Sequential(
                nn.Conv2d(in_planes,
                          out_planes,
                          kernel_size=kernel_size,
                          stride=stride,
                          dilation=dilation,
                          padding=((kernel_size - 1) * dilation) // 2,
                          bias=True), 
nn.BatchNorm2d(out_planes, affine=IN_affine)