import sys
import os.path as osp
this_dir = osp.dirname(__file__)
lib_path = osp.join(this_dir, '..')
sys.path.insert(0, lib_path)
from NCF.dataprocess import DataProcess
from NCF.network import GMF,MLP,NeuMF
from NCF.trainer import Trainer
import numpy as np
import torch

gmf_config = {'num_epoch': 100,
              'batch_size': 1024,
              'optimizer': 'adam',
              'adam_lr': 1e-3,
              'num_users': 6040,
              'num_items': 3706,
              'latent_dim_gmf': 8,
              'num_negative': 4,
              'layers': [],
              'l2_regularization': 0, # 0.01
              'pretrain': False, # do not modify this
              'use_cuda': True,
              'device_id': 0,
              'model_name': '../TrainedModels/NCF_GMF.model'
              }

mlp_config = {'num_epoch': 100,
              'batch_size': 1024,  # 1024,
              'optimizer': 'adam',
              'adam_lr': 1e-3,
              'num_users': 6040,
              'num_items': 3706,
              'latent_dim_mlp': 8,
              'latent_dim_gmf': 8,
              'num_negative': 4,
              'layers': [16,64,32,16,8],  # layers[0] is the concat of latent user vector & latent item vector
              'l2_regularization': 0.0000001,  # MLP model is sensitive to hyper params
              'use_cuda': True,
              'device_id': 0,
              'pretrain': True,
              'gmf_config': gmf_config,
              'pretrain_gmf': '../TrainedModels/NCF_GMF.model',
              'model_name': '../TrainedModels/NCF_MLP.model'
              }

neumf_config = {'num_epoch': 100,
                'batch_size': 1024, #1024
                'optimizer': 'adam',
                'adam_lr': 1e-3,
                'num_users': 6040,
                'num_items': 3706,
                'latent_dim_gmf': 8,
                'latent_dim_mlp': 8,
                'num_negative': 4,
                'layers': [16,32,16,8],  # layers[0] 是用户和物品隐层表示concat的维度
                'l2_regularization': 0.01,
                'alpha': 0.5, # 用于控制GMF和MLP模型参数的权重
                'use_cuda': True,
                'device_id': 0,
                'pretrain': False,
                'gmf_config': gmf_config,
                'pretrain_gmf': '../TrainedModels/NCF_GMF.model',
                'mlp_config': mlp_config,
                'pretrain_mlp': '../TrainedModels/NCF_MLP.model',
                'model_name': '../TrainedModels/NCF_NeuMF.model'
                }

if __name__ == "__main__":
    ####################################################################################
    # NCF 神经协同过滤算法
    ####################################################################################

    # 加载和预处理数据
    dp = DataProcess("../Data/ml-1m/ratings.dat")

    # 初始化GMP模型
    config = gmf_config
    model = GMF(config, config['latent_dim_gmf'])

    # # 初始化MLP模型
    config = mlp_config
    model = MLP(config, config['latent_dim_mlp'])

    # 初始化NeuMF模型
    config = neumf_config
    model = NeuMF(config, config['latent_dim_gmf'], config['latent_dim_mlp'])

    # ###############################################################
    # 模型训练阶段
    # ###############################################################
    trainer = Trainer(model=model, config=config)
    trainer.train(dp.sample_generator)
    trainer.save()

    # ###############################################################
    # 模型测试阶段
    # ###############################################################

    # 加载数据集
    # dp = DataProcess("../Data/ml-1m/ratings.dat")

    config = neumf_config
    neumf = NeuMF(config, config['latent_dim_gmf'], config['latent_dim_mlp'])
    state_dict = torch.load("../TrainedModels/NCF_NeuMF.model", map_location=torch.device('cpu'))
    neumf.load_state_dict(state_dict, strict=False)

    # 对用户User_id喜好度进行预测
    User_id = 1
    result = np.zeros((3706))
    for j in range(3706):
        socre = neumf.forward(torch.LongTensor([User_id]), torch.LongTensor([j]))
        score = socre.detach().numpy()
        result[j] = socre[0][0]
    # 选取User_id喜好度最高的N个电影id进行推荐
    N = 5
    indexs = np.argsort(-result)[:N]
    print(indexs)