"""Generate the README index for the DeepSeek paper collection.

Usage: python make_index.py <plan.json> <base_dir>
"""
import json
import os
import sys

CAT_ORDER = [
    '01-基础模型', '02-代码模型', '03-数学模型', '04-推理模型', '05-多模态',
    '06-形式化证明', '07-架构', '08-训练方法', '09-基础设施', '10-其他',
]
CAT_DESC = {
    '01-基础模型': '通用大语言模型系列技术报告（LLM / MoE / V2 / V3 / V3.2 / V4 / V4.1-Flash）',
    '02-代码模型': '代码智能模型与代码驱动推理（DeepSeek-Coder / Coder-V2 / CodeI/O）',
    '03-数学模型': '数学推理模型（DeepSeekMath / DeepSeekMath-V2）',
    '04-推理模型': '强化学习驱动的推理模型（DeepSeek-R1）',
    '05-多模态': '视觉-语言与统一多模态模型（VL / VL2 / Janus / Janus-Pro / JanusFlow / OCR 系列 / DreamCraft3D）',
    '06-形式化证明': '定理证明与形式化数学（DeepSeek-Prover 系列）',
    '07-架构': '基础架构创新（mHC 超连接 / Engram 条件记忆 / NSA 原生稀疏注意力 / 无辅助损失负载均衡）',
    '08-训练方法': '训练与对齐方法（ESFT 专家专精微调 / GRM 通用奖励模型）',
    '09-基础设施': '训练与推理基础设施（Fire-Flyer AI-HPC / 硬件洞察 / DualPath 存储带宽）',
}


def clean_date(p):
    """Return a short human date like '11 Feb 2025'."""
    import re as _re
    d = p.get('date', '') or ''
    if d.startswith('[') or len(d) > 24:
        m = _re.search(r'on\s+(\d{1,2}\s+\w+\s+\d{4})', d)
        if m:
            return m.group(1)
    return d


def main():
    plan = json.load(open(sys.argv[1], encoding='utf-8'))
    base = sys.argv[2]
    bycat = {}
    for p in plan:
        bycat.setdefault(p['dir'], []).append(p)
    for c in bycat:
        bycat[c].sort(key=lambda x: x['stem'])
    out = []
    out.append('# DeepSeek 官方论文合集（原文 + 中文翻译）')
    out.append('')
    out.append('本目录收录 **深度求索（DeepSeek / DeepSeek-AI）自成立以来以论文形式公开发布的'
               '全部研究成果**，共 **%d 篇**（截至 2026-09-20）。' % len(plan))
    out.append('')
    out.append('每篇论文均包含两部分：')
    out.append('')
    out.append('- **原文 PDF**：从 arXiv 下载的官方预印本（未做改动）')
    out.append('- **中文 Markdown**：摘要全文翻译 + 核心贡献 + 方法要点 + 实验结果 + 结论，'
               '便于快速掌握论文内容')
    out.append('')
    out.append('> 说明：本次中文版本为「摘要 + 核心贡献」精读版。如需某篇论文的全文逐段翻译，'
               '可在此目录基础上追加。')
    out.append('')
    out.append('## 收录范围与判定依据')
    out.append('')
    out.append('论文清单通过以下渠道交叉核验后确定，凡判定为 DeepSeek 官方成果者全部收录：')
    out.append('')
    out.append('1. **DeepSeek 官网「研究索引」**（https://www.deepseek.com/news ）——官方自行列出的论文条目')
    out.append('2. **arXiv 作者检索**：作者字段含 `DeepSeek-AI` 的全部论文')
    out.append('3. **论文署名单位核验**：逐篇抓取 arXiv 摘要页，确认作者单位为 `DeepSeek-AI` / `DeepSeek AI`')
    out.append('4. **社区索引交叉比对**：Hugging Face `DeepSeek Papers` 合集、第三方综述')
    out.append('')
    out.append('判定标准：论文作者单位标注为 DeepSeek-AI / DeepSeek AI，'
               '或作者主体为 DeepSeek 在职研究员。'
               '仅「以 DeepSeek 模型为研究对象」的第三方论文（数量众多）**不予收录**。')
    out.append('')
    out.append('> **另有一项非论文形式的研究发布**：[LPLB（Linear-Programming-Based Load Balancer）]'
               '(https://github.com/deepseek-ai/LPLB)——DeepSeek 于 2025 年 11 月 1 日发布的 '
               'MoE 专家并行负载均衡器，以 GitHub 仓库形式发布（早期研究阶段），'
               '未上传 arXiv 预印本，故不在本目录论文之列。')
    out.append('')
    out.append('## 目录结构')
    out.append('')
    out.append('```')
    out.append('大模型/deepseek/')
    out.append('├── README.md                   本索引')
    out.append('├── 01-基础模型/ … 09-基础设施/    各篇论文的原文 PDF + 中文 Markdown（同名配对）')
    out.append('├── _text/                      论文正文纯文本（供后续「全文逐段翻译」使用）')
    out.append('├── _digest/                    摘要 + 章节标题 + 正文首尾（精读写作素材）')
    out.append('└── _scripts/                   清单核验、PDF 下载、正文提取、索引生成脚本')
    out.append('```')
    out.append('')
    out.append('`_text/` 与 `_digest/` 已按 arXiv 编号逐篇存好论文全文，'
               '后续若要把某几篇扩写为**全文逐段翻译**，可直接基于这些文件进行，无需重新抓取。')
    out.append('')
    out.append('> **DreamCraft3D**（2310.16818）为 DeepSeek 参与度有限的早期合作论文：'
               '作者单位列入 DeepSeek AI，论文注明「实习期间部分完成」，'
               '为完整起见一并收录并单独标注。')
    out.append('')
    out.append('---')
    out.append('')
    out.append('## 目录')
    out.append('')
    for c in CAT_ORDER:
        if c in bycat:
            out.append('- [%s](#%s)（%d 篇）' % (c[3:], c.lower().replace('.', ''), len(bycat[c])))
    out.append('')
    out.append('---')
    out.append('')
    out.append('## 论文总览')
    out.append('')
    out.append('| # | 论文 | arXiv | 发布日期 | 类别 |')
    out.append('|---|---|---|---|---|')
    n = 0
    for c in CAT_ORDER:
        for p in bycat.get(c, []):
            n += 1
            out.append('| %d | [%s](%s) | [%s](%s) | %s | %s |' % (
                n, p['title'], p['md_path'].replace('\\', '/'), p['id'], p['abs_url'],
                clean_date(p), c[3:]))
    out.append('')
    out.append('---')
    out.append('')
    for c in CAT_ORDER:
        if c not in bycat:
            continue
        out.append('## %s' % c[3:])
        out.append('')
        out.append('> %s' % CAT_DESC.get(c, ''))
        out.append('')
        for p in bycat[c]:
            stem = p['stem']
            done = os.path.exists(os.path.join(base, p['md_path']))
            out.append('### %s' % p['title'])
            out.append('')
            out.append('- **arXiv**：%s ｜ **发布**：%s' % (p['id'], clean_date(p)))
            a = p.get('authors', '')
            out.append('- **作者**：%s' % (a[:300] + ('…' if len(a) > 300 else '')))
            out.append('- **原文 PDF**：[%s.pdf](%s)' % (stem, p['pdf_path'].replace('\\', '/')))
            if done:
                out.append('- **中文翻译**：[%s](%s)' % (
                    os.path.basename(p['md_path'].replace('\\', '/')),
                    p['md_path'].replace('\\', '/')))
            else:
                out.append('- **中文翻译**：（待生成）')
            out.append('')
        out.append('')
    with open(os.path.join(base, 'README.md'), 'w', encoding='utf-8') as f:
        f.write('\n'.join(out))
    print('README written, papers:', len(plan))


if __name__ == '__main__':
    main()