import os
import argparse

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--method", "-m", default="Vanilla", help="Method")
    parser.add_argument("--dataset", "-d", default="wildcam", help="Dataset")
    # parser.add_argument("--gpu", "-g", default=0, type=int, help="Gpu ID")
    parser.add_argument(
        "--backbone", "-b", default="resnet50", help="Backbone"
    )
    parser.add_argument(
        "--n_trials", "-n", default=1, type=int, help="Repeat times"
    )

    args = parser.parse_args()

    ###############################################################################

    data_root = "/home/zhaoxin/data/DG"
    method_name = args.method.lower()

    ###############################################################################

    output_dir = [
        os.path.join(
            'output', args.method, args.dataset + '_' + args.backbone,
            str(i + 1)
        ) for i in range(3)
    ]
    os.system(
        f'CUDA_VISIBLE_DEVICES=0 '
        f'python tools/train.py '
        f'--root {data_root} '
        f'--trainer {args.method} '
        f'--backbone {args.backbone} '
        f'--dataset-config-file configs/datasets/dg/{args.dataset}.yaml '
        f'--config-file configs/trainers/dg/{method_name}/{args.dataset}.yaml '
        f'--output-dir {output_dir[0]} '
        f'--seed 0 '
        '& '
        f'CUDA_VISIBLE_DEVICES=1 '
        f'python tools/train.py '
        f'--root {data_root} '
        f'--trainer {args.method} '
        f'--backbone {args.backbone} '
        f'--dataset-config-file configs/datasets/dg/{args.dataset}.yaml '
        f'--config-file configs/trainers/dg/{method_name}/{args.dataset}.yaml '
        f'--output-dir {output_dir[1]} '
        f'--seed 1 '
        '& '
        f'CUDA_VISIBLE_DEVICES=2 '
        f'python tools/train.py '
        f'--root {data_root} '
        f'--trainer {args.method} '
        f'--backbone {args.backbone} '
        f'--dataset-config-file configs/datasets/dg/{args.dataset}.yaml '
        f'--config-file configs/trainers/dg/{method_name}/{args.dataset}.yaml '
        f'--output-dir {output_dir[2]} '
        f'--seed 2'
    )
