import torch
import torch.optim as optim
import numpy as np
import itertools

def singleton(class_):
    instances = {}
    def getinstance(*args, **kwargs):
        if class_ not in instances:
            instances[class_] = class_(*args, **kwargs)
        return instances[class_]
    return getinstance

class get_optimizer(object):
    def __init__(self):
        self.optimizer = {}
        self.register(optim.SGD, 'sgd')
        self.register(optim.Adam, 'adam')
        self.register(optim.AdamW, 'adamw')

    def register(self, optim, name):
        self.optimizer[name] = optim

    def __call__(self, net, cfg):
        if cfg is None:
            return None
        t = cfg.type
        if isinstance(net, (torch.nn.DataParallel,
                            torch.nn.parallel.DistributedDataParallel)):
            netm = net.module
        else:
            netm = net
        pg = getattr(netm, 'parameter_group', None)

        if pg is not None:
            params = []
            for group_name, module_or_para in pg.items():
                if not isinstance(module_or_para, list):
                    module_or_para = [module_or_para]

                grouped_params = [mi.parameters() if isinstance(mi, torch.nn.Module) else [mi] for mi in module_or_para]
                grouped_params = itertools.chain(*grouped_params)
                pg_dict = {'params':grouped_params, 'name':group_name}
                params.append(pg_dict)
        else:
            params = net.parameters()
        return self.optimizer[t](params, lr=0, **cfg.args)