import os
from os.path import join as pjoin
import shutil
import time
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-t', '--task', type=str, default='interactable')
parser.add_argument('-sp', '--split_path', type=str, default='./dataset')
args = parser.parse_args()

# SPLIT = '613'
# TRAIN_FOLD = '0,1,4,5,7,9'
# VAL_FOLD = '3'
# TEST_FOLD = '2,6,8'
# FOLD_FILE = 'fold_app.csv'

# SPLIT = 'genre'
# TRAIN_FOLD = '0'
# VAL_FOLD = '1'
# TEST_FOLD = '2'
# FOLD_FILE = 'fold_genre.csv'

SPLIT = 'cat'
TRAIN_FOLD = '0'
VAL_FOLD = '1'
TEST_FOLD = '2'
FOLD_FILE = 'fold_cat.csv'

DATASET_FOLDER_PATH = '../../dataset/'
NUM_CLASSES_DICT = {'interaction': 53, 'semantics': 766, 'interactable': 1}

split_script = pjoin(DATASET_FOLDER_PATH, 'split_coco.py')
det_script = './detect.py'
region2coco_script = './xianyu2coco.py'


def generate_dataset(args):
    ann_file = pjoin(DATASET_FOLDER_PATH, f'data/coco_det/annotations/{args.task}.json')
    img_dir = pjoin(DATASET_FOLDER_PATH, 'data/coco_merged/images', args.task)
    fold_file = pjoin(DATASET_FOLDER_PATH, FOLD_FILE)
    split_path = args.split_path
    if os.path.exists(split_path):
        shutil.rmtree(split_path)
    cli = f'python {split_script} ' + \
            f'--ann_file {ann_file} ' + \
            f'--img_dir {img_dir} ' + \
            f'--output_path {split_path} ' + \
            f'--fold_file {fold_file} ' + \
            f'--train_fold {TRAIN_FOLD} ' + \
            f'--val_fold {VAL_FOLD} ' + \
            f'--test_fold {TEST_FOLD} '
    os.system(cli)


def main(args):
    generate_dataset(args)
    output_dir = f'./output/det/{args.task}/{SPLIT}'
    os.makedirs(output_dir, exist_ok=True)
    split_path = args.split_path
    det_cli = f'python {det_script} ' + \
            f'--test_folder {pjoin(split_path, "images/instances_test2017")} ' + \
            f'--output_folder {pjoin(output_dir, "regions")} '
    to_coco_cli = f'python {region2coco_script} ' + \
            f'--xianyu_dir {pjoin(output_dir, "regions")} ' + \
            f'--ann {pjoin(split_path, "annotations/instances_test2017.json")} ' + \
            f'--output_file {pjoin(output_dir, "results.json")} ' + \
            f'--filter_oof '
    os.system(det_cli)
    os.system(to_coco_cli)
    if os.path.exists(split_path):
        shutil.rmtree(split_path)


if __name__ == '__main__':
    assert args.task in ['interaction', 'semantics', 'interactable'], 'expected task: interaction, semantics or interactable'
    assert args.task == 'interactable', 'only support interactable'
    main(args)