import math
from rouge import Rouge
from nltk.translate.bleu_score import sentence_bleu
import warnings

warnings.filterwarnings("ignore")  # 忽略警告信息，避免输出干扰


class LLMAssessment:
    """语言模型评估类，包含困惑度、ROUGE和BLEU分数的计算方法"""

    @staticmethod
    def calculate_perplexity(sentences, unigram):
        """
        计算语料库的平均困惑度（Perplexity）
        :param sentences: 语料库，列表形式，每个元素是一个句子的单词列表
        :param unigram: 一元模型的单词概率字典，key为单词，value为概率
        :return: 平均困惑度
        """
        perplexity = 0
        for sentence in sentences:
            sentence_prob = 1
            for word in sentence:
                if word not in unigram:
                    raise ValueError(f"单词 '{word}' 不在一元模型中")
                sentence_prob *= unigram[word]
            temp = -math.log(sentence_prob, 2) / len(sentence)
            perplexity += 2 ** temp
        return perplexity / len(sentences)

    @staticmethod
    def calculate_rouge(generated_text, reference_text):
        """
        计算生成文本与参考文本之间的ROUGE分数
        :param generated_text: 模型生成的文本（字符串）
        :param reference_text: 参考文本（字符串）
        :return: ROUGE分数字典，包含ROUGE-1, ROUGE-2, ROUGE-L的precision、recall和F1分数
        """
        rouge = Rouge()
        scores = rouge.get_scores(generated_text, reference_text)
        return scores[0]

    @staticmethod
    def calculate_bleu(reference_texts, candidate_text):
        """
        计算候选句子与参考句子之间的BLEU分数（1-gram到4-gram）
        :param reference_texts: 参考文本列表（list of list of str），每个参考文本是分词后的列表
        :param candidate_text: 生成文本（list of str），分词后的单词列表
        :return: 1-gram, 2-gram, 3-gram, 4-gram的BLEU分数
        """
        bleu_1_gram = sentence_bleu(reference_texts, candidate_text, weights=(1, 0, 0, 0))
        bleu_2_gram = sentence_bleu(reference_texts, candidate_text, weights=(0.5, 0.5, 0, 0))
        bleu_3_gram = sentence_bleu(reference_texts, candidate_text, weights=(0.33, 0.33, 0.33, 0))
        bleu_4_gram = sentence_bleu(reference_texts, candidate_text, weights=(0.25, 0.25, 0.25, 0.25))
        return bleu_1_gram, bleu_2_gram, bleu_3_gram, bleu_4_gram


if __name__ == "__main__":
    # 示例：困惑度计算
    sentences = [
        ['I', 'have', 'a', 'pen'],
        ['He', 'has', 'a', 'book'],
        ['She', 'has', 'a', 'cat']
    ]
    unigram = {
        'I': 1 / 12, 'have': 1 / 12, 'a': 3 / 12, 'pen': 1 / 12,
        'He': 1 / 12, 'has': 2 / 12, 'book': 1 / 12, 'She': 1 / 12, 'cat': 1 / 12
    }
    perplexity = LLMAssessment.calculate_perplexity(sentences, unigram)
    print("困惑度为：", perplexity)

    # 示例：ROUGE分数计算
    generated_text = "This is some generated text."
    reference_text = "This is a reference text."
    rouge_scores = LLMAssessment.calculate_rouge(generated_text, reference_text)
    print("ROUGE分数：", rouge_scores)

    # 示例：BLEU分数计算
    candidate_text = ["This", "is", "some", "generated", "text"]
    reference_texts = [
        ["This", "is", "a", "reference", "text"],
        ["This", "is", "another", "reference", "text"]
    ]
    bleu_scores = LLMAssessment.calculate_bleu(reference_texts, candidate_text)
    print("BLEU分数：", bleu_scores)
