import os
import tqdm
import shutil


import pandas as pd


def split_3d_front(data_path, out_path, split_csv):
    df = pd.read_csv(split_csv, header=None, names=['file', 'type'])
    file_type_map = {f: t for f, t in zip(df['file'], df['type'])}
    for file in tqdm.tqdm(os.listdir(data_path)):
        filename = os.path.splitext(file)[0]
        if filename not in file_type_map:
            continue
        data_type = file_type_map[filename]
        if not os.path.exists(os.path.join(out_path, data_type)):
            os.makedirs(os.path.join(out_path, data_type))
        shutil.copy(os.path.join(data_path, file), os.path.join(out_path, data_type, file))


if __name__ == '__main__':
    # https://github.com/tangjiapeng/DiffuScene  config/bedroom_threed_front_splits.csv
    split_csv = './bedroom_threed_front_splits.csv'
    split_3d_front('/anonymous/3D-FRONT-SIMPLE-OUTPUT/rooms/bedroom/',
                   '/anonymous/3D-FRONT-SIMPLE-OUTPUT/rooms/bedroom/', split_csv)
    # https://github.com/tangjiapeng/DiffuScene  config/livingroom_threed_front_splits.csv
    split_csv = './livingroom_threed_front_splits.csv'
    split_3d_front('/anonymous/3D-FRONT-SIMPLE-OUTPUT/rooms/living_room/',
                   '/anonymous/3D-FRONT-SIMPLE-OUTPUT/rooms/living_room/', split_csv)
