import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np

import re
import os
import json

def process_log_clean(log_file):
    memorys = []
    poison_labels = []
    with open(log_file, 'r', encoding='utf-8') as file:
        log_data = file.read()
    # 1. 分隔每个以 "**********************************" 为标志的块
    blocks = log_data.split('**********************************')
    for i, block in enumerate(blocks, 1):
        if i == 0 or i%2==1:
            continue

        pattern = re.compile(
            r'''The workflow generated for the problem is.*?(\[\{.*?\}\])''',
            re.DOTALL
        )

        # 查找所有匹配项
        matches = pattern.findall(block)
        # print(f"Block {i} - Number of matches: {len(matches)}")
        # 打印提取到的 JSON 部分（即工作流内容）
        generated_plan = matches[0]
        print(generated_plan)
        memorys.append(generated_plan)
        poison_labels.append(False)

    return memorys

def parse_plans(plans_json_strs):
    """
    将工作流的 JSON 字符串解析为文本形式。

    参数：
    - plans_json_strs: 工作流的 JSON 字符串列表。

    返回：
    - plans_texts: 解析后的工作流文本列表。
    """
    plans_texts = []
    for json_str in plans_json_strs:
        # 将单引号替换为双引号，以符合 JSON 格式
        json_str_fixed = json_str.replace("'", '"')
        try:
            plan_list = json.loads(json_str_fixed)
            # 提取每个步骤的 'message' 字段，并连接成文本
            plan_text = ' '.join([step['message'] for step in plan_list])
            plans_texts.append(plan_text)
        except json.JSONDecodeError as e:
            print(f'Error decoding JSON: {e}')
            continue

    print(f'Total plans converted to text: {len(plans_texts)}')
    return plans_texts

class PerplexityFilter:
    """
    使用 Hugging Face 模型计算文本的困惑度（Perplexity）。

    参数：
    - model: 已加载的语言模型。
    - tokenizer: 与模型对应的 tokenizer。
    - threshold: 困惑度阈值。
    - window_size: 滑动窗口大小，默认为 'all'。

    方法：
    - detect(sequence): 判断文本是否通过困惑度过滤器。
    - get_log_prob(sequence): 获取文本的对数概率（NLL）。
    - filter(sequences): 对一组文本计算 NLL，并判断是否通过过滤器。
    - filter_window(sequences): 使用滑动窗口对文本进行过滤。
    """
    def __init__(self, model, tokenizer, threshold, window_size='all'):
        self.tokenizer = tokenizer
        self.model = model.cuda()  # 确保模型在 GPU 上运行
        self.threshold = threshold
        self.window_threshold = threshold
        self.window_size = window_size
        self.cn_loss = torch.nn.CrossEntropyLoss(reduction='none')

    def detect(self, sequence):
        if self.window_size == 'all':
            return (not self.filter([sequence])[-1][0])
        elif self.window_size <= 0 or type(self.window_size) != int:
            raise ValueError(f"ERROR: window_size={self.window_size}. window_size must be a positive integer.")
        return (not self.filter_window([sequence])[-1][0])

    def get_log_prob(self, sequence):
        input_ids = self.tokenizer.encode(sequence, return_tensors='pt').cuda()
        with torch.no_grad():
            logits = self.model(input_ids, labels=input_ids, use_cache=False).logits
        logits = logits[:, :-1, :].contiguous()
        input_ids = input_ids[:, 1:].contiguous()
        log_probs = self.cn_loss(logits.view(-1, logits.size(-1)), input_ids.view(-1))
        return log_probs

    def filter(self, sequences):
        filtered_log_ppl = []
        passed_filter = []
        for sequence in sequences:
            log_probs = self.get_log_prob(sequence)
            NLL_by_token = log_probs
            if NLL_by_token.mean() <= self.threshold:
                passed_filter.append(True)
                filtered_log_ppl.append(NLL_by_token.mean().item())
            else:
                passed_filter.append(False)
                filtered_log_ppl.append(NLL_by_token.mean().item())
        # print(filtered_log_ppl, passed_filter)
        return filtered_log_ppl, passed_filter

    def filter_window(self, sequences, reverse=False):
        filtered_log_ppl_by_window = []
        passed_filter_by_window = []
        passed = []
        for sequence in sequences:
            sequence_window_scores = []
            passed_window_filter = []
            log_probs = self.get_log_prob(sequence)
            NLL_by_token = log_probs
            for i in np.arange(0, len(NLL_by_token), self.window_size):
                if not reverse:
                    window = NLL_by_token[i:i+self.window_size]
                else:
                    if i == 0:
                        window = NLL_by_token[-self.window_size:]
                    elif -(-i-self.window_size) > len(NLL_by_token) and i != 0:
                        window = NLL_by_token[:-i]
                    else:
                        window = NLL_by_token[-i-self.window_size:-i]
                if window.mean() <= self.window_threshold:
                    passed_window_filter.append(True)
                    sequence_window_scores.append(window.mean().item())
                else:
                    passed_window_filter.append(False)
                    sequence_window_scores.append(window.mean().item())
            if all(passed_window_filter):
                passed.append(True)
            else:
                passed.append(False)
            passed_filter_by_window.append(passed_window_filter)
            filtered_log_ppl_by_window.append(sequence_window_scores)
        return filtered_log_ppl_by_window, passed_filter_by_window, passed

def calculate_perplexities(plans_texts, model, tokenizer, perplexity_threshold, window_size):
    """
    对给定的文本列表计算困惑度，并返回每个文本的困惑度和平均困惑度。

    参数：
    - plans_texts: 文本列表。
    - model: 已加载的语言模型。
    - tokenizer: 与模型对应的 tokenizer。
    - perplexity_threshold: 困惑度阈值。
    - window_size: 滑动窗口大小。

    返回：
    - perplexities: 每个文本的困惑度列表。
    - average_perplexity: 所有文本的平均困惑度。
    """
    # 实例化 PerplexityFilter
    filter = PerplexityFilter(model=model, tokenizer=tokenizer, threshold=perplexity_threshold, window_size=window_size)
    perplexities = []

    for plan_text in plans_texts:
        # 计算困惑度
        log_ppl, passed_filter_list = filter.filter([plan_text])
        nll = log_ppl[0]
        # 将 NLL 转换为困惑度
        perplexity = torch.exp(torch.tensor(nll)).item()
        perplexities.append(perplexity)
        # print(f"Perplexity for plan: {perplexity}")

    # 计算平均困惑度
    average_perplexity = sum(perplexities) / len(perplexities) if perplexities else 0.0
    print(f'Average Perplexity over plans: {average_perplexity}')
    return perplexities, average_perplexity

def main():
    # 使用日志文件路径
    log_base = 'logs/clean/gpt-4o-2024-08-06/no_memory/combined_attack-non-aggressive_.log'
    # 提取工作流 JSON 字符串
    plans_json_strs = process_log_clean(log_base)
    # 将工作流解析为文本
    # plans_texts = parse_plans(plans_json_strs)
    # 加载模型和 tokenizer
    model_name = 'facebook/opt-2.7b'
    # model_name = 'EleutherAI/gpt-neo-1.3B'

    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name, clean_up_tokenization_spaces=False)

    # 定义 perplexity filter 的阈值和窗口大小
    perplexity_threshold = 2.0  # 可根据任务调整
    window_size = 10  # 可选，用于滑动窗口过滤，或者使用 'all'

    # 计算 perplexities
    perplexities, average_perplexity = calculate_perplexities(plans_json_strs, model, tokenizer, perplexity_threshold, window_size)
    # 将困惑度数据排序
    ppl_values_sorted = sorted(perplexities)
    print(f"Perplexity values: {ppl_values_sorted}")
    # 计算第99百分位的值
    threshold_ppl = np.percentile(ppl_values_sorted, 99)

    print(f"FPR为1%的困惑度阈值是: {threshold_ppl}")
    # print(f"Average Perplexity: {average_perplexity}")
if __name__ == "__main__":
    main()
