import json
import argparse
from nltk.translate.bleu_score import sentence_bleu
import re
import nltk
import sacrebleu
bleu_metric = sacrebleu.BLEU(tokenize='13a')
from scipy.stats import spearmanr, pearsonr
import numpy as np
from tqdm import tqdm
# 确保你已经下载了 NLTK 所需的数据包
def extract_numbers_and_floats(text):
    # 使用正则表达式找出文本中的所有整数和浮点数
    numbers = re.findall(r'\d+\.\d+|\d+', text)
    # 将提取的数字转换为整数或浮点数
    return float(numbers[0])


import transformers
import torch
import json

model_id = "meta-llama/Meta-Llama-3.1-70B-Instruct"

pipeline = transformers.pipeline(
    "text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, max_new_tokens=150, device_map="auto"
)

prompt = "{}\n" \
         "According to the context, please judge if SpeechA is better or SpeechB is better. Only output '[SpeechA]' or '[SpeechB]', do not write any analysis."


error = []
pre_score, ref_score = [], []
def normalize_and_calculate_bleu(jsonl_file):
    pre, gt = [], []
    total, correct = 0, 0
    with open(jsonl_file, 'r') as f:
        for line in tqdm(f):
            data = json.loads(line.strip())
            candidate = data['response']
            reference = data['label']

            if len(candidate) == 0 or len(reference) == 0:
                continue

            messages = [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt.format(candidate)}
            ]
            outputs1 = pipeline(
                messages,
                max_new_tokens=256,
            )
            output1 = outputs1[0]["generated_text"][-1]['content']

            messages = [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt.format(reference)}
            ]
            outputs2 = pipeline(
                messages,
                max_new_tokens=256,
            )
            output2 = outputs2[0]["generated_text"][-1]['content']

            total += 1
            if output1 == output2:
                correct += 1

            pre.append(candidate)
            gt.append(reference)  # BLEU 需要参考句子是列表形式

            candidate_score = re.findall(r'\d+\.\d+|\d+', candidate)
            reference_score = re.findall(r'\d+\.\d+|\d+', reference)

            if len(candidate_score) !=1 or len(reference_score) !=1:
                continue
            pre_score.append(float(candidate_score[0]))
            ref_score.append(float(reference_score[0]))
            error.append((float(candidate_score[0])-float(reference_score[0]))**2)


    # 计算每个句对的 BLEU 分数并取平均
    bleu_score = bleu_metric.corpus_score(pre, [gt]).score

    print(correct / total)
    return bleu_score


def main():
    parser = argparse.ArgumentParser(description="Calculate BLEU for normalized text.")
    parser.add_argument("jsonl_file", type=str, help="Path to the JSONL file containing the results to be evaluated.")

    args = parser.parse_args()

    bleu_result = normalize_and_calculate_bleu(args.jsonl_file)


    # print('mse is:', sum(error)/ len(error))
    # score_pearsonr, _ = pearsonr(pre_score, ref_score)
    # print('pearsonr is:', score_pearsonr)
    # corr, _ = spearmanr(pre_score, ref_score)
    # print('SRCC is:', corr)
    print(f"The average BLEU score is: {bleu_result}")


if __name__ == "__main__":
    main()