import os
import json

def analyze_sts_results(folder_path):
    """
    Read all STS task .json files in the specified folder, extract Spearman correlation coefficients, 
    and calculate total and average scores.

    This version will prioritize looking for the "validation" field, and if not found, 
    it will look for the "test" field.

    :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")

                # First try to extract score from "validation"
                spearman_score = data.get("validation", {}).get("cos_sim", {}).get("spearman")
                # If no value is obtained from "validation" (i.e., spearman_score is None), try to extract from "test"
                if spearman_score is None:
                    spearman_score = data.get("test", {}).get("cos_sim", {}).get("spearman")

                assert (dataset_name and spearman_score is not None), \
                    f"Warning: File '{filename}' does not contain 'validation/cos_sim/spearman' or 'test/cos_sim/spearman' field."

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

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

    # Note: total and avg here are calculated based on percentage scores
    total_score_percentage = sum(scores)
    avg_score_percentage = total_score_percentage / len(scores)

    print(f"\ntotal: {total_score_percentage:.2f}")
    print(f"avg: {avg_score_percentage:.2f}")


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


if __name__ == '__main__':
    main()
