"""
法律知识图谱构建程序

使用方法：
1. 运行单个模式：
   python script.py --modes 1
   
2. 运行多个模式：
   python script.py --modes 1,2,3
   
3. nohup后台运行：
   nohup python script.py --modes 2,4 --sample_size 10 > output.log 2>&1 &
   
4. 自定义路径和参数：
   python script.py --modes 1,2 --sample_size 20 --model_path /your/model/path

模式说明：
- 模式1：单个查询的综合知识图谱（案例+法条）
- 模式2：批量构建综合知识图谱（案例+法条）  
- 模式3：单个查询的仅案例知识图谱
- 模式4：批量构建仅案例知识图谱
"""

import json
import os
from typing import Dict, List, Any
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def load_retrieval_results_by_id(folder_path):
    """按ID加载预检索结果，同时记录原始文件夹路径"""
    results_dict = {}
    folder_paths = {}  # 记录每个query_id对应的文件夹路径
    
    if not os.path.exists(folder_path):
        print(f"警告: 检索结果文件夹不存在: {folder_path}")
        return results_dict, folder_paths
        
    for query_id in os.listdir(folder_path):
        query_folder = os.path.join(folder_path, query_id)
        
        if not os.path.isdir(query_folder):
            continue
            
        # 加载原始查询
        query_file = os.path.join(query_folder, 'query.json')
        if os.path.exists(query_file):
            with open(query_file, 'r', encoding='utf-8') as f:
                query_data = json.load(f)
                actual_query_id = query_data.get('id', query_id)
            
            # 加载检索摘要
            summary_file = os.path.join(query_folder, 'search_summary.json')
            retrieved_docs = []
            
            if os.path.exists(summary_file):
                with open(summary_file, 'r', encoding='utf-8') as f:
                    summary = json.load(f)
                
                # 加载每个检索到的文档
                for item in summary:
                    doc_file = os.path.join(query_folder, f"score_{item['file_name']}")
                    if os.path.exists(doc_file):
                        with open(doc_file, 'r', encoding='utf-8') as f:
                            doc_data = json.load(f)
                            retrieved_docs.append(doc_data)
            
            # 存储结果 - 使用ID作为键，同时记录文件夹路径
            results_dict[actual_query_id] = retrieved_docs
            folder_paths[actual_query_id] = query_folder  # 记录原始文件夹路径
    
    print(f"已加载 {len(results_dict)} 条检索结果(按ID索引)")
    return results_dict, folder_paths

def load_qwen32b_model(model_path: str):
    """加载qwen32B本地模型"""
    print(f"正在加载qwen32B模型: {model_path}")
    
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True
    )
    
    print("模型加载完成!")
    return model, tokenizer

def call_qwen32b_for_knowledge_graph_with_laws(case_data_list: List[Dict], law_data_list: List[Dict], model, tokenizer, query_id: str) -> Dict:
    """
    调用qwen32B模型构建包含案例和法条的知识图谱
    """
    # 构建案例信息
    cases_text = ""
    if case_data_list:
        cases_text = "## 相关案例：\n"
        for i, case in enumerate(case_data_list[:3]):  # 取top3案例
            fact = case.get('fact', '')
            reason = case.get('reason', '')
            result = case.get('result', '')
            similarity = case.get('similarity', 0)
            
            cases_text += f"""
### 案例{i+1}（相似度：{similarity:.3f}）：
**案件事实：** {fact}
**法院判决：** {reason}
**判决结果：** {result}
"""
    
    # 构建法条信息 - 修正为实际的数据结构
    laws_text = ""
    if law_data_list:
        laws_text = "## 相关法条：\n"
        for i, law in enumerate(law_data_list[:3]):  # 🎯 改为取top3法条
            law_id = law.get('id', '')                    # ✅ 法条ID
            law_content = law.get('fact', '')             # ✅ 法条内容在fact字段
            similarity = law.get('similarity', 0)         # ✅ 相似度
            rank = law.get('rank', i+1)                   # ✅ 排名
            
            laws_text += f"""
### 法条{i+1}（ID：{law_id}，相似度：{similarity:.3f}，排名：{rank}）：
**法条内容：** {law_content}
"""
    
    # 构建综合prompt
    prompt = f"""你是专业的刑法知识图谱构建专家。请基于提供的相关案例和法条，构建一个**四层关联的综合知识图谱**。这不是简单的信息抽取，而是要深度分析案例与法条之间的对应关系，构建完整的法律推理链条。

## 核心建模原则：
1. **案例与法条的对应映射**：找出每个案例事实对应的具体法条条文
2. **法律推理链条完整性**：从事实→法条依据→法院认定→最终判决的完整链条
3. **多案例法条融合**：整合多个相似案例和相关法条，形成统一的知识结构
4. **层级关系明确性**：不同层次实体间的逻辑关系要清晰准确

## 四层建模框架：

### 🏛️ 法条层（法理基础）：
**核心实体类型**：
- 【法条条文】：具体的法律条款（如"刑法第XXX条"）
- 【犯罪构成要件】：法条规定的构成犯罪的必要条件
- 【法定刑罚幅度】：法条规定的刑罚范围和标准
- 【适用条件】：法条的适用范围和前提条件
- 【量刑情节】：法条规定的加重、减轻情节

### 📋 事实层（客观事实）：
**核心实体类型**：
- 【犯罪行为】：具体的违法行为和手段
- 【行为后果】：犯罪造成的损失和影响
- 【时空要素】：犯罪的时间、地点、环境
- 【物证材料】：涉案物品、证据、检验结果
- 【主体身份】：当事人的身份、关系、状态

### ⚖️ 推理层（法律认定）：
**核心实体类型**：
- 【罪名认定】：法院确定的具体罪名
- 【构成要件匹配】：事实与法条构成要件的对应关系
- 【量刑情节认定】：法院认定的具体量刑情节
- 【法院观点】：法院的法律评价和推理过程
- 【适用法条确定】：法院选择适用的具体法条

### 📜 判决层（裁判结果）：
**核心实体类型**：
- 【确定罪名】：最终认定的罪名
- 【具体刑期】：确定的刑罚期限
- 【执行方式】：刑罚执行的具体方式
- 【附加刑罚】：罚金、没收财产等附加刑
- 【程序事项】：羁押折抵、缓刑考验等

## 知识图谱构建要求：

### 第一步：实体识别与层级分类
请按照以下格式输出四层实体：

**法条层实体：**
- 【刑法第XXX条】（属性：【条文内容：..., 适用对象：..., 刑罚范围：...】）
- 【XXX罪构成要件】（属性：【主观要件：..., 客观要件：..., 主体要件：...】）

**事实层实体：**
- 【具体犯罪行为】（属性：【行为方式：..., 行为时间：..., 涉案金额：...】）
- 【被告人】（属性：【身份特征：..., 前科情况：..., 认罪态度：...】）

**推理层实体：**
- 【XXX罪认定】（属性：【认定依据：..., 构成要件满足：..., 量刑情节：...】）
- 【法院法律评价】（属性：【推理过程：..., 适用法条：..., 量刑理由：...】）

**判决层实体：**
- 【有期徒刑X年】（属性：【刑期计算：..., 执行地点：..., 特殊规定：...】）
- 【罚金X万元】（属性：【缴纳期限：..., 缴纳方式：..., 执行措施：...】）

### 第二步：层内关系构建
**法条层关系：**
1. 【刑法第XXX条】 -> 【规定】 -> 【XXX罪构成要件】
2. 【XXX罪构成要件】 -> 【决定】 -> 【法定刑罚幅度】

**事实层关系：**
1. 【被告人】 -> 【实施】 -> 【具体犯罪行为】
2. 【具体犯罪行为】 -> 【造成】 -> 【行为后果】

**推理层关系：**
1. 【构成要件匹配】 -> 【支撑】 -> 【罪名认定】
2. 【法院观点】 -> 【基于】 -> 【适用法条确定】

**判决层关系：**
1. 【确定罪名】 -> 【对应】 -> 【具体刑期】
2. 【具体刑期】 -> 【通过】 -> 【执行方式】

### 第三步：跨层映射关系（核心）
**法条→事实映射：**
1. 【XXX罪构成要件】 -> 【匹配】 -> 【具体犯罪行为】
2. 【适用条件】 -> 【符合】 -> 【被告人身份】

**事实→推理映射：**
1. 【具体犯罪行为】 -> 【证明】 -> 【构成要件匹配】
2. 【行为后果】 -> 【影响】 -> 【量刑情节认定】

**法条→推理映射：**
1. 【刑法第XXX条】 -> 【作为依据】 -> 【罪名认定】
2. 【法定刑罚幅度】 -> 【指导】 -> 【法院观点】

**推理→判决映射：**
1. 【罪名认定】 -> 【确定】 -> 【确定罪名】
2. 【量刑情节认定】 -> 【影响】 -> 【具体刑期】

**法条→判决映射：**
1. 【法定刑罚幅度】 -> 【约束】 -> 【具体刑期】
2. 【刑法条文】 -> 【直接依据】 -> 【执行方式】

## 特别注意：
1. **具体化**：避免抽象概念，多用具体的法条号、具体金额、具体时间等
2. **一致性**：同一概念在不同层次中的表述要保持一致
3. **完整性**：每个案例的推理链条要完整，从事实到法条到判决
4. **关联性**：多个案例间的共同模式要体现出来

请根据以下材料构建知识图谱：

{cases_text}

{laws_text}

## 构建要求：
请严格按照上述四个步骤，逐步构建完整的四层知识图谱。特别要注意案例中的具体事实与相关法条条文的精确对应关系。

## 输出要求：
1. 严格按照四个步骤输出一次完整的知识图谱
2. 不要重复输出相同内容
3. 不要添加额外的总结或解释
4. 输出完成后立即结束
"""

    try:
        # 调用qwen32B模型
        inputs = tokenizer(prompt, return_tensors="pt")
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=3072,  # 增加输出长度以容纳更复杂的知识图谱
                temperature=0.1,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # 解码输出
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # 提取模型回复部分（去掉输入的prompt）
        model_output = response[len(prompt):].strip()
        
        # 构建返回结构
        knowledge_graph = {
            "query_id": query_id,
            "total_cases": len(case_data_list),
            "total_laws": len(law_data_list),
            "model_output": model_output,
            "source_data": {
                "cases": case_data_list[:3],  # 保存top3案例的完整数据
                "laws": law_data_list[:3]     # 🎯 改为保存top3法条的完整数据
            }
        }
        
        return knowledge_graph
        
    except Exception as e:
        print(f"调用qwen32B模型时出错: {str(e)}")
        return None

def call_qwen32b_for_knowledge_graph(case_data: Dict, model, tokenizer) -> Dict:
    """
    调用qwen32B模型构建单个案例的知识图谱（保留原有功能）
    """
    fact = case_data.get('fact', '')
    reason = case_data.get('reason', '')
    result = case_data.get('result', '')
    
    # 构建prompt（与原代码相同）
    prompt = f"""你是专业的刑法知识图谱构建专家。请将案件材料分为**事实层**、**推理层**、**判决层**三个层面进行建图。

## 三层建模框架：

### 事实层（基于fact字段）：
**客观事实实体**：犯罪行为、时间地点、涉案物品、证据材料
**当事人实体**：被告人、被害人、证人（客观身份）
**行为属性**：具体行为描述、数量、金额、检测结果

### 推理层（基于reason字段）：
**法律认定实体**：罪名、犯罪构成、量刑情节、适用法条
**司法判断实体**：法院观点、量刑理由、法律后果
**法律属性**：法律评价、刑罚种类、处罚依据

### 判决层（基于result字段）：
**判决结果实体**：确定罪名、具体刑期、罚金数额、执行方式
**程序规定实体**：刑期计算、羁押折抵、罚金缴纳要求

## 输出格式：
**事实层实体：**
- 【实体名】（属性：【客观属性1：值1，客观属性2：值2】）

**推理层实体：**  
- 【实体名】（属性：【法律属性1：值1，法律属性2：值2】）

**判决层实体：**
- 【实体名】（属性：【判决执行属性1：值1，属性2：值2】）

**事实层关系：**
1. 【实体1】 -> 【客观关系】 -> 【实体2】

**推理层关系：**
1. 【实体1】 -> 【法律关系】 -> 【实体2】

**判决层关系：**
1. 【实体1】 -> 【判决关系】 -> 【实体2】

**映射关系：**
1. 【事实层实体】 -> 【支撑/证明/构成】 -> 【推理层实体】
2. 【推理层实体】 -> 【依据/确定】 -> 【判决层实体】

请根据以下案件材料构建知识图谱：

**案件事实：** {fact}

**法院判决：** {reason}

**判决结果：** {result}

## 输出要求：
1. 严格按照格式输出一次完整的知识图谱
2. 不要重复输出相同内容
3. 输出完成后立即结束
"""

    try:
        # 调用qwen32B模型
        inputs = tokenizer(prompt, return_tensors="pt")
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=2048,
                temperature=0.1,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # 解码输出
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # 提取模型回复部分（去掉输入的prompt）
        model_output = response[len(prompt):].strip()
        
        # 构建返回结构
        knowledge_graph = {
            "case_id": case_data.get('id', 'unknown'),
            "similarity": case_data.get('similarity', 0),
            "model_output": model_output,
            "raw_case_data": {
                "fact": fact,
                "reason": reason, 
                "result": result
            }
        }
        
        return knowledge_graph
        
    except Exception as e:
        print(f"调用qwen32B模型时出错: {str(e)}")
        return None

def build_comprehensive_knowledge_graph(query_id: str, case_results: Dict, law_results: Dict, 
                                       case_folder_paths: Dict, model, tokenizer):
    """
    为指定查询ID构建包含案例和法条的综合知识图谱
    
    Args:
        query_id: 查询ID
        case_results: 案例检索结果字典
        law_results: 法条检索结果字典
        case_folder_paths: 查询ID到案例文件夹路径的映射
        model: qwen32B模型
        tokenizer: 对应的tokenizer
    """
    
    # 检查是否存在对应的检索结果
    case_data_list = case_results.get(query_id, [])
    law_data_list = law_results.get(query_id, [])
    
    if not case_data_list and not law_data_list:
        print(f"警告: 未找到ID为{query_id}的任何检索结果（案例或法条）")
        return None
    
    if query_id not in case_folder_paths:
        print(f"警告: 未找到ID为{query_id}的文件夹路径")
        return None
    
    # 使用案例的原始文件夹路径作为输出目录（在此目录下新增文件）
    output_folder = case_folder_paths[query_id]
    
    print(f"开始为查询ID {query_id} 构建综合知识图谱...")
    print(f"案例数量: {len(case_data_list)}, 法条数量: {len(law_data_list)}")
    print(f"输出路径: {output_folder}")
    
    # 调用模型构建综合知识图谱
    print("调用qwen32B模型构建综合知识图谱...")
    kg = call_qwen32b_for_knowledge_graph_with_laws(case_data_list, law_data_list, model, tokenizer, query_id)
    
    if kg is None:
        print(f"查询ID {query_id} 综合知识图谱构建失败")
        return None
    
    # 保存综合知识图谱到案例文件夹（新增文件，不影响原始文件）
    comprehensive_output_file = os.path.join(output_folder, "knowledge_graph_comprehensive.json")
    with open(comprehensive_output_file, 'w', encoding='utf-8') as f:
        json.dump(kg, f, ensure_ascii=False, indent=2)
    
    print(f"综合知识图谱已保存到: {comprehensive_output_file}")
    print(f"查询ID {query_id} 的综合知识图谱构建完成!")
    
    return kg

def build_knowledge_graphs_for_top3_cases(query_id: str, case_results: Dict, folder_paths: Dict, model, tokenizer):
    """
    为指定查询ID的top3案例构建知识图谱并保存到原始文件夹（保留原有功能）
    """
    
    if query_id not in case_results:
        print(f"警告: 未找到ID为{query_id}的检索结果")
        return None
    
    if query_id not in folder_paths:
        print(f"警告: 未找到ID为{query_id}的文件夹路径")
        return None
    
    # 获取top3案例
    top3_cases = case_results[query_id][:3]
    
    if not top3_cases:
        print(f"警告: ID为{query_id}的检索结果为空")
        return None
    
    # 使用原始文件夹路径作为输出目录
    output_folder = folder_paths[query_id]
    
    knowledge_graphs = []
    
    print(f"开始为查询ID {query_id} 构建top3案例的知识图谱...")
    print(f"输出路径: {output_folder}")
    
    for i, case_data in enumerate(top3_cases):
        print(f"处理第{i+1}个案例，调用qwen32B模型...")
        
        # 调用qwen32B模型构建知识图谱
        kg = call_qwen32b_for_knowledge_graph(case_data, model, tokenizer)
        
        if kg is None:
            print(f"案例{i+1}知识图谱构建失败，跳过")
            continue
            
        kg["rank"] = i + 1  # 添加排名信息
        knowledge_graphs.append(kg)
        
        # 保存单个案例的知识图谱到原始文件夹（新增文件）
        case_output_file = os.path.join(output_folder, f"knowledge_graph_case_{i+1}.json")
        with open(case_output_file, 'w', encoding='utf-8') as f:
            json.dump(kg, f, ensure_ascii=False, indent=2)
        
        print(f"案例{i+1}知识图谱已保存到: {case_output_file}")
    
    if not knowledge_graphs:
        print(f"查询ID {query_id} 没有成功构建任何知识图谱")
        return None
    
    # 保存合并的知识图谱到原始文件夹（新增文件）
    combined_output_file = os.path.join(output_folder, "knowledge_graphs_top3_combined.json")
    combined_kg = {
        "query_id": query_id,
        "folder_path": output_folder,
        "total_cases": len(knowledge_graphs),
        "knowledge_graphs": knowledge_graphs,
        "creation_time": str(torch.cuda.current_device()) if torch.cuda.is_available() else "cpu"
    }
    
    with open(combined_output_file, 'w', encoding='utf-8') as f:
        json.dump(combined_kg, f, ensure_ascii=False, indent=2)
    
    print(f"合并的知识图谱已保存到: {combined_output_file}")
    print(f"查询ID {query_id} 的知识图谱构建完成! 共处理 {len(knowledge_graphs)} 个案例")
    
    return knowledge_graphs

def batch_build_comprehensive_knowledge_graphs(case_results: Dict, law_results: Dict, case_folder_paths: Dict, 
                                             model, tokenizer, query_ids: List[str] = None):
    """
    批量构建多个查询的综合知识图谱（案例+法条）
    
    Args:
        case_results: 案例检索结果字典
        law_results: 法条检索结果字典
        case_folder_paths: 查询ID到案例文件夹路径的映射
        model: qwen32B模型
        tokenizer: 对应的tokenizer
        query_ids: 要处理的查询ID列表，如果为None则处理所有
    """
    
    if query_ids is None:
        # 合并所有可用的查询ID
        all_case_ids = set(case_results.keys())
        all_law_ids = set(law_results.keys())
        query_ids = list(all_case_ids.union(all_law_ids))
    
    print(f"开始批量处理 {len(query_ids)} 个查询的综合知识图谱...")
    
    success_count = 0
    
    for query_id in query_ids:
        try:
            print(f"\n{'='*60}")
            print(f"正在处理查询ID: {query_id}")
            print(f"{'='*60}")
            
            result = build_comprehensive_knowledge_graph(query_id, case_results, law_results, 
                                                       case_folder_paths, model, tokenizer)
            if result is not None:
                success_count += 1
        except Exception as e:
            print(f"处理查询ID {query_id} 时出错: {str(e)}")
            continue
    
    print(f"\n批量处理完成! 成功处理 {success_count}/{len(query_ids)} 个查询")

def batch_build_knowledge_graphs(case_results: Dict, folder_paths: Dict, model, tokenizer, query_ids: List[str] = None):
    """
    批量构建多个查询的知识图谱（保留原有功能 - 仅案例）
    """
    
    if query_ids is None:
        query_ids = list(case_results.keys())
    
    print(f"开始批量处理 {len(query_ids)} 个查询的知识图谱...")
    
    success_count = 0
    
    for query_id in query_ids:
        try:
            print(f"\n{'='*50}")
            print(f"正在处理查询ID: {query_id}")
            print(f"{'='*50}")
            
            result = build_knowledge_graphs_for_top3_cases(query_id, case_results, folder_paths, model, tokenizer)
            if result is not None:
                success_count += 1
        except Exception as e:
            print(f"处理查询ID {query_id} 时出错: {str(e)}")
            continue
    
    print(f"\n批量处理完成! 成功处理 {success_count}/{len(query_ids)} 个查询")

# ============== 修改后的主函数 - 支持模式选择 ==============
import argparse

def main():
    """主函数 - 支持选择运行模式"""
    
    # 添加命令行参数解析
    parser = argparse.ArgumentParser(description='法律知识图谱构建程序')
    parser.add_argument('--modes', type=str, default='1,2,3,4', 
                       help='要运行的模式，用逗号分隔。例如: --modes 1,2 或 --modes 1 或 --modes 1,2,3,4')
    parser.add_argument('--sample_size', type=int, default=5,
                       help='处理的样本数量，默认5个查询ID')
    parser.add_argument('--model_path', type=str, default='/data1/xubuqiang/QWEN32',
                       help='模型路径')
    parser.add_argument('--case_folder', type=str, default='/data2/xubuqiang/RAG_dataset/results_folder_train_5_tasks',
                       help='案例结果文件夹路径')
    parser.add_argument('--law_folder', type=str, default='/data2/xubuqiang/RAG_dataset/results_folder_law_articles_train_5_tasks',
                       help='法条结果文件夹路径')
    
    args = parser.parse_args()
    
    # 解析要运行的模式
    try:
        selected_modes = [int(mode.strip()) for mode in args.modes.split(',')]
        valid_modes = [1, 2, 3, 4]
        selected_modes = [mode for mode in selected_modes if mode in valid_modes]
        
        if not selected_modes:
            print("❌ 错误：没有有效的模式选择，请使用1-4之间的数字")
            return
            
    except ValueError:
        print("❌ 错误：模式参数格式不正确，请使用逗号分隔的数字，如: 1,2,3")
        return
    
    print(f"🚀 开始运行法律知识图谱构建程序 - 选择模式: {selected_modes}")
    print("="*80)
    
    # 使用命令行参数设置路径
    model_path = args.model_path
    case_results_folder = args.case_folder
    law_results_folder = args.law_folder
    
    print(f"📂 模型路径: {model_path}")
    print(f"📂 案例文件夹: {case_results_folder}")
    print(f"📂 法条文件夹: {law_results_folder}")
    print(f"📊 样本数量: {args.sample_size}")
    
    # 第一步：加载模型
    print("📖 第一步：加载Qwen32B模型...")
    try:
        model, tokenizer = load_qwen32b_model(model_path)
        print("✅ 模型加载成功！")
    except Exception as e:
        print(f"❌ 模型加载失败: {str(e)}")
        return
    
    # 第二步：加载检索结果
    print("\n📂 第二步：加载检索结果...")
    print("加载案例检索结果...")
    case_results, case_folder_paths = load_retrieval_results_by_id(case_results_folder)
    
    print("加载法条检索结果...")
    law_results, law_folder_paths = load_retrieval_results_by_id(law_results_folder)
    
    if not case_results and not law_results:
        print("❌ 未找到任何检索结果，程序退出")
        return
    
    # 获取可用的查询ID
    case_query_ids = set(case_results.keys())
    law_query_ids = set(law_results.keys())
    all_query_ids = list(case_query_ids.union(law_query_ids))
    
    print(f"✅ 找到 {len(case_results)} 个案例查询结果")
    print(f"✅ 找到 {len(law_results)} 个法条查询结果")
    print(f"✅ 总共 {len(all_query_ids)} 个唯一查询ID")
    
    # 取样本数量作为示例（你可以改为更多或全部）
    sample_query_ids = all_query_ids[:args.sample_size]
    print(f"📋 将处理的查询IDs: {sample_query_ids}")
    
    # ================== 根据选择运行对应模式 ==================
    
    print("\n" + "="*80)
    print(f"🔥 开始运行选定的功能模式: {selected_modes}...")
    print("="*80)
    
    # 模式1：构建单个查询的综合知识图谱（案例+法条）
    if 1 in selected_modes and sample_query_ids:
        query_id = sample_query_ids[0]
        print(f"\n🎯 模式1: 为查询ID {query_id} 构建综合知识图谱（案例+法条）...")
        print("-" * 60)
        
        try:
            comprehensive_kg = build_comprehensive_knowledge_graph(
                query_id, case_results, law_results, case_folder_paths, model, tokenizer
            )
            
            if comprehensive_kg:
                print(f"✅ 模式1完成: 成功构建了包含 {comprehensive_kg['total_cases']} 个案例和 {comprehensive_kg['total_laws']} 个法条的综合知识图谱")
            else:
                print("❌ 模式1失败: 综合知识图谱构建失败")
        except Exception as e:
            print(f"❌ 模式1出错: {str(e)}")
    
    # 模式2：批量构建综合知识图谱（案例+法条）
    if 2 in selected_modes:
        print(f"\n🎯 模式2: 批量构建综合知识图谱（案例+法条）...")
        print("-" * 60)
        
        try:
            batch_build_comprehensive_knowledge_graphs(
                case_results, law_results, case_folder_paths, model, tokenizer, sample_query_ids
            )
            print("✅ 模式2完成: 批量综合知识图谱构建完成")
        except Exception as e:
            print(f"❌ 模式2出错: {str(e)}")
    
    # 模式3：构建单个查询的仅案例知识图谱
    if 3 in selected_modes and sample_query_ids:
        query_id = sample_query_ids[0]
        print(f"\n🎯 模式3: 为查询ID {query_id} 构建仅案例知识图谱...")
        print("-" * 60)
        
        try:
            case_kg = build_knowledge_graphs_for_top3_cases(
                query_id, case_results, case_folder_paths, model, tokenizer
            )
            
            if case_kg:
                print(f"✅ 模式3完成: 成功构建了 {len(case_kg)} 个案例的知识图谱")
            else:
                print("❌ 模式3失败: 案例知识图谱构建失败")
        except Exception as e:
            print(f"❌ 模式3出错: {str(e)}")
    
    # 模式4：批量构建仅案例知识图谱
    if 4 in selected_modes:
        print(f"\n🎯 模式4: 批量构建仅案例知识图谱...")
        print("-" * 60)
        
        try:
            batch_build_knowledge_graphs(
                case_results, case_folder_paths, model, tokenizer, sample_query_ids
            )
            print("✅ 模式4完成: 批量案例知识图谱构建完成")
        except Exception as e:
            print(f"❌ 模式4出错: {str(e)}")
    
    # 最终总结
    print("\n" + "="*80)
    print(f"🎉 选定模式 {selected_modes} 运行完成！")
    print("="*80)
    print(f"📊 处理统计:")
    print(f"   - 案例查询数量: {len(case_results)}")
    print(f"   - 法条查询数量: {len(law_results)}")
    print(f"   - 处理样本数量: {len(sample_query_ids)}")
    print(f"   - 运行模式: {selected_modes}")
    print("\n✨ 所有知识图谱文件已保存到对应的查询文件夹中！")

if __name__ == "__main__":
    main()