#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.

import megengine as mge
import megengine.functional as F
import megengine.module as M

__all__ = ["fuse_conv_and_bn", "fuse_model", "replace_module"]


def fuse_conv_and_bn(conv, bn):
    # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
    fusedconv = M.Conv2d(
        conv.in_channels,
        conv.out_channels,
        kernel_size=conv.kernel_size,
        stride=conv.stride,
        padding=conv.padding,
        groups=conv.groups,
        bias=True,
    )

    # fused_conv.weight = bn.weight / running_var * conv.weight
    w_conv = conv.weight.reshape(conv.out_channels, -1)
    factor = (bn.weight / F.sqrt(bn.eps + bn.running_var)).reshape(-1)
    # diag_factor = diag(factor)
    fusedconv.weight = mge.Parameter(
        (factor.reshape(-1, 1) * w_conv).reshape(fusedconv.weight.shape)
    )

    # fused_conv.bias = bn.bias + (conv.bias - running_mean) * bn.weight / runing_var
    conv_bias = 
F.zeros(bn.running_mean.shape)