import json
import csv
from pathlib import Path

dataset='sudoku_large'
split='test'

# Load test data
with open(f'{dataset}/{split}.jsonl', 'r') as f:
    data = [json.loads(line) for line in f]

# Write to CSV
with open(f'{dataset}/{split}.csv', 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    writer.writerow(['input_image', 'video', 'prompt'])
    
    for idx, item in enumerate(data):
        writer.writerow([
            'images/' + item['image'],
            'videos/' + item['video'],
            item['prompt'],
        ])

# Rename `{split}.jsonl' to `{split}_info.jsonl`
Path(f'{dataset}/{split}.jsonl').rename(Path(f'{dataset}/{split}_info.jsonl'))