# Copyright (c) Alibaba, Inc. and its affiliates.
# The ZenNAS implementation is also open-sourced by the authors, and available at https://github.com/idstcv/ZenNAS.

from torch import nn

from . import (basic_blocks, super_blocks, super_res_idwexkx, super_res_k1kxk1,
               super_res_kxkx)
from .global_utils import create_netblock_list_from_str_inner


class PlainNet(nn.Module):

    def __init__(self,
                 argv=None,
                 opt=None,
                 num_classes=None,
                 plainnet_struct=None,
                 no_create=False,
                 **kwargs):
        super(PlainNet, self).__init__()
        self.argv = argv
        self.opt = opt
        self.num_classes = num_classes
        self.plainnet_struct = plainnet_struct

        self.module_opt = None

        if self.num_classes is None:
            self.num_classes = self.module_opt.num_classes

        if self.plainnet_struct is None and self.module_opt.plainnet_struct is not None:
            self.plainnet_struct = self.module_opt.plainnet_struct

        if self.plainnet_struct is None:
            if hasattr(opt, 'plainnet_struct_txt'
                       ) and opt.plainnet_struct_txt is not None:
                plainnet_struct_txt = opt.plainnet_struct_txt
            else:
                plainnet_struct_txt = self.module_opt.plainnet_struct_txt

            if plainnet_struct_txt is not None:
                with open(plainnet_struct_txt, 'r', encoding='utf-8') as fid:
                    the_line = fid.readlines()[0].strip()
                    self.plainnet_struct = the_line
                pass

        if self.plainnet_struct is None:
            return

        the_s = self.plainnet_struct

        block_list, remaining_s = create_netblock_list_from_str_inner(
            the_s,
            netblocks_dict=_all_netblocks_dict_,
            no_create=no_create,
            **kwargs)
        assert len(remaining_s) == 0

        self.block_list = block_list
        if not no_create:
            self.module_list = nn.ModuleList(block_list)

    def forward(self, x):
        output = x
        for the_block in self.block_list:
            output = the_block(output)
        return output

    def __str__(self):
        s = ''
        for the_block in self.block_list:
            s += str(the_block)
        return s

    def __repr__(self):
        return str(self)


_all_netblocks_dict_ = {}
_all_netblocks_dict_ = basic_blocks.register_netblocks_dict(
    _all_netblocks_dict_)
_all_netblocks_dict_ = super_blocks.register_netblocks_dict(
    _all_netblocks_dict_)
_all_netblocks_dict_ = super_res_kxkx.register_netblocks_dict(
    _all_netblocks_dict_)
_all_netblocks_dict_ = super_res_k1kxk1.register_netblocks_dict(
    _all_netblocks_dict_)
_all_netblocks_dict_ = super_res_idwexkx.register_netblocks_dict(
    _all_netblocks_dict_)