import argparse
import json
import random
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--questions", default="data/train/verified-mcq.jsonl")
    parser.add_argument("--exclude", default="data/train/star2/overlap/excluded_ids.txt",
                        help="ids that review_overlap.py found similar to a dev or test question")
    parser.add_argument("--out_dir", default="data/train/star2")
    parser.add_argument("--shards", type=int, default=2)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    lines = Path(args.questions).read_text(encoding="utf-8").splitlines()
    excluded = set(Path(args.exclude).read_text(encoding="utf-8").split())
    kept = [line for line in lines if json.loads(line)["id"] not in excluded]
    assert len(lines) - len(kept) == len(excluded), "an excluded id is not in the question file"

    random.Random(args.seed).shuffle(kept)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    for n in range(args.shards):
        shard = kept[n :: args.shards]
        (out_dir / f"shard-{n}.jsonl").write_text("".join(line + "\n" for line in shard), encoding="utf-8")
        print(f"shard-{n}.jsonl: {len(shard)} questions")
    print(f"Kept {len(kept)} of {len(lines)} questions; excluded {len(excluded)} listed in {args.exclude}")


if __name__ == "__main__":
    main()