import copy
from torch import nn


class Registry(dict):
    r""" A helper class for managing registering modules, it extends a dictionary
    and provides a register function.
    """
    def __init__(self, *args, **kwargs):
        super(Registry, self).__init__(*args, **kwargs)

    def register_module(self, module_class, module_name=None):
        if module_name is None:
            module_name = module_class.__name__
        assert module_name not in self, module_name
        self[module_name] = module_class
        return module_class

    def register(self, module_class):
        if isinstance(module_class, str):
            def _register(_module_class):
                return self.register_module(_module_class, module_class)
            return _register
        else:
            return self.register_module(module_class)


class ConvertMappingRegistry(dict):
    r""" A helper class for managing registering modules, it extends a dictionary
    and provides a register function.
    """
    def __init__(self, *args, **kwargs):
        super(ConvertMappingRegistry, self).__init__(*args, **kwargs)

    def register_module(self, convert_module_class, module_class):
        self[module_class] = convert_module_class
        return convert_module_class

    def register(self, module_class):
        def _register(_module_class):
            return self.register_module(_module_class, module_class)
        return _register

class ParamsMappingRegistry(dict):
    r""" A helper class for managing registering modules, it extends a dictionary
    and provides a register function.
    """
    def __init__(self, *args, **kwargs):
        super(ParamsMappingRegistry, self).__init__(*args, **kwargs)

    def register(self, module_class):
        if hasattr(module_class, '_PARAMS_DICT'):
            d = module_class._PARAMS_DICT
        else:
            d = {}
        self[module_class.__name__] = d
        return module_class


class AbbrevMappingRegistry(dict):
    r""" A helper class for managing registering modules, it extends a dictionary
    and provides a register function.
    """
    def __init__(self, *args, **kwargs):
        super(AbbrevMappingRegistry, self).__init__(*args, **kwargs)

    def register_module(self, module_class, abbrev_name):
        self[abbrev_name] = module_class.__name__
        return module_class

    def register(self, module_class):
        def _register(_module_class):
            return self.register_module(_module_class, module_class)
        return _register

def merge_params(params_or_type, **kwargs):
    if isinstance(params_or_type, str):
        assert 'type' not in kwargs
        kwargs['type'] = params_or_type
    elif isinstance(params_or_type, dict):
        params = copy.deepcopy(params_or_type)
        for name, value in params.items():
            assert name not in kwargs
            kwargs[name] = value
    else:
        assert params_or_type is None
    return kwargs


def build_module(registry, params_or_type, *args, **kwargs):
    params = merge_params(params_or_type, **kwargs)
    if 'type' in params:
        obj_type = params.pop('type')
        assert obj_type in registry, '%s not in registry' % obj_type
        return registry[obj_type](*args, **params)
    else:
        return None
    
LAYERS = Registry()
CONVS = Registry()
NORMS = Registry()
ACTIVATIONS = Registry()
POOLS = Registry()
LOSSES = Registry()
UPSAMPLES = Registry()
DROPOUTS = Registry()
LINEARS = Registry()
PADDINGS = Registry()
RNNS = Registry()
PLUGINS = Registry()
BLOCKS = Registry()
# _all_build_fn = dict()

PARAMS_MAPPING = ParamsMappingRegistry()
ABBREVS = AbbrevMappingRegistry()


def build_block(params_or_type, *args, **kwargs):
    return build_module(BLOCKS, params_or_type, *args, **kwargs)

def build_layer(params_or_type, *args, **kwargs):
    return build_module(LAYERS, params_or_type, *args, **kwargs)

_build_fn = {
    'block': build_block,
    'layer': build_layer,
}

def build_nn(op_type, params_or_type, *args, **kwargs):
    return _build_fn[op_type](params_or_type, *args, **kwargs)


def decode_params_str(params, convert_dict=None, name=''):
    module_config = []
    
    def trace_back(params_idx, total_params_num, name, config):
        if params_idx == total_params_num:
            config['op_name'] = name
            module_config.append(copy.deepcopy(config))
            return 
        else:
            s = params[params_idx]
            if ':' in s:
                kv = s.split(':')
                assert len(kv) >= 2
                k, v = kv[0], kv[1]
                if v.startswith('[') and v.endswith(']'):
                    v = v[1:-1].split(',')
                else:
                    v = [v]

                for vi in v:
                    config[convert_dict[k]] = vi
                    if len(name) == 0:
                        new_name = vi
                    else:
                        new_name = name + '_' + vi
                    trace_back(params_idx+1, total_params_num, new_name, config)
                    config.pop(convert_dict[k])
                
            elif s in convert_dict:
                config[convert_dict[s]] = True
                if len(name) == 0:
                    new_name = s
                else:
                    new_name = name + '_' + s
                trace_back(params_idx+1, total_params_num, new_name, config)
                config.pop(convert_dict[s])
            elif s[0] in convert_dict:
                if s[1:].startswith('[') and s[1:].endswith(']'):
                    l = [int(v) for v in s[2:-1].split(',')] if '.' not in v else [float(v) for v in s[2:-1].split(',')]
                else:
                    l = [int(s[1:])] if '.' not in s else [float(s[1:])]

                for param in l:
                    if len(name) == 0:
                        new_name = s[0] + str(param)
                    else:
                        new_name = name + '_' + s[0] + str(param)
                    config[convert_dict[s[0]]] = param
                    trace_back(params_idx+1, total_params_num, new_name, config)
                    config.pop(convert_dict[s[0]])
                    
            else:
                trace_back(params_idx+1, total_params_num, name, config)
    if len(params) > 0:
        trace_back(0, len(params), name, {})
    return module_config

def decode(params):
    if isinstance(params, dict):
        decode_params = copy.deepcopy(params)
        for k in params:
            if k == 'params_str':
                v = decode_params.pop('params_str')
                op_type = params.get('op_type', None)
                type = params.get('type', None)
                if type is not None:
                    if op_type != 'supernet':
                        params_dict = copy.deepcopy(PARAMS_MAPPING[type])
                        decode_params['op_type'] = params_dict.pop('op_type', op_type)
                        if isinstance(v, list):
                            params_list = []
                            for vi in v:
                                params_list += decode_params_str(vi.split('_'), params_dict)
                        else:
                            params_list = decode_params_str(v.split('_'), params_dict)

                        if len(params_list) == 1:
                            decode_params = merge_params(params_list[0], **decode_params)
                            decode_params.pop('op_name', None)
                        else:
                            op_params_list = []
                            for param in params_list:
                                param.update({'op_type': op_type, 'type': type})
                                op_params_list.append(param)
                            decode_params = {
                                'op_type': 'supernet',
                                'type': 'SupernetModule',
                                'op_params_list': op_params_list,
                            }
                    else:
                        op_params_list = decode_params.pop('op_params_list', [])
                        decode_params['op_params_list'] = decode(v).get('op_params_list', []) + op_params_list
                else:
                    decode_params = merge_params(decode(v), **decode_params)
            else:
                decode_params[k] = decode(params[k])
        return decode_params
    elif isinstance(params, list) or isinstance(params, str):
        if isinstance(params, list):

            # if params can be decoded
            if not isinstance(params[0], str) or params[0].split('_')[0] not in ABBREVS:
                return params

            params_list = []
            for p in params:
                p_list = p.split('_')
                type = ABBREVS[p_list[0]]
                params_dict = copy.deepcopy(PARAMS_MAPPING[type])
                decode_params = {
                    'op_type': params_dict.pop('op_type', 'block'),
                    'type': type
                }
                params_list += [merge_params(dp, **decode_params) for dp in decode_params_str(p_list[1:], params_dict)]
        else:
            p_list = params.split('_')
            if p_list[0] in ABBREVS:
                type = ABBREVS[p_list[0]]
                params_dict = copy.deepcopy(PARAMS_MAPPING[type])
                decode_params = {
                    'op_type': params_dict.pop('op_type', 'block'),
                    'type': type
                }
                params_list = [merge_params(dp, **decode_params) for dp in decode_params_str(p_list[1:], params_dict)]
            else:
                return params

        if len(params_list) == 1:
            decode_params = params_list[0]
            decode_params.pop('op_name', None)
            return decode_params
        else:
            return {
                'op_type': 'supernet',
                'type': 'SupernetModule',
                'op_params_list': params_list,
            }
    else:
        return params

