#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""PubMed 反推：中国机构 mRNA/LNP/circRNA 论文 → 通讯作者 + 单位 + 国自然批准号
输出：数据/mrna_pubmed_raw.json
"""
import io
import json
import re
import time
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

QUERY = ('("mRNA vaccine"[tiab] OR "mRNA therapeutic"[tiab] OR "mRNA delivery"[tiab] '
         'OR "lipid nanoparticle"[tiab] OR LNP[tiab] OR circRNA[tiab] '
         'OR "self-amplifying RNA"[tiab] OR "in vitro transcription"[tiab]) AND China[ad] AND 2024:2025[pdat]')
E = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/'
OUT = '/Users/liuxinyuan/Desktop/Hermes输出-工作类/数据/mrna_pubmed_raw.json'

PROV = {
    '广东': ['广东', '广州', '深圳', '珠海', '东莞', '佛山', '中山', '汕头', '湛江', '惠州', '江门', '肇庆', '澳门'],
    '浙江': ['浙江', '杭州', '宁波', '温州', '绍兴', '嘉兴', '金华', '台州', '湖州', '丽水', '衢州', '舟山'],
}


def get(url, tries=4):
    for i in range(tries):
        try:
            with urllib.request.urlopen(url, timeout=60) as r:
                return r.read()
        except Exception as e:
            if i == tries - 1:
                raise
            time.sleep(2 + i)


# 1) esearch
u = E + 'esearch.fcgi?db=pubmed&retmode=json&retmax=5000&term=' + urllib.parse.quote(QUERY)
ids = json.loads(get(u))['esearchresult']['idlist']
print('命中 PMID', len(ids), flush=True)

# 2) efetch 分批
recs = []
B = 200
for i in range(0, len(ids), B):
    chunk = ids[i:i + B]
    u = E + 'efetch.fcgi?db=pubmed&retmode=xml&id=' + ','.join(chunk)
    xml = get(u)
    root = ET.fromstring(xml)
    for art in root.findall('.//PubmedArticle'):
        try:
            pmid = art.findtext('.//PMID') or ''
            title = ''.join(art.find('.//ArticleTitle').itertext()).strip() if art.find('.//ArticleTitle') is not None else ''
            jour = art.findtext('.//Journal/Title') or ''
            year = art.findtext('.//JournalIssue/PubDate/Year') or art.findtext('.//ArticleDate/Year') or ''
            authors = art.findall('.//AuthorList/Author')
            names, affs = [], []
            for a in authors:
                ln = a.findtext('LastName') or a.findtext('CollectiveName') or ''
                fn = a.findtext('ForeName') or ''
                aff = a.findtext('AffiliationInfo/Affiliation') or ''
                names.append((ln, fn, aff))
                if aff:
                    affs.append(aff)
            if not names:
                continue
            last = names[-1]
            grants = []
            for g in art.findall('.//GrantList/Grant'):
                gid = (g.findtext('GrantID') or '').strip()
                ag = (g.findtext('Agency') or '').strip()
                if gid or ag:
                    grants.append({'id': gid, 'agency': ag})
            text = ' | '.join(affs)
            prov = ''
            for p, kws in PROV.items():
                if any(k in text for k in kws):
                    prov = p
                    break
            recs.append({
                'pmid': pmid, 'title': title[:300], 'journal': jour[:60], 'year': year,
                'last_ln': last[0], 'last_fn': last[1], 'last_aff': last[2][:400],
                'first_ln': names[0][0], 'first_fn': names[0][1],
                'prov': prov, 'china': 'China' in text or '中国' in text,
                'grants': grants, 'n_auth': len(names),
            })
        except Exception as e:
            continue
    print('批次 %d/%d 累计 %d' % (i // B + 1, (len(ids) + B - 1) // B, len(recs)), flush=True)
    time.sleep(0.4)

io.open(OUT, 'w', encoding='utf-8').write(json.dumps(recs, ensure_ascii=False))
print('✅ 保存', len(recs), '条 →', OUT, flush=True)
cn = [r for r in recs if r['china']]
print('中国机构', len(cn), '｜广东', sum(1 for r in cn if r['prov'] == '广东'), '｜浙江', sum(1 for r in cn if r['prov'] == '浙江'))