# ***************************************************************
# Copyright (c) 2021 Jittor. All Rights Reserved.
# Maintainers:
#     Guoye Yang <498731903@qq.com>
#     Dun Liang <randonlang@gmail.com>.
#
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
# ***************************************************************
# This model is generated by pytorch converter.
import jittor as jt
from jittor import nn
from jdet.utils.registry import BACKBONES

__all__ = [
    'SSD_VGG16', 'SSD_VGG19',
]

@BACKBONES.register_module()
class SSD_VGG(nn.Module):

    extra_setting = {
        300: (256, 'S', 512, 128, 'S', 256, 128, 256, 128, 256),
        512: (256, 'S', 512, 128, 'S', 256, 128, 'S', 256, 128, 'S', 256, 128),
    }

    def __init__(self,
                 features,
                 input_size,
                 out_feature_indices=(22, 34)):

        super(SSD_VGG, self).__init__()
        assert input_size in (300, 512)
        self.input_size = input_size

        self.features = features
        self.features += [nn.MaxPool2d(kernel_size=3, stride=1, padding=1)]
        self.features += [nn.Conv2d(512, 1024,
                                    kernel_size=3, padding=6, dilation=6)]
        self.features += [nn.ReLU()]
        self.features += [nn.Conv2d(1024, 1024, kernel_size=1)]
        self.features += [nn.ReLU()]
        self.features = nn.Sequential(self.features)
        self.out_feature_indices = out_feature_indices

    def execute(self, x):
        outs = []
        for i, layer in enumerate(self.features):
            x = layer(x)
            if i in self.out_feature_indices:
                outs.append(x)

        if len(outs) == 1:
            return outs[0]
        else:
            return tuple(outs)

def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3
    for v in cfg:
        if v == 'M':
            layers += [nn.Pool(kernel_size=2, stride=2,
                               op="maximum", ceil_mode=True)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm(v), nn.ReLU()]
            else:
                layers += [conv2d, nn.ReLU()]
            in_channels = v
    return layers


cfgs = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512],
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512],
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512],
}


def _vgg(arch, cfg, batch_norm, **kwargs):
    model = SSD_VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
    return model

@BACKBONES.register_module()
def SSD_VGG16(pretrained=None, **kwargs):
    model = _vgg('vgg16', 'D', False, **kwargs)
    if pretrained is not None:
        model.load(pretrained)
    return model


@BACKBONES.register_module()
def SSD_VGG19(pretrained=None, **kwargs):
    model = _vgg('vgg19', 'E', False, **kwargs)
    if pretrained is not None:
        model.load(pretrained)
    return model

if __name__ == '__main__':
    model = SSD_VGG16(input_size=300)
    i = jt.random((1, 3, 300, 300))
    model(i)
    print(model)
