#!/usr/bin/env python3
"""
Generate a comprehensive JSON file with official ARC V1 task IDs
"""

import json

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 comprehensive JSON
    data = {
        "version": "ARC-AGI-1 (Official)",
        "source": "https://github.com/fchollet/ARC-AGI v1.0.2",
        "description": "Official ARC-AGI Version 1 task IDs - 400 training + 400 evaluation tasks",
        "total_tasks": len(training_ids) + len(evaluation_ids),
        "training": {
            "count": len(training_ids),
            "task_ids": training_ids
        },
        "evaluation": {
            "count": len(evaluation_ids),
            "task_ids": evaluation_ids
        },
        "all_task_ids": sorted(training_ids + evaluation_ids)
    }

    # Save to JSON
    with open('arc_v1_official_task_ids.json', 'w') as f:
        json.dump(data, f, indent=2)

    print("=" * 80)
    print("ARC V1 Official Task IDs")
    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 Training IDs:")
    for task_id in training_ids[:10]:
        print(f"  {task_id}")

    print("\nFirst 10 Evaluation IDs:")
    for task_id in evaluation_ids[:10]:
        print(f"  {task_id}")

    print("\n✓ Saved to: arc_v1_official_task_ids.json")
    print("=" * 80)

    # Also create a simple combined text file
    with open('arc_v1_all_ids.txt', 'w') as f:
        f.write("# ARC-AGI Version 1 Official Task IDs\n")
        f.write("# Source: https://github.com/fchollet/ARC-AGI v1.0.2\n")
        f.write(f"# Total: {len(training_ids) + len(evaluation_ids)} tasks\n\n")
        f.write("## Training Tasks (400)\n")
        for task_id in training_ids:
            f.write(f"{task_id}\n")
        f.write("\n## Evaluation Tasks (400)\n")
        for task_id in evaluation_ids:
            f.write(f"{task_id}\n")

    print("✓ Saved to: arc_v1_all_ids.txt")

if __name__ == '__main__':
    main()