import os
import json


def analyze_retrieval_results(folder_path):
    """
    Read all .json files in the specified folder, extract information and calculate total and average scores.
    :param folder_path: Path to the folder containing .json files
    """
    scores = []
    results = []

    assert os.path.isdir(folder_path), f"Error: Folder '{folder_path}' does not exist."

    for filename in sorted(os.listdir(folder_path)):
        if not filename.endswith('.json'):
            continue
        file_path = os.path.join(folder_path, filename)
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)

                dataset_name = data.get("mteb_dataset_name")
                # Check if 'dev' and 'ndcg_at_10' exist
                ndcg_score = data.get("dev", {}).get("ndcg_at_10")

                assert dataset_name and ndcg_score is not None, \
                    f"Warning: File '{filename}' is missing 'mteb_dataset_name' or 'ndcg_at_10' field."

                scores.append(ndcg_score)
                percentage_score = round(ndcg_score * 100, 2)
                results.append(f"{dataset_name} : {ndcg_score} => {percentage_score:.2f}")

        except json.JSONDecodeError:
            print(f"Warning: Unable to parse file '{filename}'.")

    assert len(scores) > 0, "No valid data found for statistics."
    total_score = sum(scores) * 100
    avg_score = total_score / len(scores)
    print(f"\ntotal: {total_score:.2f}")
    print(f"avg: {avg_score:.2f}")


def main():
    # --- Execute analysis ---
    folder_path = './IR/model_MiniCPM-Embedding'
    print(f"--- Analyzing folder: {folder_path} ---")
    analyze_retrieval_results(folder_path=folder_path)


if __name__ == '__main__':
    main()
