# -*- coding: utf-8 -*-
import numpy as np
import torch
import pdb
from sklearn.metrics import roc_auc_score
import pdb
import arguments

from dataset import load_data
from matrix_factorization_logistic_gpu_uidr import *
from utils import gini_index, ndcg_func, get_user_wise_ctr, rating_mat_to_sample, binarize, shuffle, minU, precision_func, recall_func
mse_func = lambda x,y: np.mean((x-y)**2)
acc_func = lambda x,y: np.sum(x == y) / len(x)
mae_func = lambda x,y: np.mean(np.abs(x-y))



def train_and_eval(dataset_name, train_args, model_args):

    top_k_list = [5]
    top_k_names = ("precision_5", "recall_5", "ndcg_5", "f1_5")
    if dataset_name == "coat":
        train_mat, test_mat = load_data("coat")        
        x_train, y_train = rating_mat_to_sample(train_mat)
        x_test, y_test = rating_mat_to_sample(test_mat)
        num_user = train_mat.shape[0]
        num_item = train_mat.shape[1]

    elif dataset_name == "yahoo":
        x_train, y_train, x_test, y_test = load_data("yahoo")
        x_train, y_train = shuffle(x_train, y_train)
        num_user = x_train[:,0].max() + 1
        num_item = x_train[:,1].max() + 1

    elif dataset_name == "kuai":
        x_train, y_train, x_test, y_test = load_data("kuai", True)
        num_user = x_train[:,0].max() + 1
        num_item = x_train[:,1].max() + 1
        top_k_list = [50]
        top_k_names = ("precision_50", "recall_50", "ndcg_50", "f1_50")

    np.random.seed(2020)
    torch.manual_seed(2020)
    num_user = int(num_user)
    num_item = int(num_item)
    print("# user: {}, # item: {}".format(num_user, num_item))
    # binarize
    if dataset_name == "kuai":
        y_train = binarize(y_train, 1)
        y_test = binarize(y_test, 1)
    else:
        y_train = binarize(y_train)
        y_test = binarize(y_test)

    "MF IPS"

    mf_dr_mse = MF_DR_MSE(num_user, num_item)
    mf_dr_mse.cuda()
    # ips_idxs = np.arange(len(y_test))
    # np.random.shuffle(ips_idxs)
    # y_ips = y_test[ips_idxs[:int(0.05 * len(ips_idxs))]]
    mf_dr_mse._compute_IPS(x_train)

    mf_dr_mse.fit(x_train, y_train,
    lr=0.05,
    gamma = 0.25,
    batch_size=2048,
    lamb=1e-5,
    tol=1e-5)
    test_pred = mf_dr_mse.predict(x_test)
    mse_mf_dr_mse = mse_func(y_test, test_pred)
    auc_mf_dr_mse = roc_auc_score(y_test, test_pred)
    ndcg_res = ndcg_func(mf_dr_mse, x_test, y_test, top_k_list)
    recall_res = recall_func(mf_dr_mse, x_test, y_test, top_k_list)
    precisions = precision_func(mf_dr_mse, x_test, y_test, top_k_list)

    f1 = 2 / (1 / np.mean(precisions[top_k_names[0]]) + 1 / np.mean(recalls[top_k_names[1]]))

    print("***"*5 + "[MF]" + "***"*5)
    print("[MF] test mse:", mse_mf)
    print("[MF] test mse:", mae_mf)
    print("[MF] test auc:", auc)
    print("[MF] {}:{:.6f}".format(
            top_k_names[2].replace("_", "@"), np.mean(ndcgs[top_k_names[2]])))
    print("[MF] {}:{:.6f}".format(top_k_names[3].replace("_", "@"), f1))
    print("[MF] {}:{:.6f}".format(
            top_k_names[0].replace("_", "@"), np.mean(precisions[top_k_names[0]])))
    print("[MF] {}:{:.6f}".format(
            top_k_names[1].replace("_", "@"), np.mean(recalls[top_k_names[1]])))
    print("f1:",f1)
    print("***"*5 + "[MF-DR_MSE]" + "***"*5)


def para(args):
    if args.dataset=="coat":
        args.train_args = {"batch_size":128}
        args.model_args = {"lr":0.05, "lamb": 2e-3}
    elif args.dataset=="yahoo":
        args.train_args = {"batch_size":2048}
        args.model_args = {"lr":0.01, "lamb": 7e-5}
    elif args.dataset=="kuai":
        args.train_args = {"batch_size":2048}
        args.model_args = {"lr":0.05, "lamb": 1e-7}
    return args

if __name__ == "__main__":
    args = arguments.parse_args()
    para(args=args)

    train_and_eval(args.dataset, args.train_args, args.model_args)