import time
import argparse
import pickle
from model import *
from utils import *
import matplotlib.pyplot as plt

def init_seed(seed=None):
    if seed is None:
        seed = int(time.time() * 1000 // 1000)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


parser = argparse.ArgumentParser()
parser.add_argument('--dataset', default='diginetica', help='diginetica/Nowplaying/Tmall')
parser.add_argument('--hiddenSize', type=int, default=100)
parser.add_argument('--epoch', type=int, default=20)
parser.add_argument('--activate', type=str, default='relu')
parser.add_argument('--n_sample_all', type=int, default=12)
parser.add_argument('--n_sample', type=int, default=12)
parser.add_argument('--batch_size', type=int, default=100)
parser.add_argument('--lr', type=float, default=0.001, help='learning rate.')
parser.add_argument('--lr_dc', type=float, default=0.1, help='learning rate decay.')
parser.add_argument('--lr_dc_step', type=int, default=3, help='the number of steps after which the learning rate decay.')
parser.add_argument('--l2', type=float, default=1e-5, help='l2 penalty ')
parser.add_argument('--n_iter', type=int, default=1)                                    # [1, 2]
parser.add_argument('--dropout_gcn', type=float, default=0, help='Dropout rate.')       # [0, 0.2, 0.4, 0.6, 0.8]
parser.add_argument('--dropout_local', type=float, default=0, help='Dropout rate.')     # [0, 0.5]
parser.add_argument('--dropout_global', type=float, default=0.5, help='Dropout rate.')
parser.add_argument('--validation', action='store_true', help='validation')
parser.add_argument('--valid_portion', type=float, default=0.1, help='split the portion')
parser.add_argument('--alpha', type=float, default=0.2, help='Alpha for the leaky_relu.')
parser.add_argument('--patience', type=int, default=3)
parser.add_argument('--num_latent',type=int, default=10)

opt = parser.parse_args()


def main():
    init_seed(2020)

    if opt.dataset == 'diginetica':
        num_node = 43098
        #num_node = 5590
        opt.n_iter = 2
        opt.dropout_gcn = 0.2
        opt.dropout_local = 0.0
    elif opt.dataset == 'Nowplaying':
        num_node = 60417
        #num_node = 6384
        opt.n_iter = 1
        opt.dropout_gcn = 0.0
        opt.dropout_local = 0.0
    elif opt.dataset == 'Tmall':
        num_node = 40728
        opt.n_iter = 1
        opt.dropout_gcn = 0.6
        opt.dropout_local = 0.5
    elif opt.dataset == 'lastfm':
        num_node = 38616
        opt.n_iter = 1
        opt.dropout_global = 0.1
        opt.dropout_gcn = 0
        opt.dropout_local = 0
    elif opt.dataset == 'retailrocket':
        num_node = 36969
        opt.n_iter = 1
        opt.dropout_global = 0.5
        opt.dropout_gcn = 0.8
        opt.dropout_local = 0.0
    elif opt.dataset == 'yoochoose64':
        num_node = 37484
        opt.n_iter =1
        opt.dropout_gcn = 0.2
        opt.dropout_local = 0.0   
    else:
        num_node = 310

    train_data = pickle.load(open('pretrain/datasets/' + opt.dataset + '/train.txt', 'rb'))
    if opt.validation:
        train_data, valid_data = split_validation(train_data, opt.valid_portion)
        test_data = valid_data
    else:
        test_data = pickle.load(open('pretrain/datasets/' + opt.dataset + '/test.txt', 'rb'))

    adj = pickle.load(open('datasets/' + opt.dataset + '/adj_' + str(opt.n_sample_all) + '.pkl', 'rb'))
    num = pickle.load(open('datasets/' + opt.dataset + '/num_' + str(opt.n_sample_all) + '.pkl', 'rb'))
    train_data = Data(train_data)
    test_data = Data(test_data)

    adj, num = handle_adj(adj, num_node, opt.n_sample_all, num)
    gnn = trans_to_cuda(CombineGraph(opt, num_node, adj, num))
    model = trans_to_cuda(NP_MISR(opt, num_node, gnn))
    

    # 定义不同的 k 值
    k_values = [3, 5, 10, 20]
    
    print(opt)
    start = time.time()
    best_result = {k: [0, 0, 0] for k in k_values}
    best_epoch = {k: [0, 0, 0] for k in k_values}
    bad_counter = 0
    
    for epoch in range(opt.epoch):
        print('-------------------------------------------------------')
        print('epoch: ', epoch)
        results = train_test(model, train_data, test_data)
        path = f'retailrocketWeight/np_weights{epoch}.pth'
        torch.save(model.state_dict(), path)
        print(f"NP has been saved to '{path}'")
        for i, k in enumerate(k_values):
            hit = results[k][0]
            mrr = results[k][1]
            cov = results[k][2]
            cov = cov * 100 / (num_node - 1)
            flag = 0
            if hit >= best_result[k][0]:
                best_result[k][0] = hit
                best_epoch[k][0] = epoch
                flag = 1
            if mrr >= best_result[k][1]:
                best_result[k][1] = mrr
                best_epoch[k][1] = epoch
                flag = 1
            if cov >= best_result[k][2]:
                best_result[k][2] = cov
                best_epoch[k][2] = epoch
                flag = 1
            print(f'Current Result (k={k}):')
            print(f'\tRecall@{k}:\t{hit:.4f}\tMRR@{k}:\t{mrr:.4f}\tCov@{k}:\t{cov:.4f}')
            print(f'Best Result (k={k}):')
            print(
                f'\tRecall@{k}:\t{best_result[k][0]:.4f}\tMRR@{k}:\t{best_result[k][1]:.4f}\tCov@{k}:\t{best_result[k][2]:.4f}\t\tEpoch:\t{best_epoch[k][0]},\t{best_epoch[k][1]},\t{best_epoch[k][2]}')
        bad_counter += 1 - min([flag for k, flag in enumerate(k_values)])
        if bad_counter >= opt.patience:
            break
    
    print('-------------------------------------------------------')
    end = time.time()
    print("Run time: %f s" % (end - start))

if __name__ == '__main__':
    main()

