import torch.nn as nn
import torch.nn.functional as F
from .repvgg import create_RepVGG

class RepVGG_8_1_align(nn.Module):
    """
    RepVGG backbone, output resolution are 1/8 and 1.
    Each block has 2 layers.
    """

    def __init__(self, config):
        super().__init__()
        backbone = create_RepVGG(False)

        self.layer0, self.layer1, self.layer2, self.layer3 = backbone.stage0, backbone.stage1, backbone.stage2, backbone.stage3

        for layer in [self.layer0, self.layer1, self.layer2, self.layer3]:
            for m in layer.modules():
                if isinstance(m, nn.Conv2d):
                    nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                    nn.init.constant_(m.weight, 1)
                    nn.init.constant_(m.bias, 0)

    def forward(self, x):
        out = self.layer0(x) # 1/2
        for module in self.layer1:
            out = module(out) # 1/2
        x1 = out
        for module in self.layer2:
            out = module(out) # 1/4
        x2 = out
        for module in self.layer3:
            out = module(out) # 1/8
        x3 = out
                
        return {'feats_c': x3, 'feats_f': None, 'feats_x2': x2, 'feats_x1': x1}