import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
os.environ['HUGGINGFACE_HUB_URL'] = 'https://hf-mirror.com'
os.environ['HF_HUB_BASE_URL'] = 'https://hf-mirror.com'

import glob
import pickle
import requests
import json
import time
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed
import torch
import uuid
import sys
import ijson
from decimal import Decimal
from neo4j import GraphDatabase
from collections import defaultdict
import numpy as np
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
import re
from collections import Counter
import random
from sentence_transformers import SentenceTransformer
from json_repair import repair_json

# api url; api key是您的专属设定
# 对于具体使用的llm模型：
# 1. 在总结归纳seed concepts阶段，我们使用强模型MEDXPERT_THINKER=o3对现有的medxpert心血管领域100+问题进行扩增，抽取其分析问题的思考过程。我们的目的并非让o3正确回答每个问题，而是想在其思考分析每个心血管问题之后，由其思考过程总结心血管种子概念。
# 2. 在提取出o3的思维过程后，我们用MEDXPERT_THINK_EXTRACTOR=gpt-4o对其思考过程抽取心血管概念名词作为种子概念
# 对以上两环节，我们的目的是构建种子概念库；如果您有您的专属领域文档库，您大可以跳过这一阶段，直接由您的领域文档抽取出您感兴趣的种子概念。
# 3. 在种子概念库建立完成后，我们将进行领域概念图扩增。这一环节我们建议根据我们的实验结果和您的成本预算酌情选择您的模型。在我们的试验阶段，我们以CONCEPT_GRAPH_EXTENDER=gpt-4omini为参考。
# 4. 构建完成领域概念图后，我们将对(a)每个概念进行单独的qa生成(b)每个有意义的概念对进行qa生成。由于qa生成依赖模型能力，我们建议您使用强模型。在我们的实验阶段，我们以QA_SYNTHESIZER=gpt-4o为参考。
# 我们对所有使用的embedding模型应用英文embedding开源模EMBEDDING_MODEL=nomic-ai/nomic-embed-text-v1.5
api_url="your api url"
api_key="your api key"

#各种可能的model：
MEDXPERT_THINKER = 'o3'
MEDXPERT_THINK_EXTRACTOR = 'gpt-4o'
CONCEPT_GRAPH_EXTENDER = 'gpt-4o-mini'
QA_SYNTHESIZER = 'gpt-4o'
EMBEDDING_MODEL = 'nomic-ai/nomic-embed-text-v1.5'

# 从最新的medxpert数据集提取种子概念，其优势在有对应领域的划分
import json
from collections import Counter

data = []
with open("medxpertqa_text_input.jsonl", "r", encoding="utf-8") as f:
    for line in f:
        data.append(json.loads(line.strip()))

# 取出心血管领域试题
body_system_counts = Counter(entry["body_system"] for entry in data)
heart_data = []
for i in range(len(data)):
    if data[i]['body_system']=='Cardiovascular':
        heart_data.append(data[i])

@dataclass
class AnalysisResult:
    """分析结果数据类"""
    status: str  # success, api_error
    question_id: str
    input_data: Dict
    response: Optional[str] = None
    error_details: Optional[str] = None
    processing_time: Optional[float] = None

class APIClient:
    """API调用客户端"""

    def __init__(self, api_url: str, api_key: str, model_name: str):
        self.api_url = api_url
        self.api_key = api_key
        self.model_name = model_name

        # 创建session并配置连接池和重试策略
        self.session = requests.Session()

        retry_strategy = Retry(
            total=3,
            backoff_factor=1,
            status_forcelist=[500, 502, 503, 504],
        )

        adapter = HTTPAdapter(pool_connections=10, pool_maxsize=20, max_retries=retry_strategy)
        self.session.mount('http://', adapter)
        self.session.mount('https://', adapter)

        self.session.headers.update({
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        })

    def call_api(self, messages: List[Dict], timeout: int = 120) -> Dict:
        """API调用"""
        data = {
            "model": self.model_name,
            "messages": messages,
            "stream": False
        }

        try:
            response = self.session.post(url=self.api_url, json=data, timeout=timeout)
            response.raise_for_status()
            result = response.json()
            return {
                "status": "success",
                "content": result['choices'][0]['message']['content']
            }
        except requests.exceptions.RequestException as e:
            return {
                "status": "error",
                "error": str(e)
            }

class PromptTemplate:
    """提示词模板类"""

    @staticmethod
    def get_system_prompt() -> str:
        return """You are a world-renowned clinical professor at a top teaching hospital with over 20 years of experience. Your thinking is grounded in evidence-based medicine, characterized by rigorous logic and clear reasoning.

Your core mission extends beyond solving clinical problems—you must **teach young doctors and medical students your decision-making process**. Therefore, when analyzing any case, you must:

1. **Systematic Deconstruction**: Begin by breaking down the problem from a macro perspective, identifying core clinical contradictions and key information.
2. **Comprehensive Evaluation**: Provide independent and thorough analysis of all possibilities (including every option), without skipping any.
3. **Clear Reasoning**: Explicitly articulate the "because-therefore" logic behind each judgment, clearly stating which specific clinical indicators, guideline consensus, or pathophysiological principles your decisions are based on.
4. **Principle Extraction**: After analysis, skillfully distill complex individual case decision processes into reusable, instructive core principles.

Your language should combine authority with clarity, enabling listeners to fully replicate your thought process."""

    @staticmethod
    def get_analysis_prompt(question_data: Dict) -> str:
        question = question_data['question']
        options = question_data['options']
  
        options_text = ""
        for opt in options:
            options_text += f"({opt['letter']}) {opt['content']}\n"
        return f"""Analyze this cardiovascular medicine multiple-choice question systematically and select the SINGLE CORRECT ANSWER. Provide a comprehensive analysis that demonstrates expert clinical reasoning. 

    **[Clinical Problem]**
    ---
    {question}

    Answer Choices:
    {options_text}
    ---


    **[Analysis Structure]**

    **Part 1: Clinical Context Analysis**

    Begin by establishing the clinical foundation for this question:

    * **Clinical Scenario Identification**: What is the primary clinical situation being presented? (e.g., diagnostic workup, treatment decision, risk stratification, pathophysiology question, etc.)

    * **Key Clinical Elements**: What are the most important clinical details, patient characteristics, findings, or parameters mentioned in the question stem? Why are these details clinically significant?

    * **Question Focus**: What specific aspect of clinical medicine is this question testing? What clinical knowledge or decision-making skill is being assessed?

    * **Relevant Clinical Framework**: What established clinical guidelines, diagnostic criteria, or treatment algorithms are relevant to answering this question?

    **Part 2: Systematic Option Analysis**

    Now analyze each answer choice methodically:

    **Option (A): **
    * **Clinical Evaluation**: How does this option relate to the clinical scenario? What would be the clinical implications if this were the correct choice?
    * **Evidence-Based Assessment**: Based on current guidelines, evidence, and pathophysiology, is this option clinically appropriate? Why or why not?

    **Option (B): **
    * **Clinical Evaluation**: [Same analysis format]
    * **Evidence-Based Assessment**: [Same analysis format]

    [Continue this systematic analysis for each option through the last one]

    **Part 3: Final Answer and Clinical Synthesis**

    * **Clinical Summary**: Briefly synthesize the key clinical scenario from the question stem and the critical findings from my analysis of each option.

    * **Selected Answer**: Based on my systematic analysis, the correct answer is: **(Letter) [Brief restatement of the correct option]**

    * **Answer Justification**: Concisely explain why this is the best answer, focusing on the most compelling clinical evidence and reasoning.

    * **Option Comparison Summary**: Provide a brief comparative overview of why the chosen option is superior to the other alternatives, highlighting the key clinical distinctions.

    * **Clinical Teaching Point**: Detailedly summarize the essential clinical medicine principle demonstrated by this question as a practical clinical pearl.

    **CRITICAL REQUIREMENT**: End with a clear statement: "**FINAL ANSWER: (a [single] Letter, DO NOT GIVE DETAILED CONTENT IN OPTION)**"

    Begin your analysis now."""

# 抽取llm思考过程，为后续的种子概念做准备
class CardioAnalysisProcessor:
    """心血管问题分析处理器"""

    def __init__(self, api_client: APIClient):
        self.api_client = api_client
        self.template = PromptTemplate()

    def process_single_question(self, question_data: Dict, question_id: str) -> AnalysisResult:
        """处理单个心血管问题"""
        start_time = time.time()
  
        # 构建消息
        system_prompt = self.template.get_system_prompt()
        user_prompt = self.template.get_analysis_prompt(question_data)
  
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        # 调用API
        api_result = self.api_client.call_api(messages, timeout=120)
        processing_time = time.time() - start_time

        if api_result["status"] != "success":
            return AnalysisResult(
                status="api_error",
                question_id=question_id,
                input_data=question_data,
                error_details=api_result["error"],
                processing_time=processing_time
            )

        return AnalysisResult(
            status="success",
            question_id=question_id,
            input_data=question_data,
            response=api_result["content"],
            processing_time=processing_time
        )

class BatchProcessor:
    """批量处理器"""

    def __init__(self, processor: CardioAnalysisProcessor, output_dir: str = "cardio_analysis"):
        self.processor = processor
        self.output_dir = output_dir
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

    def process_questions_list(self, heart_data: List[Dict], max_workers: int = 10, 
                              batch_size: int = 50, batch_delay: int = 1) -> Dict:
        """批量处理心血管问题"""
        total_questions = len(heart_data)
        print(f"开始批量处理 {total_questions} 个心血管临床问题")
        print(f"批次大小: {batch_size}, 最大并发: {max_workers}")

        all_results = {}
        batch_num = 1

        # 分批处理
        for i in range(0, total_questions, batch_size):
            batch_data = heart_data[i:i + batch_size]
            print(f"\n处理批次 {batch_num}: 问题 {i+1}-{min(i+batch_size, total_questions)} ({len(batch_data)} 个)")

            batch_start_time = time.time()
            batch_results = self._process_batch(batch_data, max_workers, i)
            batch_end_time = time.time()

            # 保存批次结果
            self._save_batch_results(batch_results, batch_num, batch_start_time)

            all_results.update(batch_results)

            print(f"批次 {batch_num} 完成，耗时: {batch_end_time - batch_start_time:.2f} 秒")

            batch_num += 1

            # 批次间休息
            if i + batch_size < total_questions:
                print(f"批次间休息 {batch_delay} 秒...")
                time.sleep(batch_delay)

        print(f"\n所有批次处理完成！总计处理 {total_questions} 个问题")
        return all_results

    def _process_batch(self, batch_data: List[Dict], max_workers: int, start_index: int) -> Dict:
        """处理单个批次"""
        batch_results = {}

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # 提交任务
            future_to_question = {}
            for idx, question_data in enumerate(batch_data):
                # 使用原始ID或生成新ID
                question_id = question_data.get('id', f"cardio_{start_index + idx:06d}")
                future = executor.submit(self.processor.process_single_question, question_data, question_id)
                future_to_question[future] = (question_id, question_data)

            # 收集结果
            completed = 0
            for future in as_completed(future_to_question):
                question_id, question_data = future_to_question[future]
                #try:
                result = future.result()
                batch_results[question_id] = result

                # 简单状态显示
                status_symbol = "✓" if result.status == "success" else "✗"
                completed += 1
        
                if completed % 5 == 0 or completed == len(batch_data):
                    success_count = sum(1 for r in batch_results.values() if r.status == "success")
                    print(f"  已完成: {completed}/{len(batch_data)} (成功: {success_count}) {status_symbol}")

        return batch_results

    def _save_batch_results(self, batch_results: Dict, batch_num: int, start_time: float):
        """保存批次结果"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"cardio_analysis_batch_{batch_num:03d}_{timestamp}.pkl"
        filepath = os.path.join(self.output_dir, filename)

        # 统计信息
        total_count = len(batch_results)
        success_count = sum(1 for r in batch_results.values() if r.status == "success")
        error_count = total_count - success_count

        # 构建保存数据
        save_data = {
            "metadata": {
                "batch_num": batch_num,
                "timestamp": timestamp,
                "start_time": start_time,
                "total_questions": total_count,
                "successful_analyses": success_count,
                "failed_analyses": error_count,
                "success_rate": success_count / total_count if total_count > 0 else 0
            },
            "results": batch_results
        }

        # 保存到文件
        with open(filepath, 'wb') as f:
            pickle.dump(save_data, f)

        print(f"  批次结果已保存: {filename}")
        print(f"  成功: {success_count}/{total_count} ({success_count/total_count*100:.1f}%)")

        return filepath

def process_heart_data(heart_data: List[Dict], api_url: str, api_key: str, model_name: str, max_workers: int = 10, 
                      batch_size: int = 50, output_dir: str = "cardio_analysis"):
    """处理heart_data的便捷函数"""
    print(f"准备处理 {len(heart_data)} 个心血管临床问题")

    # 初始化API客户端
    api_client = APIClient(
        api_url=api_url,
        api_key=api_key,
        model_name=model_name
    )

    # 初始化处理器
    processor = CardioAnalysisProcessor(api_client)
    batch_processor = BatchProcessor(processor, output_dir=output_dir)

    # 批量处理
    results = batch_processor.process_questions_list(
        heart_data=heart_data,
        max_workers=max_workers,
        batch_size=batch_size,
        batch_delay=1
    )

    return results

# 批量抽取每个qa的思考过程
# max_workers: 访问api并发数
# batch_size：多少次api访问完成后进行一次访问结果文件保存

results = process_heart_data(heart_data, max_workers=100, batch_size=400, output_dir="cookbooktest", 
                             api_url = api_url,
                             api_key = api_key,
                             model_name = MEDXPERT_THINKER)
textlist = [results[i].response for i in results.keys()]

class PromptTemplate:

    @staticmethod
    def get_system_prompt() -> str:

        return """You are an experienced cardiovascular specialist, skilled in identifying and extracting medical concept terms from clinical texts.

Your task is to extract all relevant concept terms from cardiovascular clinical texts.

Extraction principles:
- Only extract cardiovascular-related medical concepts, terms, and nouns
- Only concept terms, not complete definitions or explanations
- Prefer single core terms (e.g., "myocardial infarction", "hypertension", "echocardiography")
- Use phrases only when they represent standard medical terminology that cannot be meaningfully separated (e.g., "atrial fibrillation", "coronary artery disease")
- Avoid descriptive combinations (e.g., "severe hypertension" → "hypertension")
- Avoid overly vague terms (e.g., "heart problem")
- Include but not limited to disease names, examination methods, drug treatments, anatomical structures, physiological indicators, clinical manifestations, assessment tools, and all other related concepts
- Remove duplicate concepts
- Sort by importance

Please ensure the output format strictly follows JSON format requirements."""

    @staticmethod
    def get_extraction_prompt(text_content: str) -> str:

        return f"""**Task: Extract concept terms from cardiovascular clinical text**

**Please extract all relevant cardiovascular concept terms from the following text:**

---
{text_content}
---

**Output format (strictly follow JSON format):**
```json
{{
"concepts": [
    "concept1",
    "concept2",
    "concept3",
    "..."
]
}}"""

seed_concepts = [
    'blood pressure', 'step-up', 'glucagon', 'therapeutic anticoagulation'...
]

class ConceptGraph:
  
    @classmethod
    def from_graph_dict(cls, graph_dict: Dict[str, List[str]], concept_mapping, model, similarity_threshold):
        """
        从保存的图字典重新构建ConceptGraph
        Args:
            graph_dict: 保存的邻接字典
            model: SentenceTransformer模型实例
            similarity_threshold：判为相似概念的threshold
            concept_mapping: 保存的所有见过的相似概念mapping，例如{'hearts':'heart'}
        Returns:
            ConceptGraph实例
        """
        if model is None:
            raise ValueError("模型参数不能为None，请先使用load_embedding_model()加载模型")
  
        # 创建实例但不初始化
        instance = cls.__new__(cls)
        instance.model = model
        instance.graph = graph_dict.copy()
        instance.concept_embeddings = {}
        instance.concept_mapping = concept_mapping  
        instance.similarity_threshold = similarity_threshold
  
        # 重新计算所有概念的embedding并建立自映射
        all_concepts = list(graph_dict.keys())
        if all_concepts:
            print(f"正在为 {len(all_concepts)} 个概念重新计算embedding...")
            all_embeddings = model.encode(all_concepts)
    
            for concept, embedding in zip(all_concepts, all_embeddings):
                instance.concept_embeddings[concept] = embedding
                #instance.concept_mapping[concept] = concept  # 建立自映射
    
            print("ConceptGraph重建完成")
  
        return instance
  
    def __init__(self, seed_concepts: List[str], model, similarity_threshold):
        """
        从种子概念初始化图，并建立embedding库
        Args:
            seed_concepts: 已经过外部去重的种子概念列表
            model: SentenceTransformer模型实例（必需）
            similarity_threshold：判为相似概念的threshold
        """
        if model is None:
            raise ValueError("模型参数不能为None，请先使用load_embedding_model()加载模型")
    
        self.model = model
        self.graph = {}
        self.concept_embeddings = {}  # 维护concept -> embedding的映射
        self.concept_mapping = {}     # 添加概念映射表
        self.similarity_threshold = similarity_threshold  # embedding相似度阈值
  
        # 清理种子概念
        cleaned_seeds = [concept.strip() for concept in seed_concepts if concept.strip()]
  
        print(f"正在为 {len(cleaned_seeds)} 个种子概念计算embedding...")
  
        # 批量计算embedding
        if cleaned_seeds:
            seed_embeddings = self.model.encode(cleaned_seeds)
    
            # 建立初始图、embedding库和映射表
            for concept, embedding in zip(cleaned_seeds, seed_embeddings):
                self.graph[concept] = []
                self.concept_embeddings[concept] = embedding
                self.concept_mapping[concept] = concept  # 建立自映射
  
        print(f"初始化概念图完成，种子概念数: {len(cleaned_seeds)}")
  
    def _get_target_concept(self, concept: str) -> Optional[str]:
        """统一的概念映射查找"""
        return self.concept_mapping.get(concept)
  
    def _is_similar_to_existing(self, new_concept: str, new_embedding: np.ndarray) -> Optional[str]:
        """
        检查新概念是否与已有概念相似
        Returns:
            如果相似，返回相似的已有概念；否则返回None
        """
        if not self.concept_embeddings:
            return None
    
        # 计算与所有已有概念的相似度
        existing_concepts = list(self.concept_embeddings.keys())
        existing_embeddings = np.array([self.concept_embeddings[concept] for concept in existing_concepts])
  
        # 计算cosine相似度
        similarities = self.model.similarity([new_embedding], existing_embeddings)[0]
  
        # 找到最相似的概念
        max_similarity_idx = np.argmax(similarities)
        max_similarity = similarities[max_similarity_idx]
  
        if max_similarity >= self.similarity_threshold:
            return existing_concepts[max_similarity_idx]
  
        return None
  
    def get_current_adjacency(self) -> Dict[str, List[str]]:
        """获取当前邻接字典"""
        return self.graph.copy()
  
    def calculate_metrics(self, expansion_results: Dict) -> Dict[str, float]:
        """计算扩增指标"""
        # 收集所有新概念
        all_new_concepts = []
        for result in expansion_results.values():
            if result.status == "success" and result.new_concepts:
                all_new_concepts.extend(result.new_concepts)
  
        if not all_new_concepts:
            return {"connectivity_rate": 0.0}
  
        existing_concepts = set(self.graph.keys())
  
        # 连通度：新生成的边（彼此都是旧节点）/ 前一轮已有总边数
        old_total_edges = sum(len(neighbors) for neighbors in self.graph.values()) // 2
  
        # 统计新生成的边（彼此都是旧节点）
        new_edges_between_old_nodes = 0
        for result in expansion_results.values():
            if result.status == "success" and result.new_concepts:
                center_concept = result.center_concept
                for new_concept in result.new_concepts:
                    if (new_concept in existing_concepts and 
                        center_concept != new_concept and
                        new_concept not in self.graph.get(center_concept, [])):
                        new_edges_between_old_nodes += 1
  
        connectivity_rate = new_edges_between_old_nodes / old_total_edges if old_total_edges > 0 else float('inf') 
  
        return {
            "connectivity_rate": connectivity_rate
        }
  
  
    def update_graph(self, expansion_results: Dict):
        """
        更新图结构 - 使用mapping机制进行去重
        """
        nodes_added = 0
        edges_added = 0
        embedding_duplicates = 0
  
        # 收集所有新概念
        all_new_concepts = []
        concept_to_centers = {}  # 记录每个新概念对应的中心概念
  
        for result in expansion_results.values():
            if result.status == "success" and result.new_concepts:
                center_concept = result.center_concept
                for new_concept in result.new_concepts:
                    if new_concept.strip():
                        cleaned_concept = new_concept.strip()
                        all_new_concepts.append(cleaned_concept)
                        if cleaned_concept not in concept_to_centers:
                            concept_to_centers[cleaned_concept] = []
                        concept_to_centers[cleaned_concept].append(center_concept)
  
        if not all_new_concepts:
            return nodes_added, edges_added, embedding_duplicates
  
        num_all_new_concepts = len(all_new_concepts)
        print(f"收到 {num_all_new_concepts} 个新概念（未去重）")
  
        # 使用mapping快速过滤已知概念
        concepts_need_embedding = []
        concept_targets = {}  # concept -> target_concept 映射
  
        for concept in all_new_concepts:
            target = self._get_target_concept(concept)
            if target is not None:
                # 已知概念，直接使用映射
                concept_targets[concept] = target
                if target != concept:
                    embedding_duplicates += 1
            else:
                # 需要计算embedding的新概念
                concepts_need_embedding.append(concept)
  
        # 只对未知概念计算embedding
        if concepts_need_embedding:
            unique_concepts = list(set(concepts_need_embedding))
            print(f"正在对 {len(unique_concepts)} 个新概念进行embedding去重...")
            new_embeddings = self.model.encode(unique_concepts)
    
            # 逐个处理需要embedding的概念
            total_concepts = len(unique_concepts)
            for idx, (new_concept, new_embedding) in enumerate(zip(unique_concepts, new_embeddings), 1):
                # 每500个输出一次进度
                if idx % 500 == 0 or idx == total_concepts:
                    print(f"  处理进度: {idx}/{total_concepts} ({idx/total_concepts*100:.1f}%)")
                # 检查是否与已有概念相似
                similar_concept = self._is_similar_to_existing(new_concept, new_embedding)
        
                if similar_concept:
                    # 发现相似概念，建立映射
                    self.concept_mapping[new_concept] = similar_concept
                    concept_targets[new_concept] = similar_concept
                    embedding_duplicates += 1
                else:
                    # 全新概念，添加到图中并建立自映射
                    self.graph[new_concept] = []
                    self.concept_embeddings[new_concept] = new_embedding
                    self.concept_mapping[new_concept] = new_concept
                    concept_targets[new_concept] = new_concept
                    nodes_added += 1
  
        # 添加边（连接到所有相关的中心概念）
        for concept in all_new_concepts:
            target_concept = concept_targets[concept]
    
            for center_concept in concept_to_centers[concept]:
                # 确保中心概念存在于图中
                if center_concept in self.graph:
                    # 双向连接
                    if target_concept not in self.graph[center_concept]:
                        self.graph[center_concept].append(target_concept)
                        edges_added += 1
            
                    if center_concept not in self.graph[target_concept]:
                        self.graph[target_concept].append(center_concept)
                        edges_added += 1
  
        print(f"去重完成: 新增节点 {nodes_added}, 新增边 {edges_added//2}, 去重概念 {embedding_duplicates}")
  
        return nodes_added, edges_added // 2, embedding_duplicates, nodes_added / num_all_new_concepts  # 无向图，边数除以2
  
    def get_graph_stats(self) -> Dict[str, int]:
        """获取图统计信息"""
        node_count = len(self.graph)
        edge_count = sum(len(neighbors) for neighbors in self.graph.values()) // 2
        return {"nodes": node_count, "edges": edge_count}

# JSON修复库导入
try:
    from json_repair import repair_json
    HAS_JSONREPAIR = True
    print("✓ jsonrepair库已加载，JSON修复功能已启用")
except ImportError:
    HAS_JSONREPAIR = False
    print("⚠ jsonrepair库未安装，将使用基础修复策略。运行 'pip install jsonrepair' 启用高级JSON修复")
    def repair_json(text):
        return text

class ResponseValidator:
    """响应验证器"""
  
    @staticmethod
    def validate_json_response(response_text: str, expected_keys: List[str]) -> Dict:
        """
        验证API返回的内容是否为有效JSON，包含预处理容错
  
        Returns:
            dict: {
                "is_valid_json": bool,
                "parsed_json": dict or None,
                "error_type": str,
                "raw_response": str
            }
        """
        if not response_text or not response_text.strip():
            return {
                "is_valid_json": False,
                "parsed_json": None,
                "error_type": "empty_response",
                "raw_response": response_text
            }
  
        repair_attempts = []
  
        try:
            # 预处理清理
            text = response_text.strip()
    
            # 1. 处理markdown代码块 ```json...``` 或 ```...```
            code_block_match = re.search(r'```(?:json)?\s*\n?(.*?)\n?```', text, re.DOTALL)
            if code_block_match:
                text = code_block_match.group(1).strip()
    
            # 2. 处理引号包装 '...' 或 "..."
            if (text.startswith("'") and text.endswith("'")) or (text.startswith('"') and text.endswith('"')):
                text = text[1:-1]
    
            # 3. 移除前后的反引号
            text = text.strip('`').strip()
    
            # 4. 查找JSON部分 - 第一个{到最后一个}
            json_match = re.search(r'(\{.*\})', text, re.DOTALL)
            if json_match:
                text = json_match.group(1)
    
            # 第一次尝试：直接解析
            try:
                parsed = json.loads(text)
                repair_attempts.append("direct_parse_success")
            except json.JSONDecodeError as e:
                repair_attempts.append(f"direct_parse_failed: {str(e)}")
        
                # 第二次尝试：使用jsonrepair修复后解析
                if HAS_JSONREPAIR:
                    try:
                        repaired_text = repair_json(text)
                        parsed = json.loads(repaired_text)
                        repair_attempts.append("jsonrepair_success")
                    except Exception as e:
                        repair_attempts.append(f"jsonrepair_failed: {str(e)}")
                        raise
                else:
                    repair_attempts.append("jsonrepair_not_available")
                    raise
    
            # 检查是否符合预期的结构
            if isinstance(parsed, dict) and all(key in parsed for key in expected_keys):
                return {
                    "is_valid_json": True,
                    "parsed_json": parsed,
                    "error_type": None,
                    "raw_response": response_text,
                    "repair_attempts": repair_attempts
                }
            else:
                missing_keys = [key for key in expected_keys if key not in parsed] if isinstance(parsed, dict) else expected_keys
                return {
                    "is_valid_json": False,
                    "parsed_json": parsed,
                    "error_type": f"missing_keys: expected {expected_keys}, missing {missing_keys}",
                    "raw_response": response_text,
                    "repair_attempts": repair_attempts
                }
        
        except json.JSONDecodeError as e:
            return {
                "is_valid_json": False,
                "parsed_json": None,
                "error_type": f"json_decode_error: {str(e)}",
                "raw_response": response_text,
                "repair_attempts": repair_attempts
            }
        except Exception as e:
            return {
                "is_valid_json": False,
                "parsed_json": None,
                "error_type": f"unexpected_error: {str(e)}",
                "raw_response": response_text,
                "repair_attempts": repair_attempts
            }

# 基于种子概念的图扩增迭代代码

@dataclass
class ConceptExpansionResult:
    """概念扩增结果数据类"""
    status: str  # success, api_error, json_error
    concept_id: str
    center_concept: str
    neighbors: List[str]
    response: Optional[str] = None
    error_details: Optional[str] = None
    processing_time: Optional[float] = None
    json_validation: Optional[Dict] = None
    new_concepts: Optional[List[str]] = None
    returned_center: Optional[str] = None  # LLM返回的center_concept

class PromptTemplate:
    """提示词模板类"""

    @staticmethod
    def get_system_prompt() -> str:
        """获取系统提示词"""
        return """You are an experienced cardiovascular specialist, skilled in building comprehensive concept graphs for the cardiovascular domain.

Your task is to expand a cardiovascular concept graph by generating new related concepts based on a given center concept and its existing connections."""

    @staticmethod
    def get_expansion_prompt(center_concept: str, neighbors: List[str]) -> str:
        """生成概念扩增提示词"""
        neighbors_text = ", ".join(neighbors) if neighbors else "None"
  
        return f"""**Task: Generate new cardiovascular concepts related to the center concept**

**Domain**: Cardiovascular medicine
**Relationship requirement**: New concepts should be directly related to the center concept through strong clinical medical associations

**Center concept**: {center_concept}
**Existing neighbor concepts of the center concept**: {neighbors_text}

**Output format (strictly follow JSON format):**

{{
  "center_concept": "{center_concept}",
  "new_concepts": [
    "concept1",
    "concept2",
    "concept3",
    "..."
  ]
}}


If no new concepts can be generated:

{{
  "center_concept": "{center_concept}",
  "new_concepts": ["NO NEW CONCEPTS"]
}}

**Instructions**:

1. Instead of generate general medical concept, focus on generating new cardiovascular-domain concepts that are directly relevant in clinical scenarios to "{center_concept}" with strong clinical medical relation
2. Do not repeat any existing connected concepts listed above
3. Prefer single core terms (e.g., "myocardial infarction", "hypertension", "echocardiography")
4. Use phrases only when they represent standard medical terminology that cannot be meaningfully separated (e.g., "atrial fibrillation", "coronary artery disease")
5. Avoid descriptive combinations (e.g., "severe hypertension" → "hypertension")
6. Avoid overly vague terms (e.g., "heart problem")
7. Generate concepts that are directly related to the center concept
8. Do not repeat any existing connected concepts listed above; Avoid duplicate concepts"""

class ConceptExpander:
    """概念扩增处理器"""

    def __init__(self, api_client: APIClient):
        self.api_client = api_client
        self.template = PromptTemplate()

    def expand_single_concept(self, center_concept: str, neighbors: List[str], concept_id: str) -> ConceptExpansionResult:
        """扩增单个概念"""
        start_time = time.time()
  
        # 构建消息
        system_prompt = self.template.get_system_prompt()
        user_prompt = self.template.get_expansion_prompt(center_concept, neighbors)
  
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        # 调用API
        api_result = self.api_client.call_api(messages, timeout=120)
        processing_time = time.time() - start_time

        if api_result["status"] != "success":
            return ConceptExpansionResult(
                status="api_error",
                concept_id=concept_id,
                center_concept=center_concept,
                neighbors=neighbors,
                error_details=api_result["error"],
                processing_time=processing_time
            )

        # 验证JSON响应
        expected_keys = ["center_concept", "new_concepts"]
        json_validation = ResponseValidator.validate_json_response(
            api_result["content"], expected_keys
        )
  
        if json_validation["is_valid_json"]:
            returned_center = json_validation["parsed_json"]["center_concept"]
            new_concepts = json_validation["parsed_json"]["new_concepts"]
    
            # 检查是否为"无新增概念"的情况
            if len(new_concepts) == 1 and new_concepts[0].strip() == "NO NEW CONCEPTS":
                return ConceptExpansionResult(
                    status="success",
                    concept_id=concept_id,
                    center_concept=center_concept,
                    neighbors=neighbors,
                    response=api_result["content"],
                    processing_time=processing_time,
                    json_validation=json_validation,
                    new_concepts=[],  # 空列表，表示没有新概念
                    returned_center=returned_center
                )
    
            # 正常处理新概念 - 移除之前的过滤逻辑，交给embedding去重
            new_concepts = [concept.strip() for concept in new_concepts if concept.strip()]
    
            return ConceptExpansionResult(
                status="success",
                concept_id=concept_id,
                center_concept=center_concept,
                neighbors=neighbors,
                response=api_result["content"],
                processing_time=processing_time,
                json_validation=json_validation,
                new_concepts=new_concepts,
                returned_center=returned_center
            )
        else:
            return ConceptExpansionResult(
                status="json_error",
                concept_id=concept_id,
                center_concept=center_concept,
                neighbors=neighbors,
                response=api_result["content"],
                error_details=f"JSON validation failed: {json_validation['error_type']}",
                processing_time=processing_time,
                json_validation=json_validation
            )

class BatchConceptExpander:
    """批量概念扩增处理器"""
    def __init__(self, expander: ConceptExpander, output_dir: str = "concept_expansion"):
        self.expander = expander
        self.output_dir = output_dir
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

    def expand_concepts_batch(self, adjacency_dict: Dict[str, List[str]], max_workers: int = 10) -> Dict:
        """批量扩增概念"""
        concepts_to_expand = list(adjacency_dict.keys())
        total_concepts = len(concepts_to_expand)
        print(f"开始批量概念扩增 {total_concepts} 个概念")
        print(f"最大并发: {max_workers}")

        batch_start_time = time.time()
        batch_results = self._process_batch(concepts_to_expand, adjacency_dict, max_workers)
        batch_end_time = time.time()

        print(f"处理完成，耗时: {batch_end_time - batch_start_time:.2f} 秒")
        return batch_results

    def _process_batch(self, batch_concepts: List[str], adjacency_dict: Dict[str, List[str]], 
                      max_workers: int) -> Dict:
        """处理批次"""
        batch_results = {}

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # 提交任务
            future_to_concept = {}
            for idx, concept in enumerate(batch_concepts):
                concept_id = f"concept_{idx:06d}"
                neighbors = adjacency_dict.get(concept, [])
                future = executor.submit(self.expander.expand_single_concept, concept, neighbors, concept_id)
                future_to_concept[future] = (concept_id, concept)

            # 收集结果
            completed = 0
            for future in as_completed(future_to_concept):
                concept_id, concept = future_to_concept[future]
                try:
                    result = future.result()
                    batch_results[concept_id] = result

                    # 简单状态显示
                    status_symbol = "✓" if result.status == "success" else "✗"
                    completed += 1
            
                    if completed % 1000 == 0 or completed == len(batch_concepts):
                        success_count = sum(1 for r in batch_results.values() if r.status == "success")
                        print(f"  已完成: {completed}/{len(batch_concepts)} (成功: {success_count}) {status_symbol}")

                except Exception as e:
                    batch_results[concept_id] = ConceptExpansionResult(
                        status="exception",
                        concept_id=concept_id,
                        center_concept=concept,
                        neighbors=adjacency_dict.get(concept, []),
                        error_details=str(e)
                    )
                    print(f"  异常: {concept_id} - {str(e)}")

        return batch_results

    def _save_results(self, batch_results: Dict, start_time: float):
        """保存结果"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"concept_expansion_results_{timestamp}.pkl"
        filepath = os.path.join(self.output_dir, filename)

        # 统计信息
        total_count = len(batch_results)
        success_count = sum(1 for r in batch_results.values() if r.status == "success")
        error_count = total_count - success_count

        # 统计生成的概念数量
        total_new_concepts = sum(len(r.new_concepts) for r in batch_results.values() 
                               if r.status == "success" and r.new_concepts)
  
        # 统计跳过的概念数量（无新增概念）
        skipped_concepts = sum(1 for r in batch_results.values() 
                              if r.status == "success" and r.new_concepts is not None and len(r.new_concepts) == 0)

        # 构建保存数据
        save_data = {
            "metadata": {
                "timestamp": timestamp,
                "start_time": start_time,
                "total_concepts": total_count,
                "successful_expansions": success_count,
                "failed_expansions": error_count,
                "success_rate": success_count / total_count if total_count > 0 else 0,
                "total_new_concepts": total_new_concepts,
                "skipped_concepts": skipped_concepts
            },
            "results": batch_results
        }

        # 保存到文件
        with open(filepath, 'wb') as f:
            pickle.dump(save_data, f)

        print(f"  结果已保存: {filename}")
        print(f"  成功: {success_count}/{total_count} ({success_count/total_count*100:.1f}%)")
        print(f"  新概念总数: {total_new_concepts}")
        print(f"  跳过概念: {skipped_concepts}")

        return filepath

class ConceptGraph:
  
    @classmethod
    def from_graph_dict(cls, graph_dict: Dict[str, List[str]], concept_mapping, model, similarity_threshold):
        """
        从保存的图字典重新构建ConceptGraph
        Args:
            graph_dict: 保存的邻接字典
            model: SentenceTransformer模型实例
            similarity_threshold：判为相似概念的threshold
            concept_mapping: 保存的所有见过的相似概念mapping，例如{'hearts':'heart'}
        Returns:
            ConceptGraph实例
        """
        if model is None:
            raise ValueError("模型参数不能为None，请先使用load_embedding_model()加载模型")
  
        # 创建实例但不初始化
        instance = cls.__new__(cls)
        instance.model = model
        instance.graph = graph_dict.copy()
        instance.concept_embeddings = {}
        instance.concept_mapping = concept_mapping  
        instance.similarity_threshold = similarity_threshold
  
        # 重新计算所有概念的embedding并建立自映射
        all_concepts = list(graph_dict.keys())
        if all_concepts:
            print(f"正在为 {len(all_concepts)} 个概念重新计算embedding...")
            all_embeddings = model.encode(all_concepts)
    
            for concept, embedding in zip(all_concepts, all_embeddings):
                instance.concept_embeddings[concept] = embedding
                #instance.concept_mapping[concept] = concept  # 建立自映射
    
            print("ConceptGraph重建完成")
  
        return instance
  
    def __init__(self, seed_concepts: List[str], model, similarity_threshold):
        """
        从种子概念初始化图，并建立embedding库
        Args:
            seed_concepts: 已经过外部去重的种子概念列表
            model: SentenceTransformer模型实例（必需）
            similarity_threshold：判为相似概念的threshold
        """
        if model is None:
            raise ValueError("模型参数不能为None，请先使用load_embedding_model()加载模型")
    
        self.model = model
        self.graph = {}
        self.concept_embeddings = {}  # 维护concept -> embedding的映射
        self.concept_mapping = {}     # 添加概念映射表
        self.similarity_threshold = similarity_threshold  # embedding相似度阈值
  
        # 清理种子概念
        cleaned_seeds = [concept.strip() for concept in seed_concepts if concept.strip()]
  
        print(f"正在为 {len(cleaned_seeds)} 个种子概念计算embedding...")
  
        # 批量计算embedding
        if cleaned_seeds:
            seed_embeddings = self.model.encode(cleaned_seeds)
    
            # 建立初始图、embedding库和映射表
            for concept, embedding in zip(cleaned_seeds, seed_embeddings):
                self.graph[concept] = []
                self.concept_embeddings[concept] = embedding
                self.concept_mapping[concept] = concept  # 建立自映射
  
        print(f"初始化概念图完成，种子概念数: {len(cleaned_seeds)}")
  
    def _get_target_concept(self, concept: str) -> Optional[str]:
        """统一的概念映射查找"""
        return self.concept_mapping.get(concept)
  
    def _is_similar_to_existing(self, new_concept: str, new_embedding: np.ndarray) -> Optional[str]:
        """
        检查新概念是否与已有概念相似
        Returns:
            如果相似，返回相似的已有概念；否则返回None
        """
        if not self.concept_embeddings:
            return None
    
        # 计算与所有已有概念的相似度
        existing_concepts = list(self.concept_embeddings.keys())
        existing_embeddings = np.array([self.concept_embeddings[concept] for concept in existing_concepts])
  
        # 计算cosine相似度
        similarities = self.model.similarity([new_embedding], existing_embeddings)[0]
  
        # 找到最相似的概念
        max_similarity_idx = np.argmax(similarities)
        max_similarity = similarities[max_similarity_idx]
  
        if max_similarity >= self.similarity_threshold:
            return existing_concepts[max_similarity_idx]
  
        return None
  
    def get_current_adjacency(self) -> Dict[str, List[str]]:
        """获取当前邻接字典"""
        return self.graph.copy()
  
    def calculate_metrics(self, expansion_results: Dict) -> Dict[str, float]:
        """计算扩增指标"""
        # 收集所有新概念
        all_new_concepts = []
        for result in expansion_results.values():
            if result.status == "success" and result.new_concepts:
                all_new_concepts.extend(result.new_concepts)
  
        if not all_new_concepts:
            return {"connectivity_rate": 0.0}
  
        existing_concepts = set(self.graph.keys())
  
        # 连通度：新生成的边（彼此都是旧节点）/ 前一轮已有总边数
        old_total_edges = sum(len(neighbors) for neighbors in self.graph.values()) // 2
  
        # 统计新生成的边（彼此都是旧节点）
        new_edges_between_old_nodes = 0
        for result in expansion_results.values():
            if result.status == "success" and result.new_concepts:
                center_concept = result.center_concept
                for new_concept in result.new_concepts:
                    if (new_concept in existing_concepts and 
                        center_concept != new_concept and
                        new_concept not in self.graph.get(center_concept, [])):
                        new_edges_between_old_nodes += 1
  
        connectivity_rate = new_edges_between_old_nodes / old_total_edges if old_total_edges > 0 else float('inf') 
  
        return {
            "connectivity_rate": connectivity_rate
        }
  
  
    def update_graph(self, expansion_results: Dict):
        """
        更新图结构 - 使用mapping机制进行去重
        """
        nodes_added = 0
        edges_added = 0
        embedding_duplicates = 0
  
        # 收集所有新概念
        all_new_concepts = []
        concept_to_centers = {}  # 记录每个新概念对应的中心概念
  
        for result in expansion_results.values():
            if result.status == "success" and result.new_concepts:
                center_concept = result.center_concept
                for new_concept in result.new_concepts:
                    if new_concept.strip():
                        cleaned_concept = new_concept.strip()
                        all_new_concepts.append(cleaned_concept)
                        if cleaned_concept not in concept_to_centers:
                            concept_to_centers[cleaned_concept] = []
                        concept_to_centers[cleaned_concept].append(center_concept)
  
        if not all_new_concepts:
            return nodes_added, edges_added, embedding_duplicates
  
        num_all_new_concepts = len(all_new_concepts)
        print(f"收到 {num_all_new_concepts} 个新概念（未去重）")
  
        # 使用mapping快速过滤已知概念
        concepts_need_embedding = []
        concept_targets = {}  # concept -> target_concept 映射
  
        for concept in all_new_concepts:
            target = self._get_target_concept(concept)
            if target is not None:
                # 已知概念，直接使用映射
                concept_targets[concept] = target
                if target != concept:
                    embedding_duplicates += 1
            else:
                # 需要计算embedding的新概念
                concepts_need_embedding.append(concept)
  
        # 只对未知概念计算embedding
        if concepts_need_embedding:
            unique_concepts = list(set(concepts_need_embedding))
            print(f"正在对 {len(unique_concepts)} 个新概念进行embedding去重...")
            new_embeddings = self.model.encode(unique_concepts)
    
            # 逐个处理需要embedding的概念
            total_concepts = len(unique_concepts)
            for idx, (new_concept, new_embedding) in enumerate(zip(unique_concepts, new_embeddings), 1):
                # 每500个输出一次进度
                if idx % 500 == 0 or idx == total_concepts:
                    print(f"  处理进度: {idx}/{total_concepts} ({idx/total_concepts*100:.1f}%)")
                # 检查是否与已有概念相似
                similar_concept = self._is_similar_to_existing(new_concept, new_embedding)
        
                if similar_concept:
                    # 发现相似概念，建立映射
                    self.concept_mapping[new_concept] = similar_concept
                    concept_targets[new_concept] = similar_concept
                    embedding_duplicates += 1
                else:
                    # 全新概念，添加到图中并建立自映射
                    self.graph[new_concept] = []
                    self.concept_embeddings[new_concept] = new_embedding
                    self.concept_mapping[new_concept] = new_concept
                    concept_targets[new_concept] = new_concept
                    nodes_added += 1
  
        # 添加边（连接到所有相关的中心概念）
        for concept in all_new_concepts:
            target_concept = concept_targets[concept]
    
            for center_concept in concept_to_centers[concept]:
                # 确保中心概念存在于图中
                if center_concept in self.graph:
                    # 双向连接
                    if target_concept not in self.graph[center_concept]:
                        self.graph[center_concept].append(target_concept)
                        edges_added += 1
            
                    if center_concept not in self.graph[target_concept]:
                        self.graph[target_concept].append(center_concept)
                        edges_added += 1
  
        print(f"去重完成: 新增节点 {nodes_added}, 新增边 {edges_added//2}, 去重概念 {embedding_duplicates}")
  
        return nodes_added, edges_added // 2, embedding_duplicates, nodes_added / num_all_new_concepts  # 无向图，边数除以2
  
    def get_graph_stats(self) -> Dict[str, int]:
        """获取图统计信息"""
        node_count = len(self.graph)
        edge_count = sum(len(neighbors) for neighbors in self.graph.values()) // 2
        return {"nodes": node_count, "edges": edge_count}

def load_embedding_model(model_name: str = "nomic-ai/nomic-embed-text-v1.5"):
    """
    加载embedding模型
    Args:
        model_name: 模型名称
    Returns:
        SentenceTransformer模型实例
    """
    print(f"正在加载embedding模型: {model_name}")
    model = SentenceTransformer(model_name, trust_remote_code=True)
    print("模型加载完成")
    return model

def extract_seed_concepts(results):
    """从批量处理结果中提取种子概念"""
    all_concepts = []
  
    for result in results.values():
        if result.status == "success" and result.extracted_concepts:
            all_concepts.extend(result.extracted_concepts)
  
    # 去空格并去重
    seed_concepts = list(set(concept.strip() for concept in all_concepts if concept.strip()))
  
    return seed_concepts

def run_concept_expansion_iteration(api_url: str, api_key: str, model_name: str, concept_graph: ConceptGraph, max_workers: int = 10):
    """运行单次概念扩增迭代"""
    # 初始化API客户端和处理器
    api_client = APIClient(
        api_url=api_url,
        api_key=api_key,
        model_name=model_name
    )
  
    expander = ConceptExpander(api_client)
    batch_expander = BatchConceptExpander(expander)
  
    # 获取当前邻接字典
    current_adjacency = concept_graph.get_current_adjacency()
  
    # 批量扩增概念
    expansion_results = batch_expander.expand_concepts_batch(
        adjacency_dict=current_adjacency,
        max_workers=max_workers
    )
  
    # 计算指标
    metrics = concept_graph.calculate_metrics(expansion_results)
  
    # 更新图 - 使用embedding去重
    nodes_added, edges_added, embedding_duplicates, concept_add_rate = concept_graph.update_graph(expansion_results)

    # 获取更新后的图统计
    graph_stats = concept_graph.get_graph_stats()
  
    # 统计跳过的概念数量
    skipped_count = sum(1 for r in expansion_results.values() 
                       if r.status == "success" and r.new_concepts is not None and len(r.new_concepts) == 0)
  
    # 打印结果
    print(f"\n=== 迭代完成 ===")
    print(f"概念更新度: {concept_add_rate:.3f}")
    print(f"概念连通度: {metrics['connectivity_rate']:.3f}")
    print(f"迭代结束节点数: {graph_stats['nodes']}")
    print(f"迭代结束边数: {graph_stats['edges']}")
    print(f"本轮新增节点: {nodes_added}")
    print(f"本轮新增边数: {edges_added}")
    print(f"跳过的概念: {skipped_count}")
  
    return {
        "concept_add_rate": concept_add_rate,
        "connectivity_rate": metrics['connectivity_rate'],
        "graph_stats": graph_stats,
        "nodes_added": nodes_added,
        "edges_added": edges_added,
        "embedding_duplicates": embedding_duplicates,
        "skipped_count": skipped_count,
        "expansion_results": expansion_results
    }