#!/usr/bin/env python3
"""Fetch recent cs.LG arXiv submissions (sorted by submission date) and write
a flat CSV of the Atom feed entries.

Usage:
    python3 make_arxiv_csv.py [OUT_CSV]

Default output: /tmp/arxiv-lg-recent.csv
"""
import csv
import re
import subprocess
import sys
from xml.etree import ElementTree as ET

NS = {
    "atom": "http://www.w3.org/2005/Atom",
    "arxiv": "http://arxiv.org/schemas/atom",
    "os": "http://a9.com/-/spec/opensearch/1.1/",
}

URL = (
    "https://export.arxiv.org/api/query?search_query=cat%3Acs.LG"
    "&start=0&max_results=20&sortBy=submittedDate&sortOrder=descending"
)


def fetch_xml():
    proc = subprocess.run(
        ["curl", "-s", "--max-time", "30", URL],
        check=True,
        capture_output=True,
        text=True,
    )
    return proc.stdout


def _text(element):
    if element is None:
        return ""
    return "".join(element.itertext()).strip()


def parse(xml_text):
    root = ET.fromstring(xml_text)
    rows = []
    for entry in root.findall("atom:entry", NS):
        def link_href(rel):
            for link in entry.findall("atom:link", NS):
                if link.get("rel") == rel:
                    return link.get("href")
            return ""

        arxiv_id = _text(entry.find("atom:id", NS))
        arxiv_id = re.sub(r"^http://arxiv.org/abs/", "", arxiv_id)
        arxiv_id = re.sub(r"v\d+$", "", arxiv_id)

        title = _text(entry.find("atom:title", NS))
        summary = _text(entry.find("atom:summary", NS))
        summary = re.sub(r"[\r\n]+", " ", summary)

        published = _text(entry.find("atom:published", NS))[:10]

        authors = "; ".join(
            _text(a.find("atom:name", NS))
            for a in entry.findall("atom:author", NS)
            if _text(a.find("atom:name", NS))
        )

        rows.append(
            {
                "arxiv_id": arxiv_id,
                "title": title,
                "authors": authors,
                "published": published,
                "abstract": summary,
            }
        )
    return rows


def main():
    out_path = sys.argv[1] if len(sys.argv) > 1 else "/tmp/arxiv-lg-recent.csv"
    xml_text = fetch_xml()
    rows = parse(xml_text)
    fieldnames = ["arxiv_id", "title", "authors", "published", "abstract"]
    with open(out_path, "w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            writer.writerow(row)
    print("rows=%d url=%s" % (len(rows), URL))
    print("wrote=%s" % out_path)


if __name__ == "__main__":
    main()