#!/usr/bin/env python3
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
"""VGG Series

VGG: `"Very Deep Convolutional Networks for Large-Scale Image Recognition"
<https://arxiv.org/abs/1409.1556>`_
"""
from typing import Any, Mapping, Sequence

import megengine as mge
import megengine.hub as hub
import megengine.module as M

from basecls.layers import activation, build_head, conv2d, init_weights, norm2d
from basecls.utils import recursive_update, registers

__all__ = ["VGGStage", "VGG"]


class VGGStage(M.Module):
    """VGG stage (sequence of blocks w/ the same output shape)."""

    def __init__(self, w_in: int, w_out: int, depth: int, norm_name: str, act_name: str):
        super().__init__()
        self.depth = depth
        for i in range(depth):
            block = M.Sequential(
                conv2d(w_in, w_out, 3), norm2d(norm_name, w_out), activation(act_name)
            )
            setattr(self, f"b{i + 1}", block)
            w_in = w_out
        self.max_pool = 
M.MaxPool2d(kernel_size=2, stride=2)