""" Model Registry
Hacked together by / Copyright 2020 Ross Wightman
"""

import fnmatch
import re
import sys
from collections import defaultdict, deque
from copy import deepcopy
from typing import List, Optional, Union, Tuple

from ._pretrained import PretrainedCfg, DefaultCfg, split_model_name_tag

__all__ = [
    'list_models', 'is_model', 'model_entrypoint', 'list_modules', 'is_model_in_modules',
    'get_pretrained_cfg_value', 'is_model_pretrained', 'get_arch_name']

_module_to_models = defaultdict(set)  # dict of sets to check membership of model in module
_model_to_module = {}  # mapping of model names to module names
_model_entrypoints = {}  # mapping of model names to architecture entrypoint fns
_model_has_pretrained = set()  # set of model names that have pretrained weight url present
_model_default_cfgs = dict()  # central repo for model arch -> default cfg objects
_model_pretrained_cfgs = dict()  # central repo for model arch + tag -> pretrained cfgs
_model_with_tags = defaultdict(list)  # shortcut to map each model arch to all model + tag names


def get_arch_name(model_name: str) -> Tuple[str, Optional[str]]:
    return split_model_name_tag(model_name)[0]


def register_model(fn):
    # lookup containing module
    mod = sys.modules[fn.__module__]
    module_name_split = fn.__module__.split('.')
    module_name = module_name_split[-1] if len(module_name_split) else ''

    # add model to __all__ in module
    model_name = fn.__name__
    if hasattr(mod, '__all__'):
        mod.__all__.append(model_name)
    else:
        mod.__all__ = [model_name]

    # add entries to registry dict/sets
    _model_entrypoints[model_name] = fn
    _model_to_module[model_name] = module_name
    _module_to_models[module_name].add(model_name)
    if hasattr(mod, 'default_cfgs') and model_name in mod.default_cfgs:
        # this will catch all models that have entrypoint matching cfg key, but miss any aliasing
        # entrypoints or non-matching combos
        cfg = mod.default_cfgs[model_name]
        if not isinstance(cfg, DefaultCfg):
            # new style default cfg dataclass w/ multiple entries per model-arch
            assert isinstance(cfg, dict)
            # old style cfg dict per model-arch
            cfg = PretrainedCfg(**cfg)
            cfg = DefaultCfg(tags=deque(['']), cfgs={'': cfg})

        for tag_idx, tag in enumerate(cfg.tags):
            is_default = tag_idx == 0
            pretrained_cfg = cfg.cfgs[tag]
            if is_default:
                _model_pretrained_cfgs[model_name] = pretrained_cfg
                if pretrained_cfg.has_weights:
                    # add tagless entry if it's default and has weights
                    _model_has_pretrained.add(model_name)
            if tag:
                model_name_tag = '.'.join([model_name, tag])
                _model_pretrained_cfgs[model_name_tag] = pretrained_cfg
                if pretrained_cfg.has_weights:
                    # add model w/ tag if tag is valid
                    _model_has_pretrained.add(model_name_tag)
                _model_with_tags[model_name].append(model_name_tag)
            else:
                _model_with_tags[model_name].append(model_name)  # has empty tag (to slowly remove these instances)

        _model_default_cfgs[model_name] = cfg

    return fn


def _natural_key(string_):
    return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]


def list_models(
        filter: Union[str, List[str]] = '',
        module: str = '',
        pretrained=False,
        exclude_filters: str = '',
        name_matches_cfg: bool = False,
        include_tags: Optional[bool] = None,
):
    """ Return list of available model names, sorted alphabetically

    Args:
        filter (str) - Wildcard filter string that works with fnmatch
        module (str) - Limit model selection to a specific submodule (ie 'vision_transformer')
        pretrained (bool) - Include only models with valid pretrained weights if True
        exclude_filters (str or list[str]) - Wildcard filters to exclude models after including them with filter
        name_matches_cfg (bool) - Include only models w/ model_name matching default_cfg name (excludes some aliases)
        include_tags (Optional[boo]) - Include pretrained tags in model names (model.tag). If None, defaults
            set to True when pretrained=True else False (default: None)
    Example:
        model_list('gluon_resnet*') -- returns all models starting with 'gluon_resnet'
        model_list('*resnext*, 'resnet') -- returns all models with 'resnext' in 'resnet' module
    """
    if include_tags is None:
        # FIXME should this be default behaviour? or default to include_tags=True?
        include_tags = pretrained

    if module:
        all_models = list(_module_to_models[module])
    else:
        all_models = _model_entrypoints.keys()

    if include_tags:
        # expand model names to include names w/ pretrained tags
        models_with_tags = []
        for m in all_models:
            models_with_tags.extend(_model_with_tags[m])
        all_models = models_with_tags

    if filter:
        models = []
        include_filters = filter if isinstance(filter, (tuple, list)) else [filter]
        for f in include_filters:
            include_models = fnmatch.filter(all_models, f)  # include these models
            if len(include_models):
                models = set(models).union(include_models)
    else:
        models = all_models

    if exclude_filters:
        if not isinstance(exclude_filters, (tuple, list)):
            exclude_filters = [exclude_filters]
        for xf in exclude_filters:
            exclude_models = fnmatch.filter(models, xf)  # exclude these models
            if len(exclude_models):
                models = set(models).difference(exclude_models)

    if pretrained:
        models = _model_has_pretrained.intersection(models)

    if name_matches_cfg:
        models = set(_model_pretrained_cfgs).intersection(models)

    return list(sorted(models, key=_natural_key))


def list_pretrained(
        filter: Union[str, List[str]] = '',
        exclude_filters: str = '',
):
    return list_models(
        filter=filter,
        pretrained=True,
        exclude_filters=exclude_filters,
        include_tags=True,
    )


def is_model(model_name):
    """ Check if a model name exists
    """
    arch_name = get_arch_name(model_name)
    return arch_name in _model_entrypoints


def model_entrypoint(model_name):
    """Fetch a model entrypoint for specified model name
    """
    arch_name = get_arch_name(model_name)
    return _model_entrypoints[arch_name]


def list_modules():
    """ Return list of module names that contain models / model entrypoints
    """
    modules = _module_to_models.keys()
    return list(sorted(modules))


def is_model_in_modules(model_name, module_names):
    """Check if a model exists within a subset of modules
    Args:
        model_name (str) - name of model to check
        module_names (tuple, list, set) - names of modules to search in
    """
    arch_name = get_arch_name(model_name)
    assert isinstance(module_names, (tuple, list, set))
    return any(arch_name in _module_to_models[n] for n in module_names)


def is_model_pretrained(model_name):
    return model_name in _model_has_pretrained


def get_pretrained_cfg(model_name):
    if model_name in _model_pretrained_cfgs:
        return deepcopy(_model_pretrained_cfgs[model_name])
    raise RuntimeError(f'No pretrained config exists for model {model_name}.')


def get_pretrained_cfg_value(model_name, cfg_key):
    """ Get a specific model default_cfg value by key. None if key doesn't exist.
    """
    if model_name in _model_pretrained_cfgs:
        return getattr(_model_pretrained_cfgs[model_name], cfg_key, None)
    raise RuntimeError(f'No pretrained config exist for model {model_name}.')