import torch.nn as nn
from mmcv.cnn import ConvModule, bias_init_with_prob, normal_init

from ..builder import HEADS
from .anchor_head import AnchorHead

from mmdet.models.backbones.effnet.utils import MemoryEfficientSwish, Swish, SeparableConvBlock

@HEADS.register_module()
class RetinaEffHead(AnchorHead):
    """"RetinaHead with efficientnet.
    add by chen
    """

    def __init__(self,
                 num_classes,
                 num_ins,
                 in_channels,
                 stacked_convs=4,
                 conv_cfg=None,
                 norm_cfg=None,
                 **kwargs):
        self.stacked_convs = stacked_convs
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.num_ins = num_ins
        super(RetinaEffHead, self).__init__(num_classes, in_channels,
                                              **kwargs)

    def _init_layers(self):
        # self.relu = nn.ReLU(inplace=True)
        # self.cls_convs = nn.ModuleList()
        # self.reg_convs = nn.ModuleList()
        # for i in range(self.num_ins):
        #     cls_convs = nn.ModuleList()
        #     reg_convs = nn.ModuleList()
        #     for i in range(self.stacked_convs):
        #         chn = self.in_channels if i == 0 else self.feat_channels
        #         cls_convs.append(
        #             ConvModule(
        #                 chn,
        #                 self.feat_channels,
        #                 3,
        #                 stride=1,
        #                 padding=1,
        #                 conv_cfg=self.conv_cfg,
        #                 norm_cfg=self.norm_cfg))
        #         reg_convs.append(
        #             ConvModule(
        #                 chn,
        #                 self.feat_channels,
        #                 3,
        #                 stride=1,
        #                 padding=1,
        #                 conv_cfg=self.conv_cfg,
        #                 norm_cfg=self.norm_cfg))
        #     self.cls_convs.append(cls_convs)
        #     self.reg_convs.append(reg_convs)
        # for i in range(self.stacked_convs):
        #     for j in range(1, self.num_ins):
        #         self.cls_convs[j][i].conv = self.cls_convs[0][i].conv
        #         self.reg_convs[j][i].conv = self.reg_convs[0][i].conv
        # self.retina_cls = nn.Conv2d(
        #     self.feat_channels,
        #     self.num_anchors * self.cls_out_channels,
        #     3,
        #     padding=1)
        # self.retina_reg = nn.Conv2d(
        #     self.feat_channels, self.num_anchors * 4, 3, padding=1)


        self.cls_convs = nn.ModuleList()
        self.reg_convs = nn.ModuleList()
        for i in range(self.num_ins):
            cls_convs = nn.ModuleList()
            reg_convs = nn.ModuleList()
            for i in range(self.stacked_convs):
                cls_convs.append(SeparableConvBlock(self.in_channels, self.in_channels, norm= True, activation=True))
                reg_convs.append(SeparableConvBlock(self.in_channels, self.in_channels, norm= True, activation=True))
            self.cls_convs.append(cls_convs)
            self.reg_convs.append(reg_convs)
        
        self.retina_cls = SeparableConvBlock(self.in_channels, self.num_anchors * self.cls_out_channels, norm=False, activation=False)
        self.retina_reg = SeparableConvBlock(self.in_channels, self.num_anchors * 4, norm=False, activation=False)

    def init_weights(self):
        # for m in self.cls_convs[0]:
        #     normal_init(m.conv, std=0.01)
        # for m in self.reg_convs[0]:
        #     normal_init(m.conv, std=0.01)
        # bias_cls = bias_init_with_prob(0.01)
        # normal_init(self.retina_cls, std=0.01, bias=bias_cls)
        # normal_init(self.retina_reg, std=0.01)

        import math
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def forward(self, feats):
        cls_scores = []
        bbox_preds = []
        for i, x in enumerate(feats):
            cls_feat = feats[i]
            reg_feat = feats[i]
            for cls_conv in self.cls_convs[i]:
                cls_feat = cls_conv(cls_feat)
            for reg_conv in self.reg_convs[i]:
                reg_feat = reg_conv(reg_feat)
            cls_score = self.retina_cls(cls_feat)
            bbox_pred = self.retina_reg(reg_feat)
            cls_scores.append(cls_score)
            bbox_preds.append(bbox_pred)
        return cls_scores, bbox_preds
