import json

# Read subset classes
with open('subset_classes.txt', 'r') as f:
    subset_classes = [line.strip() for line in f.readlines()]

# Read class names mapping
class_id_to_name = {}
with open('imagenet_class_id_index_name.txt', 'r') as f:
    for line in f:
        parts = line.strip().split(' ', 2)
        if len(parts) == 3:
            class_id, _, class_name = parts
            class_id_to_name[class_id] = class_name

# Create the subset index mapping with comments
subset_index = {}
for idx, class_id in enumerate(subset_classes):
    class_name = class_id_to_name.get(class_id, "Unknown")
    subset_index[str(idx)] = f"{class_name} # corresponding to {class_id}"

# Format the JSON string manually to include comments
json_str = "{\n"
for idx in range(len(subset_index)):
    value = subset_index[str(idx)]
    json_str += f'    "{idx}": "{value}"'
    if idx < len(subset_index) - 1:
        json_str += ","
    json_str += "\n"
json_str += "}"

# Save to JSON file
with open('subset_imagenet_class_index.json', 'w') as f:
    f.write(json_str)