#!/usr/bin/env python3
"""
Generate a CSV file with ARC V1 task IDs and their set (training/evaluation)
"""

import csv

def main():
    # Read the training IDs
    with open('arc_v1_training_ids.txt', 'r') as f:
        training_ids = [line.strip() for line in f if line.strip()]

    # Read the evaluation IDs
    with open('arc_v1_evaluation_ids.txt', 'r') as f:
        evaluation_ids = [line.strip() for line in f if line.strip()]

    # Create CSV file
    with open('arc_v1_task_ids.csv', 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)

        # Write header
        writer.writerow(['id', 'set'])

        # Write training tasks
        for task_id in training_ids:
            writer.writerow([task_id, 'training'])

        # Write evaluation tasks
        for task_id in evaluation_ids:
            writer.writerow([task_id, 'evaluation'])

    print("=" * 80)
    print("ARC V1 Task IDs - CSV Export")
    print("=" * 80)
    print(f"\nTraining Tasks: {len(training_ids)}")
    print(f"Evaluation Tasks: {len(evaluation_ids)}")
    print(f"Total: {len(training_ids) + len(evaluation_ids)}")

    print("\nFirst 10 rows (preview):")
    print("-" * 40)
    print("id,set")
    for task_id in training_ids[:5]:
        print(f"{task_id},training")
    for task_id in evaluation_ids[:5]:
        print(f"{task_id},evaluation")

    print("\n✓ Saved to: arc_v1_task_ids.csv")
    print("=" * 80)

if __name__ == '__main__':
    main()