##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

## Created by: Hang Zhang

## Email: zhanghang0704@gmail.com

## Copyright (c) 2020

##

## LICENSE file in the root directory of this source tree 

##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

from .resnest import *

from .ablation import *

_all__ = ['get_model', 'get_model_list']

models = {

    'resnest50': resnest50,

    'resnest101': resnest101,

    'resnest200': resnest200,

    'resnest269': resnest269,

    'resnest50_fast_1s1x64d': resnest50_fast_1s1x64d,

    'resnest50_fast_2s1x64d': resnest50_fast_2s1x64d,

    'resnest50_fast_4s1x64d': resnest50_fast_4s1x64d,

    'resnest50_fast_1s2x40d': resnest50_fast_1s2x40d,

    'resnest50_fast_2s2x40d': resnest50_fast_2s2x40d,

    'resnest50_fast_4s2x40d': resnest50_fast_4s2x40d,

    'resnest50_fast_1s4x24d': resnest50_fast_1s4x24d,

    }

def get_model(name, **kwargs):

    """Returns a pre-defined model by name

    Parameters

    ----------

    name : str

        Name of the model.

    pretrained : bool

        Whether to load the pretrained weights for model.

    root : str, default '~/.encoding/models'

        Location for keeping the model parameters.

    Returns

    -------

    Module:

        The model.

    """

    name = name.lower()

    if name in models:

        net = models[name](**kwargs)

    else:

        raise ValueError('%s\n\t%s' % (str(name), '\n\t'.join(sorted(models.keys()))))

    return net

def get_model_list():

    """Get the entire list of model names in model_zoo.

    Returns

    -------

    list of str

        Entire list of model names in model_zoo.

    """

    return models.keys()