# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257 
import re
import sys

import numpy as np
import torch
import torch.nn as nn

from .layers import (IdentityLayer, MBInvertedConvLayer,
                     MobileInvertedResidualBlock, ZeroLayer)
from .mix_ops import MixedEdge, build_candidate_ops, conv_func_by_name


class NasRecBackbone(nn.Module):

    def __init__(self, first_conv, blocks):
        super(NasRecBackbone, self).__init__()
        self.first_conv = first_conv
        self.blocks = nn.ModuleList(blocks)
        self.output_idx = [5, 11, 17, 23]

    def forward(self, x):
        x = self.first_conv(x)
        idx = 0
        out = []
        for block in self.blocks:
            x = block(x)
            if (idx + 1) % (int(len(self.blocks) / 4)) == 0:
                out.append(x)
            idx += 1
        return out[0], out[1], out[2], out[3]

    def get_bn_param(self):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
                return {
                    'momentum': m.momentum,
                    'eps': m.eps,
                }
        return None

    @property
    def config(self):
        return {
            'name': NasRecBackbone.__name__,
            'bn': self.get_bn_param(),
            'first_conv': 'conv_in3_out32_k3_s2_p1',
            'blocks': [block.config for block in self.blocks]
        }

    def set_bn_param(self, momentum, eps):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
                m.momentum = momentum
                m.eps = eps
        return

    @staticmethod
    def build_from_config(config):
        first_conv_config = config['first_conv']
        match_obj = re.match(r'conv_in(\d+)_out(\d+)_k(\d+)_s(\d+)_p(\d+)',
                             first_conv_config)
        in_channel = int(match_obj.group(1))
        out_channel = int(match_obj.group(2))
        kernel_size = int(match_obj.group(3))
        stride = int(match_obj.group(4))
        padding = int(match_obj.group(5))
        first_conv = nn.Sequential(
            nn.Conv2d(
                in_channel,
                out_channel,
                kernel_size,
                stride,
                padding,
                bias=False),
            nn.BatchNorm2d(out_channel),
            nn.ReLU(inplace=True),
        )
        blocks = []
        for block_config in config['blocks']:
            blocks.append(
                MobileInvertedResidualBlock.build_from_config(block_config))
        net = NasRecBackbone(first_conv, blocks)
        if 'bn' in config:
            net.set_bn_param(**config['bn'])
        else:
            net.set_bn_param(momentum=0.1, eps=1e-3)
        return net


class CompactDetBackbone(NasRecBackbone):
    '''
    proxyless nas backbone, 5M.
    '''

    def __init__(self,
                 width_stages,
                 input_channel=None,
                 bn_param=(0.1, 1e-3),
                 **kwargs):

        if input_channel is None:
            input_channel = width_stages[0]
        first_conv = nn.Sequential(
            nn.Conv2d(
                3,
                input_channel,
                kernel_size=(3, 3),
                stride=2,
                padding=1,
                bias=False), nn.BatchNorm2d(input_channel),
            nn.ReLU(inplace=True))

        conv_candidates = [
            '5x5_MBConv2', '5x5_MBConv4', '3x3_MBConv2', '3x3_MBConv4',
            '13_MixConv2', '13_MixConv4', '35_MixConv2', '35_MixConv4',
            '135_MixConv2', '135_MixConv4', '13_LinMixConv', '35_LinMixConv',
            '135_LinMixConv', '13_RepConv2', '13_RepConv4', '35_RepConv2',
            '35_RepConv4', '135_RepConv2', '135_RepConv4', 'Zero'
        ]

        se_candidates = ['SE_2', 'SE_4', 'SE_8', 'Zero']

        conv_op_ids = [
            15, 17, 17, 17, 17, 0, 16, 16, 18, 18, 16, 2, 16, 18, 16, 18, 18,
            2, 1, 18, 18, 18, 16, 2
        ]

        n_cell_stages = [5, 5, 5, 5]

        stride_stages = [(2, 2), (2, 2), (2, 2), (2, 2)]

        if se_candidates:
            assert len(conv_op_ids) == sum(n_cell_stages) + 4
        else:
            assert len(conv_op_ids) == sum(n_cell_stages)

        blocks = []

        for width, n_cell, s in zip(width_stages, n_cell_stages,
                                    stride_stages):
            for i in range(n_cell):
                if i == 0:
                    stride = s
                else:
                    stride = (1, 1)
                block_i = len(blocks)
                conv_op = conv_func_by_name(
                    conv_candidates[conv_op_ids[block_i]])(input_channel,
                                                           width, stride)

                if stride == (1, 1) and input_channel == width:
                    shortcut = IdentityLayer()
                else:
                    shortcut = None
                inverted_residual_block = MobileInvertedResidualBlock(
                    conv_op, shortcut)
                blocks.append(inverted_residual_block)
                input_channel = width

            if se_candidates is not None:
                block_i = len(blocks)
                se_op = conv_func_by_name(se_candidates[conv_op_ids[block_i]])(
                    input_channel, width, stride)
                shortcut = IdentityLayer()
                inverted_residual_block = MobileInvertedResidualBlock(
                    se_op, shortcut)
                blocks.append(inverted_residual_block)

        self.out_channel = input_channel

        super(CompactDetBackbone, self).__init__(first_conv, blocks)
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(
                    m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)