import marimo


__generated_with = "0.17.7"
app = marimo.App()


@app.cell
def _():
    """
    标题：使用 Pydantic AI 进行 PubMed 文献分析

    描述：
    本 notebook 演示如何使用 Pydantic AI 分析 PubMed 文献。
    相比旧版 LangChain 实现，新版提供强类型输出和依赖注入。

    学习目标：
    1. 如何配置 AnalysisConfig 启用 Pydantic AI
    2. 如何使用 PaperAnalyzer 进行批量分析
    3. 如何处理分析结果
    """
    import marimo as mo

    return (mo,)


@app.cell
def _(mo):
    mo.md("""
    ## 📚 PubMed 搜索与保存示例（Pydantic AI 版本）

    本 notebook 演示如何使用 Pydantic AI 分析 PubMed 文献。

    ### 与传统方法的区别

    | 特性 | LangChain | Pydantic AI |
    |------|-----------|-------------|
    | 输出格式 | 手动解析 JSON | 自动 Pydantic 验证 |
    | 类型安全 | 运行时检查 | 静态类型 + 运行时验证 |
    | Token 追踪 | 需手动实现 | 内置支持 |
    | 错误处理 | 解析异常 | 验证错误 |

    ### 运行步骤
    1. 执行所有 cell 导入依赖
    2. 自动执行 PubMed 搜索
    3. 获取论文详情
    4. **Pydantic AI 分析**：自动分析摘要内容
    5. 查看 DataFrame 结果
    """)
    return


@app.cell
def _():
    from pathlib import Path

    from information_composer.pubmed.database import PubMedDatabase

    db_path = Path(".pubmed_marimo.db")
    db_path.parent.mkdir(parents=True, exist_ok=True)

    db = PubMedDatabase(str(db_path))
    return


@app.cell
def _():
    from information_composer.pubmed.pubmed import query_pmid

    return (query_pmid,)


@app.cell
def _(query_pmid):
    query = "(uORFs OR uORF) and (plant)"
    pmids = query_pmid(email="hello@example.com", query=query)
    return (pmids,)


@app.cell
def _(pmids):
    print(pmids[:5])
    return


@app.cell
def _():
    from information_composer.pubmed.pubmed import fetch_pubmed_details_batch

    return (fetch_pubmed_details_batch,)


@app.cell
async def _(fetch_pubmed_details_batch, pmids):
    records = await fetch_pubmed_details_batch(pmids)
    return (records,)


@app.cell
async def _(records):
    """
    使用 Pydantic AI 分析论文摘要

    相比直接使用传统 LLM API，Pydantic AI 提供：
    - 结构化输出：返回 Pydantic 模型而非原始文本
    - 自动验证：输出格式自动校验
    - Token 追踪：自动记录使用量
    """
    from information_composer.pubmed.analyzer import (
        AnalysisConfig,
        PaperAnalyzer,
        PaperInput,
    )

    # 配置 Pydantic AI 分析器
    config = AnalysisConfig(
        analysis_types=["summary"],  # 仅进行摘要分析
        llm_provider="openai",
        llm_model="anthropic/claude-sonnet-4.5",
        use_pydantic_ai=True,  # 启用 Pydantic AI
        cache_enabled=True,  # 启用缓存避免重复调用
    )

    analyzer = PaperAnalyzer(config)

    # 异步分析单篇论文并提取中文描述
    async def analyze_with_summary(record):
        if record.get("abstract"):
            paper = PaperInput(
                pmid=record.get("pmid", ""),
                title=record.get("title", ""),
                abstract=record["abstract"],
            )

            try:
                result = await analyzer.analyze_paper(paper)

                # 从 summary 中提取关键信息作为中文描述
                # 注意：实际应用中可以训练专用模型做翻译
                findings = result.summary.main_findings if result.summary else []
                record["abstract_zh"] = (
                    "; ".join(findings[:2]) if findings else "分析完成"
                )

                # 添加 token 使用信息（Pydantic AI 特性）
                if result.processing_metadata.tokens_used:
                    record["tokens_used"] = result.processing_metadata.tokens_used

            except Exception as e:
                record["abstract_zh"] = f"分析失败: {str(e)}"
        else:
            record["abstract_zh"] = "无摘要"

        return record

    # 批量处理
    import asyncio

    tasks = [analyze_with_summary(r) for r in records]
    analyzed_records = await asyncio.gather(*tasks)

    # 更新原始 records
    for i, analyzed in enumerate(analyzed_records):
        records[i] = analyzed

    return (records,)


@app.cell
def _():
    import pandas as pd

    return (pd,)


@app.cell
def _(pd, records):
    df = pd.DataFrame(records)
    return (df,)


@app.cell
def _(df):
    df
    return


@app.cell
def _(records, mo):
    # 统计 token 使用量（Pydantic AI 特性）
    total_tokens = sum(r.get("tokens_used", 0) for r in records)
    mo.md(f"**Token 使用统计**: 共使用 {total_tokens} tokens")
    return


@app.cell
def _():
    return


if __name__ == "__main__":
    app.run()
