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

parser = argparse.ArgumentParser()
parser.add_argument('-t', '--task', type=str, required=True)
parser.add_argument('-f', '--format', type=str, default='seg')
parser.add_argument('-ft', '--finetune', action='store_true')
parser.add_argument('-d', '--device', type=str, default='0')
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')
trainer_script = './train_net_mrcnn.py'

def generate_dataset(args):
    ann_file = pjoin(DATASET_FOLDER_PATH, f'data/coco_{args.format}/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/{args.format}/{args.task}{"_finetune" if args.finetune else ""}/{SPLIT}'
    split_path = args.split_path
    device = args.device.split(',')
    cli = f'python {trainer_script} ' + \
            f'--config-file ./configs/my_mask_rcnn_R_101_FPN_3x.yaml ' + \
            f'--dataset_root {split_path} ' + \
            f'--num_classes {NUM_CLASSES_DICT[args.task]} ' + \
            f'--manual_device {",".join(device)} ' + \
            f'--num-gpus {len(device)} ' + \
            f'--resume ' + \
            f'OUTPUT_DIR {output_dir} '
    if args.finetune and not os.path.exists(output_dir):
        cli += f'MODEL.WEIGHTS ./mask_rcnn_R_101_FPN_3x.pkl '
    os.system(cli)
    if os.path.exists(split_path):
        shutil.rmtree(split_path)
    

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