# coding: UTF-8
import torch
import torch.nn as nn
from utils import *


# 权重初始化，默认xavier
def init_network(model, method='xavier', exclude='embedding', seed=123):
    for name, w in model.named_parameters():
        if exclude not in name:
            if len(w.size()) < 2:
                continue
            if 'weight' in name:
                if method == 'xavier':
                    nn.init.xavier_normal_(w)
                elif method == 'kaiming':
                    nn.init.kaiming_normal_(w)
                else:
                    nn.init.normal_(w)
            elif 'bias' in name:
                nn.init.constant_(w, 0)
            else:
                pass


def train(config, model, train_iter, dev_iter):
    start_time = time.time()
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
    total_batch = 0  # 记录进行到多少batch
    dev_best_auc = float('-inf')
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升
    model.train()
    for epoch in range(config.num_epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))
        for i, batch in enumerate(train_iter):
            # print(trains)
            dense_features, single_id_concat, multi_id_concat, mask_concat, labels, _ = \
                (x.to(config.device) for x in batch)
            # print(dense_features, single_id_concat, multi_id_concat, mask_concat, labels)
            outputs = model(dense_features, single_id_concat, multi_id_concat, mask_concat)

            model.zero_grad()
            loss = weighted_loss(outputs, labels.float())
            loss.backward()
            optimizer.step()
            if total_batch % 100 == 0:
                # 每多少轮输出在训练集和验证集上的效果
                dev_auc, dev_loss = evaluate(config, model, dev_iter)
                if dev_auc > dev_best_auc:
                    dev_best_auc = dev_auc
                    torch.save(model.state_dict(), config.save_path)
                    improve = '*'
                    last_improve = total_batch
                else:
                    improve = ''
                time_dif = get_time_dif(start_time)
                msg = 'Iter: {0:>6},  Train Loss: {1:>5.3},  Val Loss: {2:>5.3},  Val AUC: {3:>6.3},Time: {4} {5}'
                print(msg.format(total_batch, loss.item(), dev_loss.item(), dev_auc, time_dif, improve))
                model.train()
            total_batch += 1
            # print(total_batch)
            if total_batch - last_improve > config.require_improvement:
                # 验证集loss超过1000batch没下降，结束训练
                print("No optimization for a long time, auto-stopping...")
                flag = True
                break
        if flag:
            break


def test(config, model, test_iter):
    # test
    model.load_state_dict(torch.load(config.save_path))
    model.eval()
    predict_all, auc, _ = evaluate(config, model, test_iter, test=True)
    print('test auc:' + str(auc))
    save_pred_file(config, predict_all)


def evaluate(config, model, data_iter, test=False):
    model.eval()
    loss_total = 0

    predict_all = np.empty([0, 4])
    labels_all = np.empty([0, 4])
    userid_all = np.array([], dtype=int)
    with torch.no_grad():
        for batch in data_iter:
            dense_features, single_id_concat, multi_id_concat, mask_concat, labels, userids = \
                (x.to(config.device) for x in batch)
            outputs = model(dense_features, single_id_concat, multi_id_concat, mask_concat)
            loss = weighted_loss(outputs, labels.float())
            loss_total += loss

            labels = labels.data.cpu().numpy()
            predict = outputs.data.cpu()
            userids = userids.data.cpu()

            labels_all = np.vstack((labels_all, labels))
            predict_all = np.vstack((predict_all, predict))
            userid_all = np.append(userid_all, userids)

    uauc = weighted_auc(labels_all, predict_all, userid_all)
    if test:
        return predict_all, uauc, loss_total / len(data_iter)
    return uauc, loss_total / len(data_iter)
