"""Collect DeepSeek papers from the arXiv API.

Usage:  python collect.py <output.json>
"""
import urllib.request
import urllib.parse
import json
import sys
import time
import re
import os
import xml.etree.ElementTree as ET

NS = {'a': 'http://www.w3.org/2005/Atom',
      'arxiv': 'http://arxiv.org/schemas/atom'}

UA = 'Mozilla/5.0 (compatible; paperbot/1.0)'


def fetch(q, start=0, max_results=100, sort='submittedDate'):
    params = {
        'search_query': q,
        'start': start,
        'max_results': max_results,
        'sortBy': sort,
        'sortOrder': 'descending',
    }
    url = 'https://export.arxiv.org/api/query?' + urllib.parse.urlencode(params)
    req = urllib.request.Request(url, headers={'User-Agent': UA})
    last = None
    for attempt in range(5):
        try:
            with urllib.request.urlopen(req, timeout=90) as r:
                return r.read()
        except Exception as e:  # noqa: BLE001
            last = e
            time.sleep(4)
    print('FETCH FAILED', q, last, file=sys.stderr)
    return None


def parse(xml):
    root = ET.fromstring(xml)
    out = []
    for e in root.findall('a:entry', NS):
        aid = e.find('a:id', NS).text
        m = re.search(r'arxiv\.org/abs/(.+)$', aid)
        comment = e.find('arxiv:comment', NS)
        primary = e.find('arxiv:primary_category', NS)
        journal = e.find('arxiv:journal_ref', NS)
        doi = e.find('arxiv:doi', NS)
        out.append({
            'id': (m.group(1) if m else aid),
            'version_id': aid,
            'title': ' '.join(e.find('a:title', NS).text.split()),
            'published': e.find('a:published', NS).text,
            'updated': e.find('a:updated', NS).text,
            'authors': [a.find('a:name', NS).text for a in e.findall('a:author', NS)],
            'summary': ' '.join(e.find('a:summary', NS).text.split()),
            'comment': (' '.join(comment.text.split()) if comment is not None and comment.text else ''),
            'journal_ref': (' '.join(journal.text.split()) if journal is not None and journal.text else ''),
            'doi': (doi.text.strip() if doi is not None and doi.text else ''),
            'primary': (primary.get('term') if primary is not None else ''),
        })
    return out


def search(q, cap=200):
    allr = []
    for start in range(0, cap, 100):
        xml = fetch(q, start=start)
        if not xml:
            break
        res = parse(xml)
        allr.extend(res)
        if len(res) < 100:
            break
        time.sleep(3)
    return allr


if __name__ == '__main__':
    queries = [
        'all:"DeepSeek"',
        'abs:"DeepSeek"',
        'ti:"DeepSeek"',
        'all:"DeepSeek-V3" OR all:"DeepSeek-R1"',
        'all:"DeepSeekMath" OR all:"DeepSeek-Coder" OR all:"DeepSeek-VL"',
        'all:"DeepSeekMoE" OR all:"DeepSeek-Prover" OR all:"Janus"',
        'all:"Fire-Flyer" OR all:"DeepSeek-OCR"',
    ]
    seen = {}
    for q in queries:
        print('query:', q, file=sys.stderr)
        for r in search(q):
            key = r['id'].split('v')[0]
            if key not in seen:
                seen[key] = r
        time.sleep(3)
    res = sorted(seen.values(), key=lambda x: x['published'])
    out = sys.argv[1] if len(sys.argv) > 1 else 'raw_deepseek.json'
    os.makedirs(os.path.dirname(out), exist_ok=True)
    with open(out, 'w', encoding='utf-8') as f:
        json.dump(res, f, ensure_ascii=False, indent=1)
    print('total', len(res), file=sys.stderr)