# Copyright (c) 2014-2022 Megvii Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import copy

import megengine
import megengine.functional as F
import megengine.module as nn
import numpy as np
from basecls.layers import DropPath, init_weights
from basecls.utils import registers


def _fuse_prebn_conv1x1(bn, conv):
    module_output = copy.deepcopy(conv)
    module_output.bias = megengine.Parameter(np.zeros(module_output._infer_bias_shape(), dtype=np.float32))
    assert conv.groups == 1
    kernel = conv.weight
    running_mean = bn.running_mean
    running_var = bn.running_var
    gamma = bn.weight
    beta = bn.bias
    eps = bn.eps
    std = F.sqrt(running_var + eps)
    t = (gamma / std).reshape(1, -1, 1, 1)
    module_output.weight[:] = kernel * t
    module_output.bias[:] = F.conv2d(beta - running_mean * gamma / std, kernel, conv.bias)
    return module_output


def _fuse_conv_bn(conv, bn):
    module_output = copy.deepcopy(conv)
    module_output.bias = megengine.Parameter(np.zeros(module_output._infer_bias_shape(), dtype=np.float32))
    # flatten then reshape in case of group conv
    kernel = F.flatten(conv.weight, end_axis=conv.weight.ndim - 4)
    running_mean = bn.running_mean
    running_var = bn.running_var
    gamma = bn.weight
    beta = bn.bias
    eps = bn.eps
    std = 
F.sqrt(running_var + eps)