#!/usr/bin/env python3
"""
收集新能源汽车领域最新论文的脚本
关注方向：AI、故障诊断、电池技术
"""

import subprocess
import json
import os
from datetime import datetime

def fetch_papers():
    """使用 arXiv API 获取论文列表"""
    
    # 定义查询关键词
    queries = [
        "electric vehicle AI machine learning",
        "battery management system fault diagnosis electric vehicle", 
        "lithium battery SOH estimation deep learning",
        "neural network battery state of charge prediction",
        "predictive maintenance electric vehicle powertrain"
    ]
    
    papers_data = []
    
    for query in queries:
        cmd = [
            'curl', '-s', 
            f'https://export.arxiv.org/api/query?search_query=all:{query}&max_results=5&start=0'
        ]
        
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
            if result.returncode == 0 and result.stdout:
                # 简单解析 XML
                import xml.etree.ElementTree as ET
                root = ET.fromstring(result.stdout)
                
                # 命名空间
                ns = {
                    'atom': 'http://www.w3.org/2005/Atom',
                    'arxiv': 'http://arxiv.org/schemas/atom'
                }
                
                entries = root.findall('atom:entry', ns)
                for entry in entries:
                    title_elem = entry.find('atom:title', ns)
                    published_elem = entry.find('atom:published', ns)
                    id_elem = entry.find('atom:id', ns)
                    summary_elem = entry.find('atom:summary', ns)
                    
                    if title_elem is not None and published_elem is not None:
                        paper = {
                            'title': title_elem.text.strip() if title_elem.text else '',
                            'published': published_elem.text,
                            'url': id_elem.text if id_elem is not None else '',
                            'pdf_url': id_elem.text.replace('abs', 'pdf') if id_elem is not None else '',
                            'abstract': summary_elem.text[:500].strip() if summary_elem is not None and summary_elem.text else ''
                        }
                        
                        # 检查是否与新能源汽车相关
                        keywords = ['electric', 'vehicle', 'battery', 'ev', 'automotive', 'powertrain']
                        text_lower = (paper['title'].lower() + paper.get('abstract', '')).lower()
                        
                        if any(kw in text_lower for kw in keywords):
                            paper['keywords'] = [kw for kw in keywords if kw in text_lower]
                            papers_data.append(paper)
                            
        except Exception as e:
            print(f"Error processing query '{query}': {e}")
            
    return papers_data

def save_to_file(papers, filename='new_energy_papers.json'):
    """保存论文数据到文件"""
    with open(filename, 'w', encoding='utf-8') as f:
        json.dump({
            'collected_date': datetime.now().isoformat(),
            'papers': papers
        }, f, ensure_ascii=False, indent=2)
    print(f"\n已保存 {len(papers)} 篇论文到 {filename}")

if __name__ == '__main__':
    print("正在从 arXiv 获取新能源汽车相关论文...")
    papers = fetch_papers()
    save_to_file(papers)
    
    print("\n=== 检索到的论文 ===\n")
    for i, p in enumerate(papers, 1):
        print(f"{i}. [{p['published'][:10]}] {p['title']}")
        if 'abstract' in p:
            print(f"   摘要：{p['abstract'][:100]}...")
        print(f"   PDF: {p['pdf_url']}")
        print()