##################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019 #
##################################################
from os import path as osp
from typing import List, Text
import torch

__all__ = [
    "change_key",
    "get_cell_based_tiny_net",
    "get_search_spaces",
    "get_cifar_models",
    "get_imagenet_models",
    "obtain_model",
    "obtain_search_model",
    "load_net_from_checkpoint",
    "CellStructure",
    "CellArchitectures",
]

# useful modules
from xautodl.config_utils import dict2config
from .SharedUtils import change_key
from .cell_searchs import CellStructure, CellArchitectures


# Cell-based NAS Models
def get_cell_based_tiny_net(config):
    if isinstance(config, dict):
        config = dict2config(config, None)  # to support the argument being a dict
        # print(config)
    super_type = getattr(config, "super_type", "basic")
    # print(super_type)
    group_names = ["DARTS-V1", "DARTS-V2", "GDAS", "SETN", "ENAS", "RANDOM", "generic"]
    if super_type == "basic" and config.name in group_names:
        from .cell_searchs import nas201_super_nets as nas_super_nets

        try:
            return nas_super_nets[config.name](
                config.C,
                config.N,
                config.max_nodes,
                config.num_classes,
                config.space,
                config.affine,
                config.track_running_stats,
            )
        except:
            return nas_super_nets[config.name](
                config.C, config.N, config.max_nodes, config.num_classes, config.space
            )
    elif super_type == "search-shape":
        from .shape_searchs import GenericNAS301Model

        genotype = CellStructure.str2structure(config.genotype)
        return GenericNAS301Model(
            config.candidate_Cs,
            config.max_num_Cs,
            genotype,
            config.num_classes,
            config.affine,
            config.track_running_stats,
        )
    elif super_type == "nasnet-super":
        from .cell_searchs import nasnet_super_nets as nas_super_nets

        return nas_super_nets[config.name](
            config.C,
            config.N,
            config.steps,
            config.multiplier,
            config.stem_multiplier,
            config.num_classes,
            config.space,
            config.affine,
            config.track_running_stats,
        )
    elif config.name == "infer.tiny":
        from .cell_infers import TinyNetwork

        if hasattr(config, "genotype"):
            genotype = config.genotype
        elif hasattr(config, "arch_str"):
            genotype = CellStructure.str2structure(config.arch_str)
        else:
            raise ValueError(
                "Can not find genotype from this config : {:}".format(config)
            )
        return TinyNetwork(config.C, config.N, genotype, config.num_classes)
    # sss 网络用到的
    elif config.name == "infer.shape.tiny":
        from .shape_infers import DynamicShapeTinyNet

        if isinstance(config.channels, str):
            channels = tuple([int(x) for x in config.channels.split(":")])
        else:
            channels = config.channels
        genotype = CellStructure.str2structure(config.genotype)
        return DynamicShapeTinyNet(channels, genotype, config.num_classes)
    elif config.name == "infer.nasnet-cifar":
        from .cell_infers import NASNetonCIFAR

        raise NotImplementedError
    else:
        raise ValueError("invalid network name : {:}".format(config.name))


# obtain the search space, i.e., a dict mapping the operation name into a python-function for this op
def get_search_spaces(xtype, name) -> List[Text]:
    if xtype == "cell" or xtype == "tss":  # The topology search space.
        from .cell_operations import SearchSpaceNames

        assert name in SearchSpaceNames, "invalid name [{:}] in {:}".format(
            name, SearchSpaceNames.keys()
        )
        return SearchSpaceNames[name]
    elif xtype == "sss":  # The size search space.
        if name in ["nats-bench", "nats-bench-size"]:
            return {"candidates": [8, 16, 24, 32, 40, 48, 56, 64], "numbers": 5}
        else:
            raise ValueError("Invalid name : {:}".format(name))
    else:
        raise ValueError("invalid search-space type is {:}".format(xtype))


def get_cifar_models(config, extra_path=None):
    super_type = getattr(config, "super_type", "basic")
    if super_type == "basic":
        from .CifarResNet import CifarResNet
        from .CifarDenseNet import DenseNet
        from .CifarWideResNet import CifarWideResNet

        if config.arch == "resnet":
            return CifarResNet(
                config.module, config.depth, config.class_num, config.zero_init_residual
            )
        elif config.arch == "densenet":
            return DenseNet(
                config.growthRate,
                config.depth,
                config.reduction,
                config.class_num,
                config.bottleneck,
            )
        elif config.arch == "wideresnet":
            return CifarWideResNet(
                config.depth, config.wide_factor, config.class_num, config.dropout
            )
        else:
            raise ValueError("invalid module type : {:}".format(config.arch))
    elif super_type.startswith("infer"):
        from .shape_infers import InferWidthCifarResNet
        from .shape_infers import InferDepthCifarResNet
        from .shape_infers import InferCifarResNet
        from .cell_infers import NASNetonCIFAR

        assert len(super_type.split("-")) == 2, "invalid super_type : {:}".format(
            super_type
        )
        infer_mode = super_type.split("-")[1]
        if infer_mode == "width":
            return InferWidthCifarResNet(
                config.module,
                config.depth,
                config.xchannels,
                config.class_num,
                config.zero_init_residual,
            )
        elif infer_mode == "depth":
            return InferDepthCifarResNet(
                config.module,
                config.depth,
                config.xblocks,
                config.class_num,
                config.zero_init_residual,
            )
        elif infer_mode == "shape":
            return InferCifarResNet(
                config.module,
                config.depth,
                config.xblocks,
                config.xchannels,
                config.class_num,
                config.zero_init_residual,
            )
        elif infer_mode == "nasnet.cifar":
            genotype = config.genotype
            if extra_path is not None:  # reload genotype by extra_path
                if not osp.isfile(extra_path):
                    raise ValueError("invalid extra_path : {:}".format(extra_path))
                xdata = torch.load(extra_path)
                current_epoch = xdata["epoch"]
                genotype = xdata["genotypes"][current_epoch - 1]
            C = config.C if hasattr(config, "C") else config.ichannel
            N = config.N if hasattr(config, "N") else config.layers
            return NASNetonCIFAR(
                C, N, config.stem_multi, config.class_num, genotype, config.auxiliary
            )
        else:
            raise ValueError("invalid infer-mode : {:}".format(infer_mode))
    else:
        raise ValueError("invalid super-type : {:}".format(super_type))


def get_imagenet_models(config):
    super_type = getattr(config, "super_type", "basic")
    if super_type == "basic":
        from .ImageNet_ResNet import ResNet
        from .ImageNet_MobileNetV2 import MobileNetV2

        if config.arch == "resnet":
            return ResNet(
                config.block_name,
                config.layers,
                config.deep_stem,
                config.class_num,
                config.zero_init_residual,
                config.groups,
                config.width_per_group,
            )
        elif config.arch == "mobilenet_v2":
            return MobileNetV2(
                config.class_num,
                config.width_multi,
                config.input_channel,
                config.last_channel,
                "InvertedResidual",
                config.dropout,
            )
        else:
            raise ValueError("invalid arch : {:}".format(config.arch))
    elif super_type.startswith("infer"):  # NAS searched architecture
        assert len(super_type.split("-")) == 2, "invalid super_type : {:}".format(
            super_type
        )
        infer_mode = super_type.split("-")[1]
        if infer_mode == "shape":
            from .shape_infers import InferImagenetResNet
            from .shape_infers import InferMobileNetV2

            if config.arch == "resnet":
                return InferImagenetResNet(
                    config.block_name,
                    config.layers,
                    config.xblocks,
                    config.xchannels,
                    config.deep_stem,
                    config.class_num,
                    config.zero_init_residual,
                )
            elif config.arch == "MobileNetV2":
                return InferMobileNetV2(
                    config.class_num, config.xchannels, config.xblocks, config.dropout
                )
            else:
                raise ValueError("invalid arch-mode : {:}".format(config.arch))
        else:
            raise ValueError("invalid infer-mode : {:}".format(infer_mode))
    else:
        raise ValueError("invalid super-type : {:}".format(super_type))


# Try to obtain the network by config.
def obtain_model(config, extra_path=None):
    if config.dataset == "cifar":
        return get_cifar_models(config, extra_path)
    elif config.dataset == "imagenet":
        return get_imagenet_models(config)
    else:
        raise ValueError("invalid dataset in the model config : {:}".format(config))


def obtain_search_model(config):
    if config.dataset == "cifar":
        if config.arch == "resnet":
            from .shape_searchs import SearchWidthCifarResNet
            from .shape_searchs import SearchDepthCifarResNet
            from .shape_searchs import SearchShapeCifarResNet

            if config.search_mode == "width":
                return SearchWidthCifarResNet(
                    config.module, config.depth, config.class_num
                )
            elif config.search_mode == "depth":
                return SearchDepthCifarResNet(
                    config.module, config.depth, config.class_num
                )
            elif config.search_mode == "shape":
                return SearchShapeCifarResNet(
                    config.module, config.depth, config.class_num
                )
            else:
                raise ValueError("invalid search mode : {:}".format(config.search_mode))
        elif config.arch == "simres":
            from .shape_searchs import SearchWidthSimResNet

            if config.search_mode == "width":
                return SearchWidthSimResNet(config.depth, config.class_num)
            else:
                raise ValueError("invalid search mode : {:}".format(config.search_mode))
        else:
            raise ValueError(
                "invalid arch : {:} for dataset [{:}]".format(
                    config.arch, config.dataset
                )
            )
    elif config.dataset == "imagenet":
        from .shape_searchs import SearchShapeImagenetResNet

        assert config.search_mode == "shape", "invalid search-mode : {:}".format(
            config.search_mode
        )
        if config.arch == "resnet":
            return SearchShapeImagenetResNet(
                config.block_name, config.layers, config.deep_stem, config.class_num
            )
        else:
            raise ValueError("invalid model config : {:}".format(config))
    else:
        raise ValueError("invalid dataset in the model config : {:}".format(config))


def load_net_from_checkpoint(checkpoint):
    assert osp.isfile(checkpoint), "checkpoint {:} does not exist".format(checkpoint)
    checkpoint = torch.load(checkpoint)
    model_config = dict2config(checkpoint["model-config"], None)
    model = obtain_model(model_config)
    model.load_state_dict(checkpoint["base-model"])
    return model