"""Organize PDFs into category folders and emit a translation work plan.

Usage: python organize.py <papers.json> <base_dir>
"""
import json
import os
import shutil
import sys

CAT_DIRS = {
    '基础模型': '01-基础模型',
    '代码模型': '02-代码模型',
    '数学模型': '03-数学模型',
    '推理模型': '04-推理模型',
    '多模态': '05-多模态',
    '形式化证明': '06-形式化证明',
    '架构': '07-架构',
    '训练方法': '08-训练方法',
    '基础设施': '09-基础设施',
    '其他': '10-其他',
}

SHORT = {
    '2401.02954': 'DeepSeek-LLM',
    '2401.06066': 'DeepSeekMoE',
    '2401.14196': 'DeepSeek-Coder',
    '2402.03300': 'DeepSeekMath',
    '2403.05525': 'DeepSeek-VL',
    '2405.04434': 'DeepSeek-V2',
    '2405.14333': 'DeepSeek-Prover',
    '2406.11931': 'DeepSeek-Coder-V2',
    '2407.01906': 'ESFT-专家专精微调',
    '2408.08152': 'DeepSeek-Prover-V1.5',
    '2408.14158': 'Fire-Flyer-AI-HPC',
    '2410.13848': 'Janus',
    '2412.10302': 'DeepSeek-VL2',
    '2412.19437': 'DeepSeek-V3',
    '2501.12948': 'DeepSeek-R1',
    '2501.17811': 'Janus-Pro',
    '2504.02495': 'DeepSeek-GRM',
    '2504.21801': 'DeepSeek-Prover-V2',
    '2505.09343': 'Insights-into-DeepSeek-V3',
    '2510.18234': 'DeepSeek-OCR',
    '2511.22570': 'DeepSeekMath-V2',
    '2512.02556': 'DeepSeek-V3.2',
    '2512.24880': 'mHC',
    '2601.07372': 'Engram-条件记忆',
    '2601.20552': 'DeepSeek-OCR-2',
    '2602.21548': 'DualPath',
    '2606.19348': 'DeepSeek-V4',
    '2609.19969': 'DeepSeek-V4.1-Flash',
}


def main():
    papers = json.load(open(sys.argv[1], encoding='utf-8'))
    base = sys.argv[2]
    pdfdir = os.path.join(base, 'pdf')
    plan = []
    for cat in CAT_DIRS:
        pass
    counters = {}
    for p in papers:
        cat = p.get('category', '其他')
        d = CAT_DIRS.get(cat, '10-其他')
        counters[d] = counters.get(d, 0) + 1
        nn = '%02d' % counters[d]
        short = SHORT.get(p['id'], p['id'])
        stem = '%s-%s-%s' % (nn, short, p['id'])
        outdir = os.path.join(base, d)
        os.makedirs(outdir, exist_ok=True)
        src = os.path.join(pdfdir, p['id'] + '.pdf')
        dst = os.path.join(outdir, stem + '.pdf')
        if os.path.exists(src) and not os.path.exists(dst):
            shutil.move(src, dst)
        p['dir'] = d
        p['stem'] = stem
        p['pdf_path'] = os.path.join(d, stem + '.pdf')
        p['md_path'] = os.path.join(d, stem + '.md')
        plan.append(p)
    with open(os.path.join(base, '_scripts', 'plan.json'), 'w', encoding='utf-8') as f:
        json.dump(plan, f, ensure_ascii=False, indent=1)
    for p in plan:
        print(p['dir'], '|', p['stem'], '|', p['title'][:60])
    print('total', len(plan))


if __name__ == '__main__':
    main()