import torch
import torch.nn as nn


def absorb_bn(module, min_var=None):
    last_conv_name = None
    last_conv = None
    for name, child in module.named_children():
        if isinstance(child, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
            last_conv_name = name
            last_conv = child
        elif isinstance(child, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d,)):
            if last_conv is None:  # only fuse BN that is after Conv
                continue
            new_conv = _absorb_bn(last_conv, child, min_var=min_var)
            module.add_module(last_conv_name, new_conv)
            module.add_module(name, nn.Identity())
            last_conv = None
        else:
            absorb_bn(child, min_var=min_var)
            last_conv = None
    return module


def _absorb_bn(conv, bn, min_var=None):
    conv_w = conv.weight
    conv_b = conv.bias if conv.bias is not None else torch.zeros_like(bn.running_mean)
    bn_shape = [conv.out_channels] + [1] * (len(conv_w.shape) - 1)
    if min_var is not None:
        zero_weight_ch = torch.where(bn.running_var < min_var)[0]
    
    factor = bn.weight / torch.sqrt(bn.running_var + bn.eps)
    conv.weight = nn.Parameter(conv_w * factor.reshape(bn_shape))
    conv.bias = nn.Parameter((conv_b - bn.running_mean) * factor + bn.bias)
    if min_var is not None:
        conv.weight.data[zero_weight_ch, :, :, :] *= 0
    return conv