import argparse

import pandas as pd
from datasets import load_dataset

parser = argparse.ArgumentParser()
parser.add_argument('output_filepath', type=str, help='Path to save the output file')
parser.add_argument(
    '--dataset_name',
    type=str,
    help='Name of the dataset to download',
    default='princeton-nlp/SWE-bench_Lite',
)
parser.add_argument('--split', type=str, help='Split to download', default='test')
args = parser.parse_args()

dataset = load_dataset(args.dataset_name, split=args.split)
output_filepath = args.output_filepath
print(
    f'Downloading gold patches from {args.dataset_name} (split: {args.split}) to {output_filepath}'
)
patches = [
    {'instance_id': row['instance_id'], 'model_patch': row['patch']} for row in dataset
]
print(f'{len(patches)} gold patches loaded')
pd.DataFrame(patches).to_json(output_filepath, lines=True, orient='records')
print(f'Patches saved to {output_filepath}')