"""Partition the PatchScholar export by reviewed date flags; requires Pybtex.

Run after the complete export has been generated. Original snapshot exports are
untouched. Every serialized key set is parsed and checked against the register.
"""
import csv
import json
from pathlib import Path
from pybtex.database import BibliographyData, parse_file

ROOT = Path(__file__).resolve().parents[1]


def export():
    with (ROOT / 'reviews/2026-09-11-expanded.csv').open() as stream:
        rows = list(csv.DictReader(stream))
    whole = ROOT / 'exports/andreev-all-2026-09-11.bib'
    bib = parse_file(str(whole))
    accepted = {f"patchscholar-{r['canonical_work_id']}" for r in rows if r['decision'] == 'include'}
    assert set(bib.entries) == accepted
    counts = {whole.name: len(bib.entries)}
    for name, field in [('andreev-expanded-2000-2026.bib', 'expanded_eligible'),
                        ('andreev-backfill-2000-2016.bib', 'backfill_eligible')]:
        keys = {f"patchscholar-{r['canonical_work_id']}" for r in rows
                if r['decision'] == 'include' and r[field] == 'true'}
        output = BibliographyData(entries={k: v for k, v in bib.entries.items() if k in keys})
        path = ROOT / 'exports' / name
        output.to_file(str(path), bib_format='bibtex')
        assert set(parse_file(str(path)).entries) == keys
        counts[name] = len(keys)
    (ROOT / 'reviews/backfill-2000/bibliography-validation.json').write_text(json.dumps(counts, indent=2)+'\n')
    print(counts)


if __name__ == '__main__':
    export()