#!/usr/bin/env python3
"""给 arxiv_cs_2022_2026.jsonl 每行追加 CS 子类标注，保持原有字段与样式不变。

新增字段（插在 primary_category 之后）：
  cs_categories      : categories 中属于 cs.* 的代码子集
  cs_category_names  : 上述代码对应的全称
  cs_primary         : 首个 cs.* 代码（该论文的主 CS 子类）
"""
import json

SRC = "arxiv_cs_2022_2026.jsonl"
DST = "arxiv_cs_2022_2026.annotated.jsonl"

CS_NAMES = {
    "cs.AI": "Artificial Intelligence",
    "cs.AR": "Hardware Architecture",
    "cs.CC": "Computational Complexity",
    "cs.CE": "Computational Engineering, Finance, and Science",
    "cs.CG": "Computational Geometry",
    "cs.CL": "Computation and Language",
    "cs.CR": "Cryptography and Security",
    "cs.CV": "Computer Vision and Pattern Recognition",
    "cs.CY": "Computers and Society",
    "cs.DB": "Databases",
    "cs.DC": "Distributed, Parallel, and Cluster Computing",
    "cs.DL": "Digital Libraries",
    "cs.DM": "Discrete Mathematics",
    "cs.DS": "Data Structures and Algorithms",
    "cs.ET": "Emerging Technologies",
    "cs.FL": "Formal Languages and Automata Theory",
    "cs.GL": "General Literature",
    "cs.GR": "Graphics",
    "cs.GT": "Computer Science and Game Theory",
    "cs.HC": "Human-Computer Interaction",
    "cs.IR": "Information Retrieval",
    "cs.IT": "Information Theory",
    "cs.LG": "Machine Learning",
    "cs.LO": "Logic in Computer Science",
    "cs.MA": "Multiagent Systems",
    "cs.MM": "Multimedia",
    "cs.MS": "Mathematical Software",
    "cs.NA": "Numerical Analysis",
    "cs.NE": "Neural and Evolutionary Computing",
    "cs.NI": "Networking and Internet Architecture",
    "cs.OH": "Other Computer Science",
    "cs.OS": "Operating Systems",
    "cs.PF": "Performance",
    "cs.PL": "Programming Languages",
    "cs.RO": "Robotics",
    "cs.SC": "Symbolic Computation",
    "cs.SD": "Sound",
    "cs.SE": "Software Engineering",
    "cs.SI": "Social and Information Networks",
    "cs.SY": "Systems and Control",
}

# 保持与原文件一致的字段顺序，把标注插在 primary_category 之后
ORDER = ["id", "submitted_date", "year", "title", "authors", "categories",
         "primary_category", "cs_categories", "cs_category_names",
         "cs_primary", "abstract", "abs_url", "pdf_url"]


def main():
    n = 0
    with open(SRC) as fin, open(DST, "w") as fout:
        for line in fin:
            line = line.strip()
            if not line:
                continue
            d = json.loads(line)
            cs = [c for c in d.get("categories", []) if c.startswith("cs.")]
            d["cs_categories"] = cs
            d["cs_category_names"] = [CS_NAMES.get(c, c) for c in cs]
            d["cs_primary"] = cs[0] if cs else None
            out = {k: d[k] for k in ORDER if k in d}
            # 原顺序之外若有其它字段，补在后面
            for k in d:
                if k not in out:
                    out[k] = d[k]
            fout.write(json.dumps(out, ensure_ascii=False,
                                  separators=(",", ":")) + "\n")
            n += 1
    print(f"DONE rows={n} -> {DST}", flush=True)


if __name__ == "__main__":
    main()