# all_labels = ['n01443537', 'n01484850', 'n01491361', 'n01494475', 'n01496331', 
# 'n01498041', 'n01629819', 'n01630670', 'n01641577', 'n01644373', 'n01644900', 'n01664065', 
# 'n01665541', 'n01667114', 'n01682714', 'n01685808', 'n01687978', 'n01688243', 'n01689811', 
# 'n01694178', 'n01695060', 'n01697457', 'n01698640', 'n01728572', 'n01728920', 'n01729322', 
# 'n01749939', 'n01751748', 'n01770393', 'n01773157', 'n01773549', 'n01773797', 'n01914609', 
# 'n01917289', 'n01924916', 'n01945685', 'n01950731', 'n01978455', 'n01980166', 'n01981276', 
# 'n01983481', 'n01984695', 'n01985128', 'n01986214', 'n02066245', 'n02071294', 'n02074367', 
# 'n02077923', 'n02056570', 'n02128385', 'n02128757', 'n02129165', 'n02129604', 'n02130308', 
# 'n02167151', 'n02168699', 'n02169497',
# 'n02172182', 'n02190166', 'n02206856', 'n02219486', 'n02264363', 'n02268443', 'n02268853',
# 'n02277742', 'n02279972', 'n02280649', 'n02281406', 'n02281787', 'n02317335', 'n02319095',
# 'n02321529', 'n02325366', 'n02326432', 'n02328150', 'n02342885', 'n02346627', 'n02356798',
# 'n02361337', 'n02363005', 'n02364673', 'n02397096', 'n02398521', 'n02422106', 'n02422699',
# 'n02423022', 'n02437312', 'n02443114', 'n02443484', 'n02444819', 'n02445715', 'n02447366',
# 'n02454379', 'n02457408', 'n02480495', 'n02480855', 'n02481823', 'n02483362', 'n02483708', 
# 'n02484975', 'n02486410', 'n02487347', 'n02488291', 'n02488702', 'n02489166', 'n02493793', 
# 'n02494079', 'n02526121', 'n02536864', 'n02606052', 'n02641379', 'n02643566', 'n02655020', 
# 'n02666196', 'n02676566', 'n02747177', 'n02783161', 'n02788148', 'n02794156', 
# 'n02795169', 'n02814860', 'n02825657', 'n02834397', 'n02840245', 'n02906734', 'n02909870', 
# 'n02971356', 'n02978881', 'n02979186', 'n02980441', 'n02988304', 'n03000134', 'n03014705',
# 'n03016953', 'n03017168', 'n03018349', 'n03026506', 'n03047690',  'n03063689',
# 'n03065424', 'n03075370', 'n03109150', 'n03124043', 'n03125729', 'n03131574', 'n03133878',
# 'n03160309', 'n03207743', 'n03220513', 'n03223299', 'n03240683', 'n03271574', 'n03272010',
# 'n03291819',  'n03337140', 'n03355925', 'n03388043', 'n03388183', 'n03400231',
# 'n03425413', 'n03445777', 'n03447721', 'n03481172', 'n03482405', 'n03485794', 'n03492542',
# 'n03494278', 'n03498962', 'n03527444', 'n03530642', 'n03532672', 'n03584829', 'n03598930',
# 'n03627232', 'n03633091', 'n03637318', 'n03657121',  'n03666591', 'n03680355',
# 'n03690938', 'n03691459', 'n03706229', 'n03709823', 'n03710193', 'n03717622', 'n03729826', 
#  'n03743016', 'n03764736', 'n03773504', 'n03777754', 'n03786901', 'n03788195', 
# 'n03804744', 'n03837869', 'n03841143', 'n03854065', 'n03871628', 'n03874599', 'n03876231', 
# 'n03908618', 'n03908714', 'n03929660', 'n03930313', 'n03933933', 'n03935335', 'n03944341', 
# 'n03954731', 'n03956157', 'n03958227', 'n03976467', 'n03998194', 'n04004767', 'n04033901',
# 'n04041544', 'n04044716', 'n04049303', 'n04069434', 'n04116512', 'n04118776', 'n04125021',
# 'n04127249',  'n04141975', 'n04149813', 'n04153751', 'n04179913', 'n04201297',
# 'n04239074', 'n04243546', 'n04254777', 'n04265275', 'n04275548', 'n04277352', 'n04326547',
# 'n04355338', 'n04366367', 'n04372370', 'n04389033', 'n04418357', 'n04423845', 'n04442312',
# 'n04465501', 'n04485082', 'n04523525', 'n04525038', 'n04532670', 'n04554684', 'n04599235', 
# 'n04604644', 'n06359193', 'n06785654', 'n07565083', 'n07684084',  'n07697313', 
# 'n07715103', 'n07716358',  'n07720875', 'n07730033', 'n07734744', 'n07745940', 
# 'n07749582', 'n07753275', 'n07802026', 'n09246464', 'n09256479', 'n09399592', 'n09468604', 
# 'n09472597', 'n11939491', 'n12267677', 'n12620546', 'n12768682', 'n12985857', 'n12998815', 
# 'n13037406', 'n13040303', 'n13044778', 'n13054560']
# from shutil import copytree
# import os
# train_root_dir= "/nobackup-slow/dataset/ILSVRC-2012/train"
# root_dir= "/nobackup-slow/dataset/ILSVRC-2012/ood_train"
# for label in all_labels:
#     src = os.path.join(train_root_dir, label)
#     dst = os.path.join(root_dir, label)
#     if not os.path.exists(dst):
#         print("copying label # ", label)
#         copytree(src, dst)
#     # else:
#     #     print(f"{label} already exists")
# print(len(all_labels)) # 277


import os
from shutil import copyfile

src_dir = '/nobackup-slow/dataset/shape/output'
type_list = [ "Triangle", "Square", "Pentagon", "Hexagon", "Heptagon", "Octagon", "Nonagon", "Circle", "Star"]
for _type in type_list:
    dst_dir_train = f'/nobackup-slow/dataset/shape/train/{_type}'
    dst_dir_val = f'/nobackup-slow/dataset/shape/val/{_type}'
    if not os.path.exists(dst_dir_train):
        os.makedirs(dst_dir_train)
    if not os.path.exists(dst_dir_val):
        os.makedirs(dst_dir_val)
    count = 0
    for f in os.listdir(src_dir):
        if count >= 2200: 
            break
        if f.startswith(_type):
            src = os.path.join(src_dir, f)
            if count < 2000:
                dst_train = os.path.join(dst_dir_train, f)
                if not os.path.exists(dst_train):
                    print("Train copying file: ", f)
                    copyfile(src, dst_train)
            else:
                dst_val = os.path.join(dst_dir_val, f)
                if not os.path.exists(dst_val):
                    print("Val copying file: ", f)
                    copyfile(src, dst_val)

            count += 1