import pandas as pd

grandparent_classes = {
    'water animal': ['aquatic mammals', 'fish'],
    'invertebrate': ['insects', 'non-insect invertebrates', 'reptiles'],
    'mammal': ['large carnivores', 'large omnivores and herbivores', 'medium-sized mammals', 'small mammals'],
    'plant': ['flowers', 'trees'],
    'household item': ['food containers', 'household electrical devices', 'household furniture'],
    'food': ['fruit and vegetables'],
    'building': ['large man-made outdoor things'],
    'nature': ['large natural outdoor scenes'],
    'people': ['people'],
    'vehicles': ['vehicles 1', 'vehicles 2']
}

# speed vs storage
# speed vs performance
# dataset_size

parent_classes = {
    'aquatic mammals': ['beaver', 'dolphin', 'otter', 'seal', 'whale'],
    'fish': ['aquarium_fish', 'flatfish', 'ray', 'shark', 'trout'],
    'flowers': ['orchid', 'poppy', 'rose', 'sunflower', 'tulip'],
    'food containers': ['bottle', 'bowl', 'can', 'cup', 'plate'],
    'fruit and vegetables': ['apple', 'mushroom', 'orange', 'pear', 'sweet_pepper'], # fruit & vegetables
    'household electrical devices': ['clock', 'keyboard', 'lamp', 'telephone', 'television'],
    'household furniture': ['bed', 'chair', 'couch', 'table', 'wardrobe'],
    'insects': ['bee', 'beetle', 'butterfly', 'caterpillar', 'cockroach'],
    'large carnivores': ['bear', 'leopard', 'lion', 'tiger', 'wolf'],
    'large man-made outdoor things': ['bridge', 'castle', 'house', 'road', 'skyscraper'], # building & infrastructure
    'large natural outdoor scenes': ['cloud', 'forest', 'mountain', 'plain', 'sea'], # split into 2 parents
    'large omnivores and herbivores': ['camel', 'cattle', 'chimpanzee', 'elephant', 'kangaroo'],
    'medium-sized mammals': ['fox', 'porcupine', 'possum', 'raccoon', 'skunk'],
    'non-insect invertebrates': ['crab', 'lobster', 'snail', 'spider', 'worm'],
    'people': ['baby', 'boy', 'girl', 'man', 'woman'],
    'reptiles': ['crocodile', 'dinosaur', 'lizard', 'snake', 'turtle'],
    'small mammals': ['hamster', 'mouse', 'rabbit', 'shrew', 'squirrel'],
    'trees': ['maple_tree', 'oak_tree', 'palm_tree', 'pine_tree', 'willow_tree'],
    'vehicles 1': ['bicycle', 'bus', 'motorcycle', 'pickup_truck', 'train'],
    'vehicles 2': ['lawn_mower', 'rocket', 'streetcar', 'tank', 'tractor']
}

# df = pd.DataFrame.from_dict(parent_classes, orient='index').transpose()
# print(df)

# Create a list of dictionaries for each row of the final DataFrame
rows = []
for grandparent, parent_list in grandparent_classes.items():
    for parent in parent_list:
        for item in parent_classes[parent]:
            rows.append({
                'class': item,
                'parent class': parent,
                # 'grandparent class': grandparent,
                'root class': 'Root',
                'seen flag': 1
            })

df = pd.DataFrame(rows)
df.to_csv('cifar100_hierarchy.csv', index=False, header=False)


