import itertools

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


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)
