# 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)
    t = (gamma / std).reshape(-1, 1, 1, 1)
    module_output.weight[:] = (kernel * t).reshape(module_output.weight.shape)
    module_output.bias[:] = beta + ((conv.bias if conv.bias is not None else 0) - running_mean) * gamma / std
    return module_output


class ConvBn2d(nn.ConvBn2d):
    def __init__(self, *args, **kwargs):
        bias = kwargs.pop("bias", False) and False
        super().__init__(*args, bias=bias, **kwargs)

    @classmethod
    def fuse_conv_bn(cls, module: nn.Module):
        module_output = module
        if isinstance(module, ConvBn2d):
            return _fuse_conv_bn(module.conv, module.bn)
        for name, child in module.named_children():
            setattr(module_output, name, cls.fuse_conv_bn(child))
        del module
        return module_output


class LargeKernelReparam(nn.Module):
    def __init__(self, channels, kernel, small_kernels=()):
        super(LargeKernelReparam, self).__init__()
        self.dw_large = ConvBn2d(channels, channels, kernel, padding=kernel // 2, groups=channels)

        self.small_kernels = small_kernels
        for k in self.small_kernels:
            setattr(self, f"dw_small_{k}", ConvBn2d(channels, channels, k, padding=k // 2, groups=channels))

    def forward(self, inp):
        outp = self.dw_large(inp)
        for k in self.small_kernels:
            outp += getattr(self, f"dw_small_{k}")(inp)
        return outp

    @classmethod
    def convert_to_deploy(cls, module: nn.Module):
        module_output = module
        if isinstance(module, LargeKernelReparam):
            module = ConvBn2d.fuse_conv_bn(module)
            module_output = copy.deepcopy(module.dw_large)
            kernel = module_output.kernel_size[0]
            for k in module.small_kernels:
                dw_small = getattr(module, f"dw_small_{k}")
                module_output.weight += F.pad(dw_small.weight, [[0, 0]] * 3 + [[(kernel - k) // 2] * 2] * 2)
                module_output.bias += dw_small.bias
            return module_output
        for name, child in module.named_children():
            setattr(module_output, name, cls.convert_to_deploy(child))
        del module
        return module_output


class Mlp(nn.Module):
    def __init__(self, in_channels, hidden_channels=None, out_channels=None, act_layer=nn.GELU, drop=0.,):
        super().__init__()
        out_features = out_channels or in_channels
        hidden_features = hidden_channels or in_channels
        self.fc1 = ConvBn2d(in_channels, hidden_features, 1, stride=1, padding=0)
        self.act = act_layer()
        self.fc2 = ConvBn2d(hidden_features, out_features, 1, stride=1, padding=0)
        self.drop = 
nn.Dropout(drop)