# ============================================================
# 环境准备：安装依赖
# ============================================================
# pip install openai pandas biopython requests

import os
import json
import pandas as pd
from openai import OpenAI

# 初始化客户端（支持OpenAI、DeepSeek、Kimi等兼容API）
# 白话：就像登录一个网站，需要API密钥作为"通行证"
client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),  # 从环境变量读取密钥
    base_url="https://api.openai.com/v1"       # 如果用DeepSeek：https://api.deepseek.com/v1
)

def ask_llm(prompt, system_msg="你是一个资深生物信息学工程师。", model="gpt-4o"):
    """
    通用LLM调用函数
    参数：
        prompt: 你的问题或指令（用户消息）
        system_msg: 系统提示词，定义LLM的角色
        model: 模型名称（gpt-4o/deepseek-chat/moonshot-v1-8k）
    返回：
        LLM的回复文本
    """
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system_msg},  # 角色设定
            {"role": "user", "content": prompt}          # 你的问题
        ],
        temperature=0.3  # 低温度=更确定性的输出，适合代码和科学分析
    )
    return response.choices[0].message.content  # 提取回复文本

# ============================================================
# 实战1：LLM生成生信分析脚本
# 思路：描述需求 → LLM生成代码 → 你审查修改 → 测试运行
# ============================================================

def generate_bioinfo_script(requirement):
    """
    输入自然语言需求，让LLM生成生信分析脚本
    参数：
        requirement: 需求描述，越具体越好
    返回：
        生成的代码字符串
    """
    # 关键：system提示词要限定"生信+Python"的上下文
    system_prompt = """你是一个资深生物信息学工程师，擅长Python/R/Shell脚本。
要求：
1. 代码每一行都加中文注释
2. 使用标准生信库（biopython, scikit-bio, pandas等）
3. 包含输入参数检查和错误处理
4. 输出格式化的结果文件
5. 代码末尾加上"使用说明"注释块"""

    user_prompt = f"""请根据以下需求生成完整的Python分析脚本：

【需求描述】
{requirement}

【输入数据格式】
OTU丰度表（CSV，行=样本，列=OTU，第一列是样本ID）

【期望输出】
1. Alpha多样性指标表（Shannon, Simpson, Chao1, Observed OTUs）
2. 分组比较的统计检验结果
3. 箱线图可视化

请直接输出完整可运行的Python脚本。"""

    # 调用LLM
    code = ask_llm(user_prompt, system_prompt)
    return code


# ---- 使用示例 ----
requirement = """
分析16S扩增子测序数据的alpha多样性：
- 输入：OTU丰度表(otu_table.csv)和分组文件(metadata.csv，含group列)
- 计算Shannon、Simpson、Chao1、Observed OTUs四个指标
- 对T2D组和Health组做Wilcoxon秩和检验
- 画分组箱线图，保存为PDF
"""

generated_code = generate_bioinfo_script(requirement)
print("=" * 60)
print("【LLM生成的代码 - 需要人工审查！】")
print("=" * 60)
print(generated_code)

# ---- 审查清单（你必须做的）----
# □ 检查导入的库是否都安装了
# □ 检查统计方法是否正确（Wilcoxon适用于非正态小样本）
# □ 检查文件路径是否和该数据一致
# □ 检查可视化参数（字体大小、颜色方案）
# □ 用小数据集测试运行一遍

# ============================================================
# 审查步骤：让LLM帮你review自己生成的代码
# ============================================================

def review_code(code_to_review):
    """
    让LLM审查生信代码，找出潜在问题
    """
    review_prompt = f"""请以代码审查者的身份检查以下生信分析脚本，指出：
1. 潜在的bug或逻辑错误
2. 统计方法是否合理
3. 缺少的边界情况处理
4. 可以优化的地方
5. 生物学上是否有不合理的假设

```python
{code_to_review}

review_result = ask_llm(review_prompt)
return review_result

---

## 四、实战2：用LLM解读分析结果

### 场景
DESeq2跑完了，得到一个差异表达基因列表，怎么快速写出解读报告？

### 完整代码

```python
# ============================================================
# 实战2：LLM自动解读DESeq2差异分析结果
# 输入：DESeq2结果CSV → 输出：中文解读报告
# ============================================================

import pandas as pd

def interpret_deseq2_results(results_file, top_n=20):
    """
    读取DESeq2结果，让LLM生成解读报告
    参数：
        results_file: DESeq2输出的CSV文件路径
        top_n: 取top N个显著基因送给LLM解读
    返回：
        解读报告文本
    """
    # 第一步：读取并筛选显著差异基因
    df = pd.read_csv(results_file)

    # 筛选条件：padj < 0.05 且 |log2FC| > 1（常用阈值）
    sig_genes = df[
        (df['padj'] < 0.05) &              # 校正后p值小于0.05
        (df['log2FoldChange'].abs() > 1)    # 倍数变化大于2倍
    ].sort_values('padj')                   # 按显著性排序

    # 取top N基因
    top_genes = sig_genes.head(top_n)

    # 第二步：构建LLM输入
    # 将数据转为文本表格格式，方便LLM理解
    genes_text = top_genes[['gene_id', 'log2FoldChange', 'padj', 'baseMean']].to_string(index=False)

    # 统计摘要信息
    summary = f"""
总基因数：{len(df)}
显著上调基因数（log2FC>1, padj<0.05）：{len(sig_genes[sig_genes['log2FoldChange'] > 1])}
显著下调基因数（log2FC<-1, padj<0.05）：{len(sig_genes[sig_genes['log2FoldChange'] < -1])}
"""

    # 第三步：让LLM生成解读
    interpretation_prompt = f"""你是一个分子生物学和生物信息学专家。请解读以下DESeq2差异表达分析结果。

【实验设计】
比较组：2型糖尿病(T2D) vs 健康对照(Health)
样本来源：肠道菌群宏转录组

【统计摘要】
{summary}

【Top {top_n}显著差异基因/转录本】
{genes_text}

请生成一份结构化的解读报告，包含：
1. **总体概述**：差异表达的整体模式（上调多还是下调多，意味着什么）
2. **关键发现**：top基因的已知功能及其与T2D的潜在关联
3. **通路推测**：这些基因可能涉及哪些生物学通路
4. **研究意义**：这些发现对T2D研究有什么启示
5. **局限性说明**：需要注意的数据局限

注意：如果你不确定某个基因的功能，请明确标注"需验证"，不要编造。"""

    report = ask_llm(interpretation_prompt)
    return report, sig_genes


# ---- 使用示例 ----
# report, sig_df = interpret_deseq2_results("deseq2_results.csv")
# print(report)

# ---- 保存报告 ----
# with open("interpretation_report.md", "w", encoding="utf-8") as f:
#     f.write(report)

def batch_interpret(results_dir):
    """
    批量解读多个DESeq2结果文件
    参数：
        results_dir: 包含多个DESeq2结果CSV的目录
    """
    import glob
    results_files = glob.glob(f"{results_dir}/*_deseq2.csv")

    all_reports = {}
    for f in results_files:
        comparison_name = os.path.basename(f).replace("_deseq2.csv", "")
        print(f"正在解读: {comparison_name}...")
        report, _ = interpret_deseq2_results(f)
        all_reports[comparison_name] = report

    # 合并所有报告
    final_report = "# 差异表达分析综合解读报告\n\n"
    for name, report in all_reports.items():
        final_report += f"## {name}\n\n{report}\n\n---\n\n"

    return final_report

# ============================================================
# 实战3：LLM + PubMed文献挖掘
# 批量下载摘要 → LLM提取基因/菌属-疾病关系 → 结构化输出
# ============================================================

import requests
import time
import json

def search_pubmed(query, max_results=20):
    """
    搜索PubMed获取文章PMID列表
    参数：
        query: 搜索词（如 "Bacteroides type 2 diabetes gut"）
        max_results: 最大返回数量
    返回：
        PMID列表
    """
    # NCBI E-utilities搜索接口（免费，不需要API key也能用，但有频率限制）
    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
    search_url = f"{base_url}/esearch.fcgi"

    params = {
        "db": "pubmed",          # 数据库：PubMed
        "term": query,           # 搜索词
        "retmax": max_results,   # 最大返回数
        "retmode": "json",       # JSON格式返回
        "sort": "relevance"      # 按相关性排序
    }

    response = requests.get(search_url, params=params)
    data = response.json()
    pmids = data["esearchresult"]["idlist"]  # 提取PMID列表
    return pmids


def fetch_abstracts(pmids):
    """
    根据PMID列表获取摘要
    参数：
        pmids: PMID列表
    返回：
        摘要文本列表（每条含标题+摘要）
    """
    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
    fetch_url = f"{base_url}/efetch.fcgi"

    params = {
        "db": "pubmed",
        "id": ",".join(pmids),      # 用逗号连接多个PMID
        "rettype": "abstract",      # 返回摘要
        "retmode": "text"           # 纯文本格式
    }

    response = requests.get(fetch_url, params=params)
    # 按空行分割不同文章
    abstracts = response.text.split("\n\n\n")
    return abstracts


def extract_gene_disease_relations(abstracts, disease="type 2 diabetes"):
    """
    用LLM从摘要中提取基因/菌属-疾病关系
    参数：
        abstracts: 摘要文本列表
        disease: 目标疾病名称
    返回：
        结构化的关系列表（JSON格式）
    """
    # 每次处理5篇（避免超出token限制）
    batch_size = 5
    all_relations = []

    for i in range(0, len(abstracts), batch_size):
        batch = abstracts[i:i+batch_size]
        batch_text = "\n---\n".join(batch)

        extraction_prompt = f"""从以下PubMed摘要中提取与"{disease}"相关的微生物-疾病关系。

【摘要内容】
{batch_text}

【提取要求】
请以JSON格式输出，每条关系包含：
- "microbe": 微生物名称（属或种水平）
- "relation": 关系类型（"正相关"/"负相关"/"因果"/"未明确"）
- "mechanism": 可能的机制（简述，如"产生短链脂肪酸"）
- "evidence_strength": 证据强度（"强-临床验证"/"中-动物实验"/"弱-相关性"）
- "pmid": 来源文章的PMID（如果能识别）

只输出JSON数组，不要其他文字。如果某篇摘要没有相关信息，跳过。"""

        result = ask_llm(extraction_prompt)

        # 尝试解析JSON
        try:
            # 清理可能的markdown代码块标记
            clean_result = result.strip().strip("```json").strip("```").strip()
            relations = json.loads(clean_result)
            all_relations.extend(relations)
        except json.JSONDecodeError:
            print(f"警告：第{i//batch_size+1}批结果解析失败，跳过")
            continue

        time.sleep(1)  # 避免API频率限制

    return all_relations


# ---- 完整工作流 ----
def literature_mining_pipeline(search_query, disease, output_file="gene_disease_relations.csv"):
    """
    文献挖掘完整流程
    """
    print(f"[1/3] 搜索PubMed: {search_query}")
    pmids = search_pubmed(search_query, max_results=30)
    print(f"  找到 {len(pmids)} 篇文章")

    print("[2/3] 下载摘要...")
    abstracts = fetch_abstracts(pmids)
    print(f"  获取到 {len(abstracts)} 篇摘要")

    print("[3/3] LLM提取关系...")
    relations = extract_gene_disease_relations(abstracts, disease)
    print(f"  提取到 {len(relations)} 条关系")

    # 转为DataFrame并保存
    if relations:
        df = pd.DataFrame(relations)
        df.to_csv(output_file, index=False, encoding="utf-8-sig")
        print(f"  结果已保存: {output_file}")
        return df
    else:
        print("  未提取到有效关系")
        return pd.DataFrame()


# ---- 使用示例 ----
# df = literature_mining_pipeline(
#     search_query="gut microbiome type 2 diabetes biomarker 2023",
#     disease="type 2 diabetes",
#     output_file="t2d_microbe_relations.csv"
# )

# ============================================================
# 实战4：生信代码Debug助手
# 关键：给LLM足够的上下文，它才能给出正确答案
# ============================================================

def debug_bioinfo_error(error_message, code_snippet, context=""):
    """
    用LLM辅助debug生信代码
    参数：
        error_message: 完整的报错信息（越完整越好）
        code_snippet: 出错的代码片段
        context: 额外上下文（如数据格式、环境信息）
    返回：
        诊断结果和修复建议
    """
    debug_prompt = f"""我在做生物信息学分析时遇到了以下错误，请帮我诊断并给出修复方案。

【报错信息】

【出错代码】
```python
{code_snippet}

solution = ask_llm(debug_prompt)
return solution

### 高级用法：自动从日志文件提取错误

```python
def parse_error_from_log(log_file):
    """
    从日志文件中自动提取最近一次错误
    """
    with open(log_file, "r") as f:
        lines = f.readlines()

    # 从后往前找Traceback
    error_lines = []
    capturing = False
    for line in reversed(lines):
        if "Error" in line or "Exception" in line:
            capturing = True
        if capturing:
            error_lines.insert(0, line)
        if "Traceback" in line and capturing:
            break

    return "".join(error_lines)

# ============================================================
# 实战5：LLM生成实验方案和分析计划
# 输入：研究目标 → 输出：完整的测序+分析方案
# ============================================================

def generate_experiment_plan(research_goal, constraints=""):
    """
    根据研究目标生成测序方案和分析计划
    参数：
        research_goal: 研究目标描述
        constraints: 限制条件（预算、时间、设备等）
    返回：
        方案文本
    """
    system_prompt = """你是一个有10年经验的生物信息学项目负责人，
擅长设计组学实验方案和分析流程。
你的建议必须基于已发表的文献和行业标准实践。
如果某个建议存在不确定性，请明确标注。"""

    plan_prompt = f"""请为以下研究目标设计完整的测序方案和生信分析计划。

【研究目标】
{research_goal}

【限制条件】
{constraints if constraints else "无特殊限制"}

请按以下结构输出方案：

## 1. 实验设计
- 样本分组和样本量建议（含功效分析依据）
- 对照设计

## 2. 测序方案
- 推荐的测序策略（全基因组/16S/RNA-seq/宏基因组等）
- 测序平台选择及理由
- 测序深度建议
- 文库制备方法

## 3. 生信分析Pipeline
- 完整的分析流程（每步列出工具和参数）
- 质控标准
- 统计分析方法
- 可视化方案

## 4. 预期结果和验证策略
- 预期的主要发现
- 实验验证建议（如qPCR验证、功能实验等）

## 5. 时间线和预算估算
- 各阶段所需时间
- 主要费用项

## 6. 潜在风险和备选方案
- 可能遇到的问题
- 对应的解决策略

请确保每个建议都有文献或行业标准支撑。"""

    plan = ask_llm(plan_prompt, system_prompt)
    return plan


# ---- 使用示例 ----
plan = generate_experiment_plan(
    research_goal="""
    探究2型糖尿病患者肠道菌群与健康人群的差异，
    寻找潜在的微生物生物标志物，
    并探索菌群功能变化与糖代谢异常的关联机制。
    """,
    constraints="""
    - 预算约30万人民币
    - 已有50例T2D患者和50例健康对照的粪便样本
    - 实验室有Illumina NovaSeq可用
    - 计算资源：128GB内存服务器，2张A100 GPU
    - 时间限制：6个月内完成
    """
)

print(plan)

# ---- 方案评审：让LLM自我审查 ----
def review_experiment_plan(plan):
    """
    让LLM以审稿人视角评审实验方案
    """
    review_prompt = f"""请以期刊审稿人的视角审查以下实验方案，指出：
1. 实验设计的不足之处
2. 样本量是否充分
3. 分析方法是否有更好的替代
4. 遗漏的重要步骤
5. 统计方法是否合理
6. 可能的偏倚来源

方案内容：
{plan}

请以表格形式列出问题、严重程度（高/中/低）和建议修改。"""

    review = ask_llm(review_prompt)
    return review

# ============================================================
# AI输出验证清单（每次使用LLM结果前必须过一遍）
# ============================================================

VERIFICATION_CHECKLIST = """
□ 基因/蛋白名称验证：去NCBI Gene或UniProt确认是否存在
□ 文献引用验证：用DOI或PMID去PubMed确认文献真实存在
□ 统计方法验证：确认推荐的检验方法适合你的数据类型和分布
□ 代码运行验证：用小数据集实际运行一遍，检查输出是否合理
□ 参数合理性：比如padj阈值、log2FC阈值是否符合领域惯例
□ 通路分析验证：去KEGG/Reactome确认通路-基因对应关系
□ 物种名称验证：去NCBI Taxonomy确认菌属/种名拼写和分类
"""

# ============================================================
# 综合示例：一键完成"文献调研→方案设计→脚本生成"
# ============================================================

def full_research_workflow(topic, disease):
    """
    完整研究启动工作流
    1. 文献挖掘 → 2. 方案设计 → 3. 脚本生成
    """
    print("=" * 60)
    print(f"研究主题: {topic}")
    print(f"目标疾病: {disease}")
    print("=" * 60)

    # 步骤1：文献调研
    print("\n[步骤1] 文献调研...")
    df_relations = literature_mining_pipeline(
        search_query=f"{topic} {disease} 2023 2024",
        disease=disease
    )

    # 步骤2：基于文献生成方案
    print("\n[步骤2] 生成实验方案...")
    literature_context = ""
    if not df_relations.empty:
        literature_context = f"已知相关发现：\n{df_relations.head(10).to_string()}"

    plan = generate_experiment_plan(
        research_goal=f"基于{topic}研究{disease}的机制和生物标志物。\n{literature_context}",
        constraints="预算20万，样本量100例，6个月"
    )

    # 步骤3：生成核心分析脚本
    print("\n[步骤3] 生成分析脚本...")
    script = generate_bioinfo_script(
        f"根据以下方案生成核心数据分析脚本：\n{plan[:2000]}"
    )

    # 保存所有结果
    with open("research_plan.md", "w", encoding="utf-8") as f:
        f.write(f"# {topic} 研究方案\n\n{plan}")
    with open("analysis_pipeline.py", "w", encoding="utf-8") as f:
        f.write(script)

    print("\n[完成] 所有文件已生成，请人工审查后使用！")
    return plan, script


# 使用示例（取消注释运行）：
# plan, script = full_research_workflow(
#     topic="gut microbiome short-chain fatty acids",
#     disease="type 2 diabetes"
# )