# -*- coding: utf-8 -*-
# Created by li huayong on 2019/10/7
import re
import math
import torch
# from model.optimization import *
import pytorch_transformers.optimization as huggingfaceOptim  # 避免和torch.optim重名
from PyToolkit.PyToolkit import get_logger, debug_print


def get_optimizer_old(name, parameters, lr, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
    if name == 'sgd':
        return torch.optim.SGD(parameters, lr=lr)
    elif name == 'adagrad':
        return torch.optim.Adagrad(parameters, lr=lr)
    elif name == 'adam':
        return torch.optim.Adam(parameters, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
    elif name == 'adamax':
        return torch.optim.Adamax(parameters)  # use default lr
    else:
        raise Exception("Unsupported optimizer: {}".format(name))


def _get_bertology_optimizer_grouped_parameters(args, model):
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
         'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
         'weight_decay': 0.0}
    ]
    return optimizer_grouped_parameters


def _get_bertology_different_lr_grouped_parameters(args, model):
    no_decay = ['bias', 'LayerNorm.weight']
    params_bert_no_decay = []
    params_bert_decay = []
    params_other_no_decay = []
    params_other_decay = []
    for n, p in model.named_parameters():
        if 'encoder.bertology.' in n:
            # 是BERTology的参数
            if any(nd in n for nd in no_decay):
                # no decay
                params_bert_no_decay.append(p)
            else:
                params_bert_decay.append(p)
        else:
            # 不是bertology的参数
            if any(nd in n for nd in no_decay):
                # no decay
                params_other_no_decay.append(p)
            else:
                params_other_decay.append(p)
    optimizer_grouped_parameters = [
        {
            'params': params_bert_decay,
            'weight_decay': args.weight_decay,
            'lr': args.bertology_lr,
        },
        {
            'params': params_bert_no_decay,
            'weight_decay': 0.0,
            'lr': args.bertology_lr,
        },
        {
            'params': params_other_decay,
            'weight_decay': args.weight_decay,
            'lr': args.other_lr,
        },
        {
            'params': params_other_no_decay,
            'weight_decay': 0.0,
            'lr': args.other_lr,
        },
    ]
    return optimizer_grouped_parameters


def get_optimizer(args, model):
    logger = get_logger(args.log_name)
    args.warmup_steps = math.ceil(args.warmup_prop * args.max_train_steps)
    if args.optimizer == 'adamw-bertology':
        if args.different_lr:
            optimizer_grouped_parameters = _get_bertology_different_lr_grouped_parameters(args, model)
        else:
            optimizer_grouped_parameters = _get_bertology_optimizer_grouped_parameters(args, model)
        optimizer = huggingfaceOptim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon,
                                           betas=(args.beta1, args.beta2))
        scheduler = huggingfaceOptim.WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps,
                                                          t_total=args.max_train_steps)
        if args.local_rank in [-1, 0]:
            logger.info('Use Huggingface\'s AdamW Optimizer')
    elif args.optimizer == 'adamw-torch':
        try:
            from torch.optim import AdamW
        except ImportError as e:
            debug_print(f'torch version: {torch.__version__}')
            raise e
        if args.different_lr:
            optimizer_grouped_parameters = _get_bertology_different_lr_grouped_parameters(args, model)
        else:
            optimizer_grouped_parameters = _get_bertology_optimizer_grouped_parameters(args, model)
        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon,
                          betas=(args.beta1, args.beta2))
        scheduler = huggingfaceOptim.WarmupLinearSchedule(optimizer,
                                                          warmup_steps=args.warmup_steps,
                                                          t_total=args.max_train_steps)
    elif args.optimizer == 'sgd':
        optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate)
        scheduler = None
    elif args.optimizer == 'adagrad':
        optimizer = torch.optim.Adagrad(model.parameters(), lr=args.learning_rate)
        scheduler = None
    elif args.optimizer == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, betas=args.betas, eps=args.eps,
                                     weight_decay=args.weight_decay)
        scheduler = None
    elif args.rnn_optimizer == 'adamax':
        optimizer = torch.optim.Adamax(model.parameters())  # use default lr
        scheduler = None
    else:
        raise Exception("Unsupported optimizer: {}".format(args.optimizer))
    return optimizer, scheduler


if __name__ == '__main__':
    pass
