import omikuji
import pickle
import sys
from timeit import default_timer as timer

def omikuji_prediction(model, test_data):
    model = omikuji.Model.load(model)
    # Manually load the test data here
    actual_labels = []
    feature_pairs = []
    with open(test_data, "r", encoding="utf-8") as w:
        all_data = w.readlines()
        # remove the header
        all_data.pop(0)
        # loop through all the test data
        # then process each of them (i.e., split the label and create the feature value pairs)
        for data in all_data:
            splitted = data.split(" ")
            # split the label
            labels = splitted[0].split(",")
            temp_label = []
            for label in labels:
                temp_label.append(int(label))
            actual_labels.append(temp_label)
            # Done processing the label here
            # Process the feature value pair here
            temp_feature = []
            for i in range(1, len(splitted)):
                feature_split = splitted[i].rstrip().split(":")
                if len(feature_split) == 2:
                    feature_tuple = (int(feature_split[0]), float(feature_split[1]))
                    temp_feature.append(feature_tuple)
            feature_pairs.append(temp_feature)
        # Save the processed labels and feature pair using pickle for easier access
        with open("test_feature_pair.pkl", "wb") as f:
            pickle.dump(feature_pairs, f)
        with open("test_label.pkl", "wb") as f:
            pickle.dump(actual_labels, f)

        sum_recall_1 = 0
        sum_recall_2 = 0
        sum_recall_3 = 0
        sum_precision_1 = 0
        sum_precision_2 = 0
        sum_precision_3 = 0
        num_test_data = len(feature_pairs)
        # Now proceed with the prediction
        list_labels = []
        list_preds = []
        for i in range(len(feature_pairs)):
            feature_pair = feature_pairs[i]
            actual_label = actual_labels[i]

            correct_prediction = 0
            label_score_pairs = model.predict(feature_pair, top_k=3)
            list_labels.append(actual_label)
            list_preds.append([n[0] for n in label_score_pairs])

            # k = 1
            if label_score_pairs[0][0] in actual_label:
                correct_prediction += 1
            sum_precision_1 += (correct_prediction / 1)
            sum_recall_1 += (correct_prediction / len(actual_label))

            # k = 2
            if label_score_pairs[1][0] in actual_label:
                correct_prediction += 1
            sum_precision_2 += (correct_prediction / 2)
            sum_recall_2 += (correct_prediction / len(actual_label))

            # k = 3
            if label_score_pairs[2][0] in actual_label:
                correct_prediction += 1
            sum_precision_3 += (correct_prediction / 3)
            sum_recall_3 += (correct_prediction / len(actual_label))

        # Divide by numer of test data
        precision_1 = sum_precision_1 / num_test_data
        precision_2 = sum_precision_2 / num_test_data
        precision_3 = sum_precision_3 / num_test_data
        recall_1 = sum_recall_1 / num_test_data
        recall_2 = sum_recall_2 / num_test_data
        recall_3 = sum_recall_3 / num_test_data
        f1_1 = 2 * precision_1 * recall_1 / (precision_1 + recall_1)
        f1_2 = 2 * precision_2 * recall_2 / (precision_2 + recall_2)
        f1_3 = 2 * precision_3 * recall_3 / (precision_3 + recall_3)

        print("K = 1")
        print("P@1 = " + precision_1.__str__())
        print("R@1 = " + recall_1.__str__())
        print("F@1 = " + f1_1.__str__())

        print("K = 2")
        print("P@2 = " + precision_2.__str__())
        print("R@2 = " + recall_2.__str__())
        print("F@2 = " + f1_2.__str__())

        print("K = 3")
        print("P@3 = " + precision_3.__str__())
        print("R@3 = " + recall_3.__str__())
        print("F@3 = " + f1_3.__str__())



if __name__ == "__main__":
    model_path = sys.argv[1]
    test_path = sys.argv[2]
    omikuji_prediction(model_path, test_path)
