import pandas as pd
from tqdm import tqdm
import torch
import os
# 设置环境变量禁用 TensorFlow
os.environ["USE_TORCH"] = "1"
os.environ["USE_TF"] = "0"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 设置设备
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
print(f"使用设备: {device}")

# 基础模型路径
base_model_path = "../Qwen2.5-7B-Instruct"
# 微调模型路径
finetuned_model_path = "../qwen-medical"  # 修改为您的微调模型路径

# 加载基础模型
base_tokenizer = AutoTokenizer.from_pretrained(base_model_path, trust_remote_code=True)
base_model = AutoModelForCausalLM.from_pretrained(
        base_model_path,
        torch_dtype=torch.float16,
        trust_remote_code=True
    ).to(device)

# 加载微调后模型
finetuned_tokenizer = AutoTokenizer.from_pretrained(finetuned_model_path, trust_remote_code=True)
finetuned_model = AutoModelForCausalLM.from_pretrained(
        finetuned_model_path,
        torch_dtype=torch.float16,
        trust_remote_code=True
    ).to(device)
print("模型加载完成")

# 准备评估数据
# 创建医疗就诊测试数据
eval_dataset = [
    {
        "Question": "我最近经常感到胸闷气短，特别是在运动后，这可能是什么问题？",
        "Answer": "您描述的症状可能与心肺功能有关。建议及时就医检查，进行心电图、胸片等相关检查以明确诊断。同时请注意：1. 暂时减少剧烈运动 2. 保持规律作息 3. 记录症状发生的具体情况"
    },
    {
        "Question": "孩子发烧38.5度，已经持续一天了，除了发烧没有其他明显症状，该怎么处理？",
        "Answer": "对于儿童发烧，建议以下措施：1. 多补充水分 2. 物理降温（温水擦浴） 3. 必要时使用退烧药 4. 如果发烧持续超过48小时或出现其他症状及时就医 5. 注意观察精神状态和体温变化"
    },
    {
        "Question": "男，十二岁。右股骨下端疼痛2个月，夜间加重。查体：右股骨下端局部肿胀、皮温增高。X线示右胫上端日光射线样改变。首先考虑的诊断是\nA.骨巨细胞瘤  \nB.慢性骨髓炎\nC.股骨结节骨软骨病 \nD.骨样骨瘤 \nE. 骨肉瘤",
        "Answer": "基于提供的症状和影像学表现，12岁男孩右股骨下端的疼痛具有夜间加重的特点，并伴有局部肿胀和温度升高，X光片显示日光射线样改变，这些都非常符合骨肉瘤的特征。\n\n虽然选项中还有其他可能性，但结合年龄、症状以及特异性的X光表现，最可能的诊断是 E. 骨肉瘤。"
    },
    {
        "Question": "经常性反酸打嗝，体重下降，腹部CT完全正常、肝肾功能正常，癌胚抗原指标正常，这种情况下最可能的诊断是什么",
        "Answer": "根据提供的症状和检查结果，最可能的诊断是胃食管反流病（GERD）。这种病常见于长期的反酸和打嗝，而CT检查和肝肾功能的正常则排除了其他更严重的疾病。当然，进一步的检查如胃镜可能有助于确诊和排除其他可能性。"
    },
    {
        "Question": "一名中年男性自述鼻涕中带血已经有三个月了，面且伴有有耳堵和听力减退，无发热。根据鼻咽部的症状和体征，这种情况下最可能的诊断是什么？",
        "Answer": "根据提供的症状以及病程的描述，这名男性患者可能患有鼻咽癌的可能性较高。以下是分析理由：1. **持续时间**：鼻涕中带血已持续三个月，这种情况排除了普通感冒或短期炎症的可能，指向一种更慢性或严重的病变。\n\n2. **症状组合**：鼻涕带血、耳堵及听力减退，这些症状在结合在一起时较为典型地指向鼻咽部的问题。耳堵和听力减退尤其提示可能是由于肿瘤压迫或阻塞了耳咽管。\n\n3. **无发热**：在长期的炎症或感染中，通常会伴随发热现象，而患者并未出现此症状，这降低了急性鼻窦炎等感染的可能性。\n\n鉴于这些症状的组合和其持续的时间，早期鼻咽癌的可能性较高。当然，最终的确诊还需通过影像学检查（如CT或MRI）以及病理活检来确认。不论如何，这种情况下，尽早就医进行详细检查是至关重要的，以便及时获得准确的诊断和适当的治疗。"
    }
]
eval_dataset = pd.DataFrame(eval_dataset)
test_questions = eval_dataset["Question"].tolist()
comparison_results = []

# 对每个测试样例进行对比评估
for question in tqdm(test_questions):
    # 使用微调模型生成回答
    formatted_prompt = f"下面列出了一个问题. 请写出问题的答案.### 问题:{question}### 答案:"
    
    # 微调模型推理
    finetuned_inputs = finetuned_tokenizer(formatted_prompt, return_tensors="pt").to(device)
    
    finetuned_outputs = finetuned_model.generate(
        **finetuned_inputs,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )
    
    finetuned_response = finetuned_tokenizer.decode(finetuned_outputs[0], skip_special_tokens=True)
    finetuned_response = finetuned_response.replace(formatted_prompt, "")  # 移除提示词
    
    # 基础模型推理
    base_inputs = base_tokenizer(formatted_prompt, return_tensors="pt").to(device)
    
    base_outputs = base_model.generate(
        **base_inputs,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )
    
    base_response = base_tokenizer.decode(base_outputs[0], skip_special_tokens=True)
    base_response = base_response.replace(formatted_prompt, "")  # 移除提示词
    
    # 保存结果
    comparison_results.append({
        "question": question,
        "finetuned_response": finetuned_response,
        "base_response": base_response,
        "reference_answer": eval_dataset[eval_dataset["Question"] == question]["Answer"].values[0] #type: ignore
    })

# 保存对比结果
comparison_df = pd.DataFrame(comparison_results)
comparison_df.to_csv("./model_comparison.csv", index=False)

print(f"对比评估完成，结果已保存至 model_comparison.csv")

# 简单统计
print("\n===== 对比结果统计 =====")
for i, result in enumerate(comparison_results):
    print(f"\n问题 {i+1}: {result['question']}")
    print(f"参考答案: {result['reference_answer']}")
    print(f"微调模型: {result['finetuned_response'][:100]}...")
    print(f"基础模型: {result['base_response'][:100]}...")
    print("-" * 50)