# 设置DashScope API密钥
dashscope.api_key = os.environ.get("DASHSCOPE_API_KEY")
# 模拟医疗文本数据
medical_texts = [
    "糖尿病是一种慢性疾病，其特征是高血糖。常见症状包括多饮、多尿和多食。",
    # ... 更多医疗文本
]
# 定义医疗实体类型
MEDICAL_ENTITY_TYPES = {
    "DISEASE": "疾病",
    "SYMPTOM": "症状", 
    "DRUG": "药物",
    "TREATMENT": "治疗方法",
    "ANATOMY": "解剖部位"
}

class MedicalKnowledgeGraph:
    def __init__(self):
        self.graph = nx.DiGraph()  # 创建有向图
        self.entities = set()      # 存储所有实体
        self.relations = []        # 存储所有关系
        self.entity_types = {}     # 存储实体类型映射

def extract_entities_with_llm(self, text: str) -> List[Tuple[str, str]]:
    # 构建详细的提示词
    prompt = f"""
请从以下医疗文本中识别出所有医疗实体，并按照以下类型分类：
- DISEASE: 疾病
- SYMPTOM: 症状
# ... 其他类型
文本内容："{text}"
请以JSON格式输出结果...
"""

def extract_relations_with_llm(self, text: str, entities: List[Tuple[str, str]]) -> List[Tuple[str, str, str]]:
    # 构建实体列表字符串
    entities_str = ", ".join([f"{ent}({typ})" for ent, typ in entities])
    
    # 构建关系提取提示词
    prompt = f"""
请分析以下医疗文本，识别其中医疗实体之间的关系：
文本内容："{text}"
文本中包含的实体：{entities_str}
请识别这些实体之间的关系，使用以下关系类型：
- TREATS: 治疗关系
- CAUSES: 引起关系
# ... 其他关系类型
"""

def build_from_texts(self, texts: List[str]):
    for i, text in enumerate(texts):
        # 提取实体
        entities = self.extract_entities_with_llm(text)
        for entity, type_ in entities:
            self.entities.add(entity)
            self.entity_types[entity] = type_
        
        # 提取关系
        relations = self.extract_relations_with_llm(text, entities)
        self.relations.extend(relations)
        
        # 添加到图结构
        for entity, type_ in entities:
            self.graph.add_node(entity, type=type_)
            
        for src, rel, dst in relations:
            self.graph.add_edge(src, dst, relationship=rel)

def visualize(self, title="医疗知识图谱"):
    # 根据节点类型设置颜色
    node_colors = []
    for node in self.graph.nodes():
        node_type = self.entity_types.get(node, "UNKNOWN")
        if node_type == "DISEASE":
            node_colors.append("lightcoral")
        # ... 其他类型颜色设置
    
    # 创建布局和绘制
    pos = nx.spring_layout(self.graph, k=2, iterations=50)
    nx.draw_networkx_nodes(self.graph, pos, node_color=node_colors, node_size=2000, alpha=0.9)
    # ... 更多绘制代码

def query(self, entity: str, relation: str = None):
    """查询知识图谱"""
    results = []
    if relation:
        # 查询特定关系
        for src, rel, dst in self.relations:
            if src == entity and rel == relation:
                results.append(dst)
            elif dst == entity and rel == relation:
                results.append(src)
    else:
        # 查询所有相关关系
        for src, rel, dst in self.relations:
            if src == entity or dst == entity:
                results.append((src, rel, dst))
    
    return results
def to_json(self, filepath: str = None):
    """将知识图谱导出为JSON格式"""
    kg_data = {
        "entities": list(self.entities),
        "relations": [{"source": s, "relationship": r, "target": t} for s, r, t in self.relations],
        "entity_types": self.entity_types
    }

class MedicalQA:
    def __init__(self, knowledge_graph):
        self.kg = knowledge_graph
    
    def answer_question(self, question: str) -> str:
        # 首先尝试从知识图谱中直接找到答案
        direct_answer = self._answer_from_kg(question)
        if direct_answer:
            return direct_answer
        
        # 如果知识图谱中没有直接答案，使用大模型生成答案
        return self._answer_with_llm(question)

if __name__ == "__main__":
    # 构建基础知识图谱
    basic_kg = build_medical_kg_with_dashscope()
    
    # 构建复杂知识图谱
    complex_kg = build_complex_medical_kg()
    
    # 创建问答系统并测试
    qa_system = MedicalQA(basic_kg)
    test_questions = [
        "糖尿病有什么症状？",
        "阿司匹林可以治疗什么病？",
        "高血压应该用什么药？"
    ]
    
    for question in test_questions:
        answer = qa_system.answer_question(question)
        print(f"回答: {answer}")
    
    # 知识图谱分析
    degree_centrality = nx.degree_centrality(basic_kg.graph)
    # ... 更多分析代码

# 知识图谱构建示例：医疗领域（使用DashScope Generation模型）
import pandas as pd
import numpy as np
import re
import json
from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
matplotlib.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号
import networkx as nx
from typing import List, Dict, Tuple, Set
import dashscope
from dashscope import Generation
import warnings
warnings.filterwarnings('ignore')
import os
import logging
import time
import datetime
# 配置日志
log_file = f"medical_kg_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.FileHandler(log_file, encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)
# 设置DashScope API密钥
dashscope.api_key = os.environ.get("DASHSCOPE_API_KEY")  # 请替换为您的实际API密钥
# 模拟医疗文本数据
medical_texts = [
    "糖尿病是一种慢性疾病，其特征是高血糖。常见症状包括多饮、多尿和多食。",
    "高血压患者需要定期服用降压药，如氨氯地平或缬沙坦。",
    "阿司匹林可用于预防心脏病发作，但可能引起胃肠道出血。",
    "肺炎通常由细菌感染引起，症状包括咳嗽、发热和胸痛。治疗常用抗生素如阿莫西林。",
    "心肌梗死是冠状动脉阻塞导致的心肌缺血坏死，主要症状为胸痛和呼吸困难。",
    "2型糖尿病患者可能需要使用二甲双胍或胰岛素来控制血糖水平。",
    "哮喘患者通常使用吸入性糖皮质激素如布地奈德来控制症状。",
    "抑郁症的治疗包括心理治疗和抗抑郁药物，如氟西汀和舍曲林。"
]
# 定义医疗实体类型
MEDICAL_ENTITY_TYPES = {
    "DISEASE": "疾病",
    "SYMPTOM": "症状", 
    "DRUG": "药物",
    "TREATMENT": "治疗方法",
    "ANATOMY": "解剖部位"
}
class MedicalKnowledgeGraph:
    def __init__(self):
        self.graph = nx.DiGraph()
        self.entities = set()
        self.relations = []
        self.entity_types = {}
        
    def extract_entities_with_llm(self, text: str) -> List[Tuple[str, str]]:
        """使用Qwen大模型提取医疗实体"""
        # print(f"开始从文本提取实体: {text[:50]}...")
        # 构建实体识别提示词
        prompt = f"""
请从以下医疗文本中识别出所有医疗实体，并按照以下类型分类：
- DISEASE: 疾病
- SYMPTOM: 症状
- DRUG: 药物
- TREATMENT: 治疗方法
- ANATOMY: 解剖部位
文本内容："{text}"
请以JSON格式输出结果，包含一个"entities"数组，每个实体对象包含"name"和"type"两个字段。
输出示例：
{{
  "entities": [
    {{"name": "糖尿病", "type": "DISEASE"}},
    {{"name": "高血糖", "type": "SYMPTOM"}}
  ]
}}
"""
        
        try:
            # 调用Qwen的大模型
            start_time = time.time()
            response = Generation.call(
                model='qwen-max',
                prompt=prompt,
                result_format='text'
            )
            elapsed_time = time.time() - start_time
            print(f"调用Qwen大模型进行实体识别，耗时: {elapsed_time:.2f}秒")
            
            entities = []
            if response and response.output:
                # 尝试从响应中提取JSON
                # result_text = response.output.choices[0].message.content
                result_text = response.output.text
                
                # 查找JSON部分
                json_match = re.search(r'\{.*\}', result_text, re.DOTALL)
                if json_match:
                    result_json = json.loads(json_match.group())
                    if 'entities' in result_json:
                        for entity in result_json['entities']:
                            if all(k in entity for k in ['name', 'type']):
                                entities.append((entity['name'], entity['type']))
            
            print(f"成功提取到 {len(entities)} 个实体：{entities}")
            return entities
        except Exception as e:
            logger.error(f"实体识别错误: {e}")
            return []
    
    def extract_relations_with_llm(self, text: str, entities: List[Tuple[str, str]]) -> List[Tuple[str, str, str]]:
        """使用Qwen大模型进行关系抽取"""
        # print(f"开始从文本提取实体关系: {text[:50]}...")
        if not entities or len(entities) < 2:
            print("实体数量不足，无法提取关系")
            return []
            
        relations = []
        
        # 构建实体列表字符串
        entities_str = ", ".join([f"{ent}({typ})" for ent, typ in entities])
        
        # 构建提示词
        prompt = f"""
请分析以下医疗文本，识别其中医疗实体之间的关系：
文本内容："{text}"
文本中包含的实体：{entities_str}
请识别这些实体之间的关系，使用以下关系类型：
- TREATS: 治疗关系，如"药物治疗疾病"
- CAUSES: 引起关系，如"疾病引起症状" 
- HAS_SYMPTOM: 有症状关系，如"疾病有症状"
- IS_A: 是一种关系，如"疾病A是疾病B的一种"
- USES: 使用关系，如"治疗使用药物"
请以JSON格式输出结果，包含一个"relations"数组，每个关系对象包含"source", "relationship", "target"三个字段。
输出示例：
{{
  "relations": [
    {{"source": "糖尿病", "relationship": "HAS_SYMPTOM", "target": "高血糖"}},
    {{"source": "阿司匹林", "relationship": "TREATS", "target": "心脏病"}}
  ]
}}
"""
        
        try:
            # 调用DashScope的大模型
            # print("调用DashScope API进行关系抽取...")
            start_time = time.time()
            response = Generation.call(
                model='qwen-max',
                prompt=prompt,
                result_format='text'
            )
            elapsed_time = time.time() - start_time
            print(f"调用Qwen大模型进行关系抽取，耗时: {elapsed_time:.2f}秒")
            logging.debug(f"API响应: {response}")
            if response and response.output:
                # 尝试从响应中提取JSON
                # result_text = response.output.choices[0].message.content
                result_text = response.output.text
                
                # 查找JSON部分
                json_match = re.search(r'\{.*\}', result_text, re.DOTALL)
                if json_match:
                    result_json = json.loads(json_match.group())
                    if 'relations' in result_json:
                        for rel in result_json['relations']:
                            if all(k in rel for k in ['source', 'relationship', 'target']):
                                relations.append((rel['source'], rel['relationship'], rel['target']))
            
            print(f"成功提取到 {len(relations)} 个关系：{relations}")
            print("-" * 60)
            return relations
        except Exception as e:
            logger.error(f"关系抽取错误: {e}")
            return []
    
    def build_from_texts(self, texts: List[str]):
        """从文本集合构建知识图谱"""
        print(f"开始从 {len(texts)} 条文本构建知识图谱")
        print("-" * 60)
        start_time = time.time()
        for i, text in enumerate(texts):
            print(f"处理文本 {i+1}/{len(texts)}: {text[:50]}...")
            
            # 提取实体
            entities = self.extract_entities_with_llm(text)
            for entity, type_ in entities:
                self.entities.add(entity)
                self.entity_types[entity] = type_
            
            # 提取关系
            relations = self.extract_relations_with_llm(text, entities)
            self.relations.extend(relations)
            
            # 添加到图
            for entity, type_ in entities:
                self.graph.add_node(entity, type=type_)
                
            for src, rel, dst in relations:
                self.graph.add_edge(src, dst, relationship=rel)
                
        total_time = time.time() - start_time
        print(f"知识图谱构建完成，总耗时: {total_time:.2f}秒")
        print(f"共提取 {len(self.entities)} 个实体, {len(self.relations)} 个关系")
    
    def visualize(self, title="医疗知识图谱"):
        """可视化知识图谱"""
        print(f"开始可视化知识图谱: {title}")
        plt.figure(figsize=(14, 10))
        
        # 根据节点类型设置颜色
        node_colors = []
        for node in self.graph.nodes():
            node_type = self.entity_types.get(node, "UNKNOWN")
            if node_type == "DISEASE":
                node_colors.append("lightcoral")
            elif node_type == "SYMPTOM":
                node_colors.append("lightblue")
            elif node_type == "DRUG":
                node_colors.append("lightgreen")
            elif node_type == "TREATMENT":
                node_colors.append("khaki")
            elif node_type == "ANATOMY":
                node_colors.append("violet")
            else:
                node_colors.append("gray")
        
        # 创建布局
        pos = nx.spring_layout(self.graph, k=2, iterations=50)
        
        # 绘制节点
        nx.draw_networkx_nodes(self.graph, pos, node_color=node_colors, node_size=2000, alpha=0.9)
        
        # 绘制边
        nx.draw_networkx_edges(self.graph, pos, width=1.5, alpha=0.7, edge_color="gray")
        
        # 绘制节点标签
        nx.draw_networkx_labels(self.graph, pos, font_size=10, font_family="SimHei")
        
        # 绘制边标签
        edge_labels = nx.get_edge_attributes(self.graph, 'relationship')
        nx.draw_networkx_edge_labels(self.graph, pos, edge_labels=edge_labels, font_size=8)
        
        # 添加图例
        legend_elements = [
            plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='lightcoral', markersize=10, label='疾病'),
            plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='lightblue', markersize=10, label='症状'),
            plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='lightgreen', markersize=10, label='药物'),
            plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='khaki', markersize=10, label='治疗方法'),
            plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='violet', markersize=10, label='解剖部位')
        ]
        plt.legend(handles=legend_elements, loc='best')
        
        plt.title(title, fontsize=16)
        plt.axis('off')
        plt.tight_layout()
        print("知识图谱可视化完成，准备显示图表")
        plt.show()
    
    def to_json(self, filepath: str = None):
        """将知识图谱导出为JSON格式"""
        print(f"开始导出知识图谱到JSON{f': {filepath}' if filepath else ''}")
        kg_data = {
            "entities": list(self.entities),
            "relations": [{"source": s, "relationship": r, "target": t} for s, r, t in self.relations],
            "entity_types": self.entity_types
        }
        
        if filepath:
            try:
                with open(filepath, 'w', encoding='utf-8') as f:
                    json.dump(kg_data, f, ensure_ascii=False, indent=2)
                print(f"知识图谱已成功导出到: {filepath}")
            except Exception as e:
                logger.error(f"导出知识图谱到JSON文件时出错: {e}")
        
        return kg_data
    
    def query(self, entity: str, relation: str = None):
        """查询知识图谱"""
        results = []
        if relation:
            for src, rel, dst in self.relations:
                if src == entity and rel == relation:
                    results.append(dst)
                elif dst == entity and rel == relation:
                    results.append(src)
        else:
            for src, rel, dst in self.relations:
                if src == entity or dst == entity:
                    results.append((src, rel, dst))
        
        return results
# 构建医疗知识图谱
def build_medical_kg_with_dashscope():
    print("开始使用Qwen大模型构建医疗知识图谱...")
    print("=" * 60)
    kg = MedicalKnowledgeGraph()
    kg.build_from_texts(medical_texts)
    
    print(f"提取到 {len(kg.entities)} 个实体")
    print("实体示例:", list(kg.entities)[:5])
    print(f"提取到 {len(kg.relations)} 个关系")
    print("关系示例:", kg.relations[:3] if kg.relations else "无")
    
    # 可视化知识图谱
    kg.visualize("基于Qwen大模型的医疗知识图谱")
    
    # 导出为JSON
    kg_data = kg.to_json("medical_kg_dashscope.json")
    print("知识图谱已导出到 medical_kg_dashscope.json")
    
    # 示例查询
    print("\n示例查询:")
    print("查询与'糖尿病'相关的所有关系:")
    diabetes_relations = kg.query("糖尿病")
    for rel in diabetes_relations:
        print(f"  {rel}")
    
    print("\n查询'阿司匹林'治疗哪些疾病:")
    aspirin_treats = kg.query("阿司匹林", "TREATS")
    for disease in aspirin_treats:
        print(f"  {disease}")
    
    return kg
# 高级功能：基于知识图谱的问答系统
class MedicalQA:
    def __init__(self, knowledge_graph):
        self.kg = knowledge_graph
    
    def extract_entities_with_llm(self, text: str) -> List[Tuple[str, str]]:
        """使用LLM从文本中提取实体"""
        # 构建实体识别提示词
        prompt = f"""
请从以下文本中识别出所有医疗实体，并按照以下类型分类：
- DISEASE: 疾病
- SYMPTOM: 症状
- DRUG: 药物
- TREATMENT: 治疗方法
- ANATOMY: 解剖部位
文本内容："{text}"
请以JSON格式输出结果，包含一个"entities"数组，每个实体对象包含"name"和"type"两个字段。
输出示例：
{{
  "entities": [
    {{"name": "糖尿病", "type": "DISEASE"}},
    {{"name": "高血糖", "type": "SYMPTOM"}}
  ]
}}
"""
        
        try:
            # 调用DashScope的大模型
            start_time = time.time()
            response = Generation.call(
                model='qwen-max',
                prompt=prompt,
                result_format='text'
            )
            elapsed_time = time.time() - start_time
            print(f"调用Qwen大模型进行实体识别，耗时: {elapsed_time:.2f}秒")
            
            entities = []
            if response and response.output:
                # 尝试从响应中提取JSON
                result_text = response.output.text
                
                # 查找JSON部分
                json_match = re.search(r'\{.*\}', result_text, re.DOTALL)
                if json_match:
                    result_json = json.loads(json_match.group())
                    if 'entities' in result_json:
                        for entity in result_json['entities']:
                            if all(k in entity for k in ['name', 'type']):
                                entities.append((entity['name'], entity['type']))
            
            print(f"成功提取到 {len(entities)} 个实体：{entities}")
            return entities
        except Exception as e:
            logger.error(f"实体识别错误: {e}")
            return []
    
    def answer_question(self, question: str) -> str:
        """基于知识图谱回答问题"""
        # 首先尝试从知识图谱中直接找到答案
        direct_answer = self._answer_from_kg(question)
        if direct_answer:
            return direct_answer
        
        # 如果知识图谱中没有直接答案，使用大模型生成答案
        return self._answer_with_llm(question)
    
    def _answer_from_kg(self, question: str) -> str:
        """从知识图谱中直接提取答案"""
        # 提取问题中的实体
        entities = self.extract_entities_with_llm(question)
        if not entities:
            return None
            
        # 查找与实体相关的关系
        answers = []
        for entity, _ in entities:
            relations = self.kg.query(entity)
            for src, rel, dst in relations:
                if rel == "TREATS" and src == entity:
                    answers.append(f"{entity}可用于治疗{dst}")
                elif rel == "TREATS" and dst == entity:
                    answers.append(f"{src}可用于治疗{entity}")
                elif rel == "HAS_SYMPTOM" and src == entity:
                    answers.append(f"{entity}的症状包括{dst}")
                elif rel == "CAUSES" and src == entity:
                    answers.append(f"{entity}可能引起{dst}")
        
        if answers:
            return "。".join(answers)
        return None
    
    def _answer_with_llm(self, question: str) -> str:
        """使用大模型生成答案"""
        # 构建知识上下文
        context_entities = self.extract_entities_with_llm(question)
        context = "相关知识：\n"
        for entity, _ in context_entities:
            relations = self.kg.query(entity)
            for src, rel, dst in relations:
                context += f"- {src} {rel} {dst}\n"
        
        prompt = f"""
你是一个专业的医疗问答助手。请基于以下知识回答问题。
{context}
问题：{question}
请提供准确、专业的回答，并注明信息来源是基于医学知识图谱。
"""
        
        try:
            response = Generation.call(
                model='qwen-max',
                prompt=prompt,
                result_format='text'
            )
            
            if response and response.output:
                return response.output.text
            return "抱歉，我无法回答这个问题。"
        except Exception as e:
            return f"生成答案时出错: {str(e)}"
# 主函数
if __name__ == "__main__":
    # 构建基础知识图谱
    print("=" * 60)
    basic_kg = build_medical_kg_with_dashscope()
    
    # 创建问答系统
    print("=" * 60)
    print("创建医疗问答系统...")
    qa_system = MedicalQA(basic_kg)
    
    # 测试问答
    test_questions = [
        "糖尿病有什么症状？",
        "阿司匹林可以治疗什么病？",
        "高血压应该用什么药？"
    ]
    
    for question in test_questions:
        print(f"\n问题: {question}")
        answer = qa_system.answer_question(question)
        print(f"回答: {answer}")
    
    # 知识图谱分析
    print("=" * 60)
    print("知识图谱分析:")
    
    # 计算节点中心度
    degree_centrality = nx.degree_centrality(basic_kg.graph)
    top_nodes = sorted(degree_centrality.items(), key=lambda x: x[1], reverse=True)[:5]
    print("\n基础图谱中心度最高的节点:")
    for node, centrality in top_nodes:
        print(f"{node}: {centrality:.3f}")
    
    # 查找最常见的疾病
    diseases = [node for node, type_ in basic_kg.entity_types.items() if type_ == "DISEASE"]
    print(f"\n识别出的疾病: {', '.join(diseases)}")
    
    # 查找最常见的药物
    drugs = [node for node, type_ in basic_kg.entity_types.items() if type_ == "DRUG"]
    print(f"识别出的药物: {', '.join(drugs)}")