import argparse


parser = argparse.ArgumentParser()

parser.add_argument("--task", type=str, choices=('finetune', 'train'))
parser.add_argument("--model", type=str, choices=('base', 'sup', 'unsup'))
parser.add_argument("--dataset", type=str, choices=('sst', 'cfimdb'), default='sst')
parser.add_argument("--train-mode", type=str, choices=('last-linear', 'full-model'), default='last-linear')


args = parser.parse_args()

if args.task == "finetune":
    if args.model == "sup":
        import finetune_scripts.sup
    elif args.model == "unsup":
        import finetune_scripts.unsup
    else:
        raise Exception("Invalid task for finetuning.")
else:
    match (args.model, args.dataset, args.train_mode):
        case ('base', 'sst', 'last-linear'):
            import train_scripts.base_sst_ll
        case ('base', 'sst', 'full-model'):
            import train_scripts.base_sst_fm
        case ('base', 'cfimdb', 'last-linear'):
            import train_scripts.base_cfimdb_ll
        case ('base', 'cfimdb', 'full-model'):
            import train_scripts.base_cfimdb_fm
        case ('unsup', 'sst', 'last-linear'):
            import train_scripts.unsup_sst_ll
        case ('unsup', 'sst', 'full-model'):
            import train_scripts.unsup_sst_fm
        case ('unsup', 'cfimdb', 'last-linear'):
            import train_scripts.unsup_cfimdb_ll
        case ('unsup', 'cfimdb', 'full-model'):
            import train_scripts.unsup_cfimdb_fm
        case ('sup', 'sst', 'last-linear'):
            import train_scripts.sup_sst_ll
        case ('sup', 'sst', 'full-model'):
            import train_scripts.sup_sst_fm
        case ('sup', 'cfimdb', 'last-linear'):
            import train_scripts.sup_cfimdb_ll
        case ('sup', 'cfimdb', 'full-model'):
            import train_scripts.sup_cfimdb_fm