import torch
import torch.nn as nn
from ..builder import BACKBONES



@BACKBONES.register_module()
class RoadsideBackBone(nn.Module):
    '''a detector backbone for roadside
    '''
    def __init__(self, pretrained=''):
        super().__init__()
        #common params
        cm_f = 3 #common filter size
        cm_s = 1 #common stride
        cm_p = 1 #common padding
        cm_g = 1 #common group

        self.conv0_1 = self.create_conv(3, 16,  cm_f, cm_s, cm_p, cm_g)
        self.conv0_2 = self.create_conv(16, 32, cm_f, 2, cm_p, cm_g) # stride=2
        
        self.conv1_1 = self.create_conv(32, 32, cm_f, cm_s, cm_p, cm_g)
        self.conv1_2 = self.create_conv(32, 32, cm_f, cm_s, cm_p, cm_g)
        self.conv1_3 = self.create_conv(32, 64, cm_f, 2, cm_p, cm_g) # stride=4

        self.conv2_1 = self.create_conv(64, 64, cm_f, cm_s, cm_p, cm_g)
        self.conv2_2 = self.create_conv(64, 64, cm_f, cm_s, cm_p, cm_g)
        self.conv2_3 = self.create_conv(64, 80, cm_f, 2, cm_p, cm_g) # stride=8
        
        self.conv2_4 = self.create_conv(80, 80, cm_f, cm_s, cm_p, cm_g)
        self.conv2_5 = self.create_conv(80, 80, cm_f, cm_s, cm_p, cm_g)
        self.conv_stage1_out = self.create_conv(80, 80, cm_f, cm_s, cm_p, cm_g) # stage 1

        self.conv3_1 = self.create_conv(80, 96, cm_f, 2, cm_p, cm_g) # stride=16
        self.conv3_2 = self.create_conv(96, 96, cm_f, cm_s, cm_p, cm_g)
        self.conv_stage2_out = self.create_conv(96, 96, cm_f, cm_s, cm_p, cm_g) # stage 2

        self.conv4_1 = self.create_conv(96, 112, cm_f, 2, cm_p, cm_g) # stride=32
        self.conv4_2 = self.create_conv(112, 112, cm_f, cm_s, cm_p, cm_g)
        self.conv_stage3_out = self.create_conv(112, 112, cm_f, cm_s, cm_p, cm_g) # stage 3

        self.conv5_1 = self.create_conv(112, 128, cm_f, 2, cm_p, cm_g) # stride=64
        self.conv5_2 = self.create_conv(128, 128, cm_f, cm_s, cm_p, cm_g)
        self.conv_stage4_out = self.create_conv(128, 128, cm_f, cm_s, cm_p, cm_g) # stage 4

    def create_conv(self, in_cn, out_cn, f_sz, stride, padding, group):
        '''
        create a conv layer with bn&relu
        '''
        conv = nn.Sequential(
            nn.Conv2d(in_cn, out_cn, kernel_size=f_sz, stride=stride, padding=padding, groups=group, bias=False),
            nn.BatchNorm2d(out_cn),
            nn.ReLU()
        )
        #init with msra
        nn.init.kaiming_normal_(conv[0].weight, nonlinearity='relu')

        return conv

    def forward(self, x):
        out = self.conv0_1(x)
        out = self.conv0_2(out)

        out = self.conv1_1(out)
        out = self.conv1_2(out)
        out = self.conv1_3(out)

        out = self.conv2_1(out)
        out = self.conv2_2(out)
        out = self.conv2_3(out)
        out = self.conv2_4(out)
        out = self.conv2_5(out)
        stage1_fm = self.conv_stage1_out(out)

        out = self.conv3_1(out)
        out = self.conv3_2(out)
        stage2_fm = self.conv_stage2_out(out)

        out = self.conv4_1(out)
        out = self.conv4_2(out)
        stage3_fm = self.conv_stage3_out(out)

        out = self.conv5_1(out)
        out = self.conv5_2(out)
        stage4_fm = self.conv_stage4_out(out)

        return (stage1_fm, stage2_fm, stage3_fm, stage4_fm)

    def init_weights(self, pretrained=None):
        #inited in create_conv()
        pass