""" Authors: Shatian Wang, Eunjin (Jinny) Cho, and JordiKai Watanabe-Inouye
This script takes the prediction dictionaries generated by kNN (euclidean, cosine, adjusted cosine & normalized, not normalized)
and outputs the rmse values for each dictionary. Recall that rmse stands for root mean square error. Thereby this value will
tell us how off our predictions were from the actual rating """

import math
import time
import pickle
from create_util_mat import UtilMat
from collections import defaultdict


def get_rmse_single_bin(dev_dict, predicted_dict, nonzero_training_indices, user_index_to_id_dict):
    """ Constructs a dictionary containing the information to compute the rmse value for each "bin".
    A bin is defined as the number of ratings a user has given a book per the trainging matrix.
    The information stored in dict:
        - number of books
        - sum of all square differences between actual and prediction
    :param dev_dict: dictionary that maps a user to a dictionary of book:rating
        represents the actual data given by AB/BX
    :param predicted_dict: dictionary that maps a user to a dictionary of book:rating
        represents the data predicted based on the training data & k
    :param nonzero_training_indices: a dictionary, whose keys is the user_id and value is a list of book_ids that have been rated
    :return rmse_single_bin:
    (P329) In order to compute the rmse value we must sum over
    the square difference between the actual and prediction for all nonzero entries """
    rmse_single_bin = defaultdict()
    num_ratings_to_user_dict = defaultdict()
    # creates a dictionary that maps num_books_rated to list of users, based on training dataset
    for user in nonzero_training_indices.keys():
        num_books_rated = len(nonzero_training_indices[user])
        if num_books_rated in num_ratings_to_user_dict.keys():
            num_ratings_to_user_dict[num_books_rated].append(user)
        else:
            num_ratings_to_user_dict[num_books_rated] = [user]
    # using this data we can then calculate the sum of the differences,
    # for users that have all rated the smae number of books
    for num_ratings in num_ratings_to_user_dict.keys():
        user_id_list = num_ratings_to_user_dict[num_ratings]
        sum_sq_diff_sum = 0.0
        sum_books = 0
        for user in user_id_list:
            sq_diff_sum = 0.0
            counter = 0
            # gets square difference of actual and prediction, based on dev dataset
            for book in dev_dict[user]:
                actual = dev_dict[user][book]
                prediction = predicted_dict[user][book]
                user_id = user_index_to_id_dict[user]
                if user_id[0] != 'A':
                    # if a BX user, convert rating to 1-5 scale
                    actual = actual/2.5 + 1

                sq_diff_sum += (actual-prediction)**2
                counter += 1

            sum_books += len(dev_dict[user].keys())
            sum_sq_diff_sum += sq_diff_sum
            #accumulate in dictionary the data
            if num_ratings in rmse_single_bin.keys():
                #changed num_books_rated to num_ratings
                rmse_single_bin[num_ratings][0] += sum_sq_diff_sum
                rmse_single_bin[num_ratings][1] += sum_books
            else:
                rmse_single_bin[num_ratings] = [sum_sq_diff_sum,sum_books]

    return rmse_single_bin


def get_rmse(rmse_single_bin):
    """
    :param rmse_single_bin: a dictionary containing the information to compute the rmse value for each "bin"
        it may be the case that the range of the bins is not consistent, size of bin is a test var
    :return rmse_dict, cummulative: the rmse value for a given "bin" and the cummulative rmse value for a given k
    """
    rmse_dict = defaultdict()
    sum_rmse = 0.0
    sum_num_ratings = 0
    for n in rmse_single_bin.keys():
        if rmse_single_bin[n][0] != 0.0:
            rmse_dict[n] = math.sqrt(rmse_single_bin[n][0]/rmse_single_bin[n][1])
            sum_rmse += rmse_single_bin[n][0]
            sum_num_ratings += rmse_single_bin[n][1]
    return rmse_dict, math.sqrt(sum_rmse/sum_num_ratings)


def calculate_rmse(dev_dict, nonzero_training_indices, user_index_to_id_dict, cumulative_rmse_dict, k):
    """
    For Development:
    Normalized, should be saved to the directory "dev_norm_pred_rmse"
    Unormalized, should be saved to the directory "dev_unnorm_pred_rmse"
    Each pair of similairity, neighbors has a single_bin_rmse dictionary
    and a rmse_single_bin dictionary

    For Evaluation:
    Should be saved to the directory "eval_norm_pred_rmse" or "eval_unnorm_pred_rmse"
    Instead of a given 'k' neighbor, we only calculate the rmse value for the 'best point'
    determined by tuning our variables in development
    """
    if k == "eval":
        #TODO: change to match best point
        filename = "eval_unnorm_pred/prediction_dict_euclidean_30.p"
        prediction_dict = pickle.load(open(filename, "rb"))
        rmse_single_bin = get_rmse_single_bin(dev_dict, prediction_dict, nonzero_training_indices, user_index_to_id_dict)
        single_bin_rmse, cumulative_rmse = get_rmse(rmse_single_bin)
        #TODO: name match best point
        pickle.dump(single_bin_rmse, open("eval_unnorm_pred/single_bin_rmse.p", "wb"))
        pickle.dump(cumulative_rmse, open("eval_unnorm_pred/cummulative_rmse.p", "wb"))
        return

    start_time = time.time()
    ## Development
    # euc_filename = "dev_unnorm_pred/prediction_dict_euclidean_" + str(k) + ".p"
    # cos_filename = "dev_unnorm_pred/prediction_dict_cosine_" + str(k) + ".p"
    # adj_cos_filename = "dev_unnorm_pred/prediction_dict_adj_cosine_" + str(k) + ".p"
    euc_filename = "dev_norm_pred/prediction_dict_euclidean_" + str(k) + ".p"
    cos_filename = "dev_norm_pred/prediction_dict_cosine_" + str(k) + ".p"
    adj_cos_filename = "dev_norm_pred/prediction_dict_adj_cosine_" + str(k) + ".p"

    prediction_dict_euclidean = pickle.load(open(euc_filename, "rb"))
    prediction_dict_cosine = pickle.load(open(cos_filename, "rb"))
    prediction_dict_adj_cosine = pickle.load(open(adj_cos_filename, "rb"))

    rmse_single_bin_euc = get_rmse_single_bin(dev_dict, prediction_dict_euclidean, nonzero_training_indices, user_index_to_id_dict)
    rmse_single_bin_cos = get_rmse_single_bin(dev_dict, prediction_dict_cosine, nonzero_training_indices, user_index_to_id_dict)
    rmse_single_bin_adj_cos = get_rmse_single_bin(dev_dict, prediction_dict_adj_cosine, nonzero_training_indices, user_index_to_id_dict)

    # euc_rmse_filename = "dev_unnorm_pred_rmse/rmse_single_bin_euc_" + str(k) + ".p"
    # cos_rmse_filename = "dev_unnorm_pred_rmse/rmse_single_bin_cos_" + str(k) + ".p"
    # adj_cos_rmse_filename = "dev_unnorm_pred_rmse/rmse_single_bin_adj_cos_" + str(k) + ".p"
    euc_rmse_filename = "dev_norm_pred_rmse/rmse_single_bin_euc_" + str(k) + ".p"
    cos_rmse_filename = "dev_norm_pred_rmse/rmse_single_bin_cos_" + str(k) + ".p"
    adj_cos_rmse_filename = "dev_norm_pred_rmse/rmse_single_bin_adj_cos_" + str(k) + ".p"

    pickle.dump(rmse_single_bin_euc, open(euc_rmse_filename, "wb"))
    pickle.dump(rmse_single_bin_cos, open(cos_rmse_filename, "wb"))
    pickle.dump(rmse_single_bin_adj_cos, open(adj_cos_rmse_filename, "wb"))

    single_bin_rmse_euc, cumulative_rmse_euc = get_rmse(rmse_single_bin_euc)
    single_bin_rmse_cos, cumulative_rmse_cos = get_rmse(rmse_single_bin_cos)
    single_bin_rmse_adj_cos, cumulative_rmse_adj_cos = get_rmse(rmse_single_bin_adj_cos)

    cumulative_rmse_dict[k] = [cumulative_rmse_euc, cumulative_rmse_cos, cumulative_rmse_adj_cos]

    # euc_single_rmse_filename = "dev_unnorm_pred_rmse/single_bin_rmse_euc_" + str(k) + ".p"
    # cos_single_rmse_filename = "dev_unnorm_pred_rmse/single_bin_rmse_cos_" + str(k) + ".p"
    # adj_cos_single_rmse_filename = "dev_unnorm_pred_rmse/single_bin_rmse_adj_cos_" + str(k) + ".p"
    euc_single_rmse_filename = "dev_norm_pred_rmse/single_bin_rmse_euc_" + str(k) + ".p"
    cos_single_rmse_filename = "dev_norm_pred_rmse/single_bin_rmse_cos_" + str(k) + ".p"
    adj_cos_single_rmse_filename = "dev_norm_pred_rmse/single_bin_rmse_adj_cos_" + str(k) + ".p"

    pickle.dump(single_bin_rmse_euc, open(euc_single_rmse_filename, "wb"))
    pickle.dump(single_bin_rmse_cos, open(cos_single_rmse_filename, "wb"))
    pickle.dump(single_bin_rmse_adj_cos, open(adj_cos_single_rmse_filename, "wb"))

    # print(euc_filename, cos_filename, adj_cos_filename)
    print("cumulative rmse: ", cumulative_rmse_euc, cumulative_rmse_cos, cumulative_rmse_adj_cos)
    print("-------- calculated rmse --------")
    print("--- RMSE Execution Time ---\n %s " % (time.time() - start_time))


def main():
    # Load data structures
    util_mat_obj = pickle.load(open("datafile/UtilMat_obj.p", "rb"))
    user_index_to_id_dict = util_mat_obj.index_to_user_dict

    # Load dev_dict
    dev_dict = pickle.load(open("datafile/dev_dict.p", "rb"))
    # dev_dict = pickle.load(open("datafile/eval_dict.p", "rb"))

    # Track nonzero entries
    nonzero_training_indices = pickle.load(open("datafile/nonzero_training_indices.p", "rb"))

    cumulative_rmse_dict = {}
    ## Run rmse for dev
    for k in range(2, 31):
        calculate_rmse(dev_dict, nonzero_training_indices, user_index_to_id_dict, cumulative_rmse_dict, k)

    # Saving dictionary
    # pickle.dump(cumulative_rmse_dict, open("dev_unnorm_pred_rmse/dev_unnorm_cumulative_rmse_dict.p", "wb"))
    pickle.dump(cumulative_rmse_dict, open("dev_norm_pred_rmse/dev_norm_cumulative_rmse_dict.p", "wb"))

    ## Run rmse for eval
    # k = "eval"
    # calculate_rmse(dev_dict, nonzero_training_indices, user_index_to_id_dict, cumulative_rmse_dict, k)


if __name__ == "__main__": main()
