import pandas as pd
import numpy as np
import json

def calculate_metrics_all(result_json_path):
    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 = 0
    total_labels = 0
    with open(result_json_path, "r", encoding="utf-8") as f:
        for line in f:
            num_test_data += 1
            local_correct_prediction = 0

            d = json.loads(line)
            labels = d["tags"]
            # get only the top k prediction
            prediction = d["predict"][0:3]
            total_labels += len(labels)
            correct_prediction = 0
            # K = 1
            if prediction[0][0] in labels:
                correct_prediction += 1
            sum_precision_1 += (correct_prediction / 1)
            sum_recall_1 += (correct_prediction / len(labels))

            # K = 2
            if prediction[1][0] in labels:
                correct_prediction += 1
            sum_precision_2 += (correct_prediction / 2)
            sum_recall_2 += (correct_prediction / len(labels))

            # K = 3
            if prediction[2][0] in labels:
                correct_prediction += 1
            sum_precision_3 += (correct_prediction / 3)
            sum_recall_3 += (correct_prediction / len(labels))


        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__())
        print("TOTAL LABELS: " + total_labels.__str__())

calculate_metrics_all("inference_result.json")
