import logging
import sys
from models.modules_search import *
from models.modules_retrain import *
from dataloader.tfloader import CriteoLoader, Avazuloader, KDD12loader
import torch
import pickle
import os

def get_dataloader(dataset, path):
    if dataset == 'Criteo':
        return CriteoLoader(path)
    elif dataset == 'Avazu':
        return Avazuloader(path)
    elif dataset == 'KDD12':
        return KDD12loader(path)

def get_stats(path):
    defaults_path = os.path.join(path + "/defaults.pkl")
    with open(defaults_path, 'rb') as fi:
        defaults = pickle.load(fi)
    return [i+1 for i in list(defaults.values())] 

def get_search(opt):
    name = opt["model"]
    if name == "fm":
        model = FM_search(opt)
    elif name == "deepfm":
        model = DeepFM_search(opt)
    elif name == "fnn":
        model = FNN_search(opt)
    elif name == "ipnn":
        model = IPNN_search(opt)
    elif name == "dcn":
        model = DCN_search(opt)
    else:
        raise ValueError("Invalid model type: {}".format(name))
    return model

def get_retrain(opt, arch):
    name = opt["model"]
    if name == "fm":
        model = FM_retrain(opt, arch)
    elif name == "deepfm":
        model = DeepFM_retrain(opt, arch)
    elif name == "fnn":
        model = FNN_retrain(opt, arch)
    elif name == "ipnn":
        model = IPNN_retrain(opt, arch)
    elif name == "dcn":
        model = DCN_retrain(opt, arch)
    else:
        raise ValueError("Invalid model type: {}".format(name))
    return model

def get_optimizer(network, opt):
    arch_params, trans_params, network_params, embedding_params = [], [], [], []
    arch_names, trans_names, network_names, embedding_names = [], [], [], []
    for name, param in network.named_parameters():
        if name == "arch":
            arch_params.append(param)
            arch_names.append(name)
        elif "trans" in name:
            trans_params.append(param)
            trans_names.append(name)
        elif name == "embedding":
            embedding_params.append(param)
            embedding_names.append(name)
        else:
            network_params.append(param)
            network_names.append(name)

    # if opt['transform'] == 'linear':
    #     for name, param in network.trans_b.named_parameters():
    #         trans_params.append(param)
    #         trans_names.append(name)
    #     for name, param in network.trans_bn.named_parameters():
    #         trans_params.append(param)
    #         trans_names.append(name)
    
    print("arch_names:", arch_names)
    print("_"*30)
    print("trans_names:", trans_names)
    print("_"*30)
    print("embedding_names:", embedding_names)
    print("_"*30)
    print("network_names:", network_names)
    print("_"*30)
    
    arch_group = {
        "params": arch_params,
        "lr": opt["arch_lr"]
    }
    arch_optimizer = torch.optim.SGD([arch_group])

    embedding_group = {
        'params': embedding_params,
        'weight_decay': opt['wd'],
        'lr': opt['lr']
    }
    trans_group = {
        'params': trans_params,
        'lr': opt['lr']
    }
    network_group = {
        'params': network_params,
        'weight_decay': opt['wd'],
        'lr': opt['lr']
    }
    if opt['optimizer'] == 'sgd':
        optimizer = torch.optim.SGD([network_group, trans_group, embedding_group])
    elif opt['optimizer'] == 'adam':
        optimizer = torch.optim.Adam([network_group, trans_group, embedding_group])
    else:
        print("Optimizer not supported.")
        sys.exit(-1)

    return arch_optimizer, optimizer

def get_cuda(enabled, device_id=0):
    if enabled:
        assert torch.cuda.is_available(), 'CUDA is not available'
        torch.cuda.set_device(device_id)

def get_log(name=""):
    FORMATTER = logging.Formatter(fmt="[{asctime}]:{message}", style= '{')
    logger = logging.getLogger(name)
    logger.setLevel(logging.DEBUG)
    ch = logging.StreamHandler(stream=sys.stdout)
    ch.setFormatter(FORMATTER)
    logger.addHandler(ch)
    return logger
