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; the api key is your specific setting
# For the specific LLM models used:
# 1. In the stage of summarizing and inducing seed concepts, we use a powerful model MEDXPERT_THINKER=o3 to augment the existing 100+ questions in the medxpert cardiovascular domain and extract its thought process for analyzing the problems. Our purpose is not to have o3 answer each question correctly, but rather, after it thinks through and analyzes each cardiovascular problem, to summarize cardiovascular seed concepts from its thought process.
# 2. After extracting o3's thought process, we use MEDXPERT_THINK_EXTRACTOR=gpt-4o to extract cardiovascular concept nouns from its thought process to serve as seed concepts.
# For the two steps above, our purpose is to build a seed concept library; if you have your own proprietary domain document library, you can feel free to skip this stage and directly extract the seed concepts you are interested in from your domain documents.
# 3. After the seed concept library is established, we will proceed with domain concept graph augmentation. For this step, we recommend that you choose your model based on our experimental results and your cost budget as appropriate. In our experimental phase, we use CONCEPT_GRAPH_EXTENDER=gpt-4omini as a reference.
# 4. After the domain concept graph is completely built, we will perform separate QA generation for (a) each concept and (b) each meaningful concept pair. Since QA generation depends on the model's capabilities, we recommend you use a powerful model. In our experimental phase, we use QA_SYNTHESIZER=gpt-4o as a reference.
# For all embedding models used, we apply the open-source English embedding model EMBEDDING_MODEL=nomic-ai/nomic-embed-text-v1.5
api_url="your api url"
api_key="your api key"

# all possibly used models：
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'

# Extract seed concepts from the latest medxpert dataset, which has the advantage of having corresponding field divisions
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()))

# Extract cardiovascular field test questions
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:
    """class for analyzing the result"""
    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:
    """class for api calling"""

    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

        # Create a session and configure the connection pool and retry strategy
        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:
        """call 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."""

# Extract the LLM thought process, to prepare for subsequent seed concepts
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_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:
        """Batch process cardiovascular questions"""
        total_questions = len(heart_data)
        print(f"Starting batch processing for {total_questions} cardiovascular clinical questions")
        print(f"Batch size: {batch_size}, Max workers: {max_workers}")

        all_results = {}
        batch_num = 1

        # Process in batches
        for i in range(0, total_questions, batch_size):
            batch_data = heart_data[i:i + batch_size]
            print(f"\nProcessing batch {batch_num}: Questions {i+1}-{min(i+batch_size, total_questions)} ({len(batch_data)} items)")

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

            # Save batch results
            self._save_batch_results(batch_results, batch_num, batch_start_time)

            all_results.update(batch_results)

            print(f"Batch {batch_num} finished, time taken: {batch_end_time - batch_start_time:.2f} seconds")

            batch_num += 1

            # Rest between batches
            if i + batch_size < total_questions:
                print(f"Resting for {batch_delay} seconds between batches...")
                time.sleep(batch_delay)

        print(f"\nAll batches processed! Total questions processed: {total_questions}")
        return all_results

    def _process_batch(self, batch_data: List[Dict], max_workers: int, start_index: int) -> Dict:
        """Process a single batch"""
        batch_results = {}

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # Submit tasks
            future_to_question = {}
            for idx, question_data in enumerate(batch_data):
                # Use original ID or generate a new one
                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)

            # Collect results
            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

                # Simple status display
                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: {completed}/{len(batch_data)} (Success: {success_count}) {status_symbol}")

        return batch_results

    def _save_batch_results(self, batch_results: Dict, batch_num: int, start_time: float):
        """Save batch results"""
        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)

        # Statistics
        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

        # Construct data to save
        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
        }

        # Save to file
        with open(filepath, 'wb') as f:
            pickle.dump(save_data, f)

        print(f"  Batch results saved to: {filename}")
        print(f"  Success: {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"):
    """A convenience function for processing heart_data"""
    print(f"Preparing to process {len(heart_data)} cardiovascular clinical questions")

    # Initialize the API client
    api_client = APIClient(
        api_url=api_url,
        api_key=api_key,
        model_name=model_name
    )

    # Initialize the processor
    processor = CardioAnalysisProcessor(api_client)
    batch_processor = BatchProcessor(processor, output_dir=output_dir)

    # Batch processing
    results = batch_processor.process_questions_list(
        heart_data=heart_data,
        max_workers=max_workers,
        batch_size=batch_size,
        batch_delay=1
    )

    return results

```python
# Batch extract each QA's thought process
# max_workers: Number of concurrent API accesses
# batch_size: Number of API accesses before saving the access results to a file

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):
        """
        Reconstructs a ConceptGraph from a saved graph dictionary
        Args:
            graph_dict: The saved adjacency dictionary
            model: An instance of the SentenceTransformer model
            similarity_threshold: The threshold for judging concepts as similar
            concept_mapping: The mapping of all previously seen similar concepts, e.g., {'hearts':'heart'}
        Returns:
            An instance of ConceptGraph
        """
        if model is None:
            raise ValueError("The model parameter cannot be None, please load the model using load_embedding_model() first")

        # Create an instance without initializing
        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

        # Recalculate embeddings for all concepts and establish self-mapping
        all_concepts = list(graph_dict.keys())
        if all_concepts:
            print(f"Recalculating embeddings for {len(all_concepts)} concepts...")
            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  # Establish self-mapping
    
            print("ConceptGraph reconstruction complete")

        return instance

    def __init__(self, seed_concepts: List[str], model, similarity_threshold):
        """
        Initializes the graph from seed concepts and builds the embedding library
        Args:
            seed_concepts: A list of seed concepts that have been externally deduplicated
            model: An instance of the SentenceTransformer model (required)
            similarity_threshold: The threshold for judging concepts as similar
        """
        if model is None:
            raise ValueError("The model parameter cannot be None, please load the model using load_embedding_model() first")
    
        self.model = model
        self.graph = {}
        self.concept_embeddings = {}  # Maintains the mapping of concept -> embedding
        self.concept_mapping = {}     # Add a concept mapping table
        self.similarity_threshold = similarity_threshold  # embedding similarity threshold

        # Clean the seed concepts
        cleaned_seeds = [concept.strip() for concept in seed_concepts if concept.strip()]

        print(f"Calculating embeddings for {len(cleaned_seeds)} seed concepts...")

        # Batch calculate embeddings
        if cleaned_seeds:
            seed_embeddings = self.model.encode(cleaned_seeds)
    
            # Build the initial graph, embedding library, and mapping table
            for concept, embedding in zip(cleaned_seeds, seed_embeddings):
                self.graph[concept] = []
                self.concept_embeddings[concept] = embedding
                self.concept_mapping[concept] = concept  # Establish self-mapping

        print(f"Concept graph initialization complete, number of seed concepts: {len(cleaned_seeds)}")

    def _get_target_concept(self, concept: str) -> Optional[str]:
        """Unified concept mapping lookup"""
        return self.concept_mapping.get(concept)

    def _is_similar_to_existing(self, new_concept: str, new_embedding: np.ndarray) -> Optional[str]:
        """
        Checks if a new concept is similar to existing ones
        Returns:
            If similar, returns the similar existing concept; otherwise, returns None
        """
        if not self.concept_embeddings:
            return None
    
        # Calculate similarity with all existing concepts
        existing_concepts = list(self.concept_embeddings.keys())
        existing_embeddings = np.array([self.concept_embeddings[concept] for concept in existing_concepts])

        # Calculate cosine similarity
        similarities = self.model.similarity([new_embedding], existing_embeddings)[0]

        # Find the most similar concept
        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]]:
        """Gets the current adjacency dictionary"""
        return self.graph.copy()

    def calculate_metrics(self, expansion_results: Dict) -> Dict[str, float]:
        """Calculates augmentation metrics"""
        # Collect all new concepts
        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())

        # Connectivity: newly generated edges (between two old nodes) / total existing edges from the previous round
        old_total_edges = sum(len(neighbors) for neighbors in self.graph.values()) // 2

        # Count newly generated edges (between two old nodes)
        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):
        """
        Updates the graph structure, using a mapping mechanism for deduplication.
        """
        nodes_added = 0
        edges_added = 0
        embedding_duplicates = 0

        # Collect all new concepts
        all_new_concepts = []
        concept_to_centers = {}  # Record the corresponding center concept for each new concept

        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"Received {num_all_new_concepts} new concepts (not deduplicated)")

        # Use mapping to quickly filter known concepts
        concepts_need_embedding = []
        concept_targets = {}  # concept -> target_concept mapping

        for concept in all_new_concepts:
            target = self._get_target_concept(concept)
            if target is not None:
                # Known concept, use mapping directly
                concept_targets[concept] = target
                if target != concept:
                    embedding_duplicates += 1
            else:
                # New concept that needs an embedding calculation
                concepts_need_embedding.append(concept)

        # Only calculate embeddings for unknown concepts
        if concepts_need_embedding:
            unique_concepts = list(set(concepts_need_embedding))
            print(f"Performing embedding-based deduplication on {len(unique_concepts)} new concepts...")
            new_embeddings = self.model.encode(unique_concepts)
    
            # Process concepts that need embeddings one by one
            total_concepts = len(unique_concepts)
            for idx, (new_concept, new_embedding) in enumerate(zip(unique_concepts, new_embeddings), 1):
                # Print progress every 500 concepts
                if idx % 500 == 0 or idx == total_concepts:
                    print(f"  Processing progress: {idx}/{total_concepts} ({idx/total_concepts*100:.1f}%)")
                # Check if it is similar to existing concepts
                similar_concept = self._is_similar_to_existing(new_concept, new_embedding)
        
                if similar_concept:
                    # Found a similar concept, create a mapping
                    self.concept_mapping[new_concept] = similar_concept
                    concept_targets[new_concept] = similar_concept
                    embedding_duplicates += 1
                else:
                    # Brand new concept, add to the graph and create a self-mapping
                    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

        # Add edges (connecting to all relevant center concepts)
        for concept in all_new_concepts:
            target_concept = concept_targets[concept]

            for center_concept in concept_to_centers[concept]:
                # Ensure the center concept exists in the graph
                if center_concept in self.graph:
                    # Bidirectional connection
                    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"Deduplication complete: Nodes added {nodes_added}, Edges added {edges_added//2}, Deduplicated concepts {embedding_duplicates}")

        return nodes_added, edges_added // 2, embedding_duplicates, nodes_added / num_all_new_concepts  # Undirected graph, so divide edge count by 2

    def get_graph_stats(self) -> Dict[str, int]:
        """Gets graph statistics"""
        node_count = len(self.graph)
        edge_count = sum(len(neighbors) for neighbors in self.graph.values()) // 2
        return {"nodes": node_count, "edges": edge_count}

# Import JSON repair library
try:
    from json_repair import repair_json
    HAS_JSONREPAIR = True
    print("✓ jsonrepair library loaded, JSON repair feature enabled")
except ImportError:
    HAS_JSONREPAIR = False
    print("⚠ jsonrepair library not installed, will use basic repair strategy. Run 'pip install jsonrepair' to enable advanced JSON repair")
    def repair_json(text):
        return text

class ResponseValidator:
    """Response Validator"""
    
    @staticmethod
    def validate_json_response(response_text: str, expected_keys: List[str]) -> Dict:
        """
        Validates if the content returned by the API is valid JSON, including pre-processing for fault tolerance.
        
        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:
            # Pre-processing and cleaning
            text = response_text.strip()
    
            # 1. Handle markdown code blocks ```json...``` or ```...```
            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. Handle quoted wrappers '...' or "..."
            if (text.startswith("'") and text.endswith("'")) or (text.startswith('"') and text.endswith('"')):
                text = text[1:-1]
    
            # 3. Remove leading/trailing backticks
            text = text.strip('`').strip()
    
            # 4. Find the JSON part - from the first { to the last }
            json_match = re.search(r'(\{.*\})', text, re.DOTALL)
            if json_match:
                text = json_match.group(1)
    
            # First attempt: direct parsing
            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)}")
        
                # Second attempt: parse after repairing with 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
    
            # Check if it conforms to the expected structure
            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
            }

# Iteration code for graph augmentation based on seed concepts

@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:
        """get prompt for concept expansion"""
        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_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
            )

        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"]
    
            # Check if it is "no new concept"
            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=[],  # Empty list, indicating no new concepts
                    returned_center=returned_center
                )
    
            # Process new concepts normally - remove the previous filtering logic and let embedding remove duplicates
            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"Starting batch concept augmentation for {total_concepts} concepts")
        print(f"Max concurrency: {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"Processing complete, time taken: {batch_end_time - batch_start_time:.2f} seconds")
        return batch_results

    def _process_batch(self, batch_concepts: List[str], adjacency_dict: Dict[str, List[str]], 
                       max_workers: int) -> Dict:
        """Processes a batch"""
        batch_results = {}

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # Submit tasks
            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)

            # Collect results
            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

                    # Simple status display
                    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: {completed}/{len(batch_concepts)} (Success: {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"  Exception: {concept_id} - {str(e)}")

        return batch_results

    def _save_results(self, batch_results: Dict, start_time: float):
        """Saves the results"""
        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)

        # Statistics
        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

        # Count the number of generated concepts
        total_new_concepts = sum(len(r.new_concepts) for r in batch_results.values() 
                                 if r.status == "success" and r.new_concepts)
    
        # Count the number of skipped concepts (no new concepts added)
        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)

        # Construct save data
        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
        }

        # Save to file
        with open(filepath, 'wb') as f:
            pickle.dump(save_data, f)

        print(f"  Results saved to: {filename}")
        print(f"  Success: {success_count}/{total_count} ({success_count/total_count*100:.1f}%)")
        print(f"  Total new concepts: {total_new_concepts}")
        print(f"  Skipped concepts: {skipped_concepts}")

        return filepath

class ConceptGraph:
  
    @classmethod
    def from_graph_dict(cls, graph_dict: Dict[str, List[str]], concept_mapping, model, similarity_threshold):
        """
        Reconstructs a ConceptGraph from a saved graph dictionary
        Args:
            graph_dict: The saved adjacency dictionary
            model: An instance of the SentenceTransformer model
            similarity_threshold: The threshold for judging concepts as similar
            concept_mapping: The mapping of all previously seen similar concepts, e.g., {'hearts':'heart'}
        Returns:
            An instance of ConceptGraph
        """
        if model is None:
            raise ValueError("The model parameter cannot be None, please load the model using load_embedding_model() first")

        # Create an instance without initializing
        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

        # Recalculate embeddings for all concepts and establish self-mapping
        all_concepts = list(graph_dict.keys())
        if all_concepts:
            print(f"Recalculating embeddings for {len(all_concepts)} concepts...")
            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  # Establish self-mapping
    
            print("ConceptGraph reconstruction complete")

        return instance

    def __init__(self, seed_concepts: List[str], model, similarity_threshold):
        """
        Initializes the graph from seed concepts and builds the embedding library
        Args:
            seed_concepts: A list of seed concepts that have been externally deduplicated
            model: An instance of the SentenceTransformer model (required)
            similarity_threshold: The threshold for judging concepts as similar
        """
        if model is None:
            raise ValueError("The model parameter cannot be None, please load the model using load_embedding_model() first")
    
        self.model = model
        self.graph = {}
        self.concept_embeddings = {}  # Maintains the mapping of concept -> embedding
        self.concept_mapping = {}     # Add a concept mapping table
        self.similarity_threshold = similarity_threshold  # embedding similarity threshold

        # Clean the seed concepts
        cleaned_seeds = [concept.strip() for concept in seed_concepts if concept.strip()]

        print(f"Calculating embeddings for {len(cleaned_seeds)} seed concepts...")

        # Batch calculate embeddings
        if cleaned_seeds:
            seed_embeddings = self.model.encode(cleaned_seeds)
    
            # Build the initial graph, embedding library, and mapping table
            for concept, embedding in zip(cleaned_seeds, seed_embeddings):
                self.graph[concept] = []
                self.concept_embeddings[concept] = embedding
                self.concept_mapping[concept] = concept  # Establish self-mapping

        print(f"Concept graph initialization complete, number of seed concepts: {len(cleaned_seeds)}")

    def _get_target_concept(self, concept: str) -> Optional[str]:
        """Unified concept mapping lookup"""
        return self.concept_mapping.get(concept)

    def _is_similar_to_existing(self, new_concept: str, new_embedding: np.ndarray) -> Optional[str]:
        """
        Checks if a new concept is similar to existing ones
        Returns:
            If similar, returns the similar existing concept; otherwise, returns None
        """
        if not self.concept_embeddings:
            return None
    
        # Calculate similarity with all existing concepts
        existing_concepts = list(self.concept_embeddings.keys())
        existing_embeddings = np.array([self.concept_embeddings[concept] for concept in existing_concepts])

        # Calculate cosine similarity
        similarities = self.model.similarity([new_embedding], existing_embeddings)[0]

        # Find the most similar concept
        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]]:
        """Gets the current adjacency dictionary"""
        return self.graph.copy()

    def calculate_metrics(self, expansion_results: Dict) -> Dict[str, float]:
        """Calculates augmentation metrics"""
        # Collect all new concepts
        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())

        # Connectivity: newly generated edges (between two old nodes) / total existing edges from the previous round
        old_total_edges = sum(len(neighbors) for neighbors in self.graph.values()) // 2

        # Count newly generated edges (between two old nodes)
        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):
        """
        Updates the graph structure, using a mapping mechanism for deduplication.
        """
        nodes_added = 0
        edges_added = 0
        embedding_duplicates = 0

        # Collect all new concepts
        all_new_concepts = []
        concept_to_centers = {}  # Record the corresponding center concept for each new concept

        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"Received {num_all_new_concepts} new concepts (not deduplicated)")

        # Use mapping to quickly filter known concepts
        concepts_need_embedding = []
        concept_targets = {}  # concept -> target_concept mapping

        for concept in all_new_concepts:
            target = self._get_target_concept(concept)
            if target is not None:
                # Known concept, use mapping directly
                concept_targets[concept] = target
                if target != concept:
                    embedding_duplicates += 1
            else:
                # New concept that needs an embedding calculation
                concepts_need_embedding.append(concept)

        # Only calculate embeddings for unknown concepts
        if concepts_need_embedding:
            unique_concepts = list(set(concepts_need_embedding))
            print(f"Performing embedding-based deduplication on {len(unique_concepts)} new concepts...")
            new_embeddings = self.model.encode(unique_concepts)
    
            # Process concepts that need embeddings one by one
            total_concepts = len(unique_concepts)
            for idx, (new_concept, new_embedding) in enumerate(zip(unique_concepts, new_embeddings), 1):
                # Print progress every 500 concepts
                if idx % 500 == 0 or idx == total_concepts:
                    print(f"  Processing progress: {idx}/{total_concepts} ({idx/total_concepts*100:.1f}%)")
                # Check if it is similar to existing concepts
                similar_concept = self._is_similar_to_existing(new_concept, new_embedding)
        
                if similar_concept:
                    # Found a similar concept, create a mapping
                    self.concept_mapping[new_concept] = similar_concept
                    concept_targets[new_concept] = similar_concept
                    embedding_duplicates += 1
                else:
                    # Brand new concept, add to the graph and create a self-mapping
                    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

        # Add edges (connecting to all relevant center concepts)
        for concept in all_new_concepts:
            target_concept = concept_targets[concept]

            for center_concept in concept_to_centers[concept]:
                # Ensure the center concept exists in the graph
                if center_concept in self.graph:
                    # Bidirectional connection
                    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"Deduplication complete: Nodes added {nodes_added}, Edges added {edges_added//2}, Deduplicated concepts {embedding_duplicates}")

        return nodes_added, edges_added // 2, embedding_duplicates, nodes_added / num_all_new_concepts  # Undirected graph, so divide edge count by 2

    def get_graph_stats(self) -> Dict[str, int]:
        """Gets graph statistics"""
        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"):
    """
    Loads the embedding model
    Args:
        model_name: Model name
    Returns:
        An instance of the SentenceTransformer model
    """
    print(f"Loading embedding model: {model_name}")
    model = SentenceTransformer(model_name, trust_remote_code=True)
    print("Model loaded successfully")
    return model

def extract_seed_concepts(results):
    """Extracts seed concepts from batch processing results"""
    all_concepts = []
    
    for result in results.values():
        if result.status == "success" and result.extracted_concepts:
            all_concepts.extend(result.extracted_concepts)
    
    # Trim whitespace and deduplicate
    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):
    """Runs a single concept augmentation iteration"""
    # Initialize API client and processors
    api_client = APIClient(
        api_url=api_url,
        api_key=api_key,
        model_name=model_name
    )
    
    expander = ConceptExpander(api_client)
    batch_expander = BatchConceptExpander(expander)
    
    # Get the current adjacency dictionary
    current_adjacency = concept_graph.get_current_adjacency()
    
    # Batch augment concepts
    expansion_results = batch_expander.expand_concepts_batch(
        adjacency_dict=current_adjacency,
        max_workers=max_workers
    )
    
    # Calculate metrics
    metrics = concept_graph.calculate_metrics(expansion_results)
    
    # Update graph - using embedding-based deduplication
    nodes_added, edges_added, embedding_duplicates, concept_add_rate = concept_graph.update_graph(expansion_results)

    # Get updated graph statistics
    graph_stats = concept_graph.get_graph_stats()
    
    # Count the number of skipped concepts
    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 results
    print(f"\n=== Iteration Complete ===")
    print(f"Concept Add Rate: {concept_add_rate:.3f}")
    print(f"Concept Connectivity Rate: {metrics['connectivity_rate']:.3f}")
    print(f"Nodes at end of iteration: {graph_stats['nodes']}")
    print(f"Edges at end of iteration: {graph_stats['edges']}")
    print(f"Nodes added this round: {nodes_added}")
    print(f"Edges added this round: {edges_added}")
    print(f"Skipped concepts: {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
    }