import re


def get_train_dataset_name(config_train_name: str, shot: int, seed: int) -> str:
    """Returns a modified version of a given train dataset with proper shot and seed numbers"""
    # We replace the given config name's shot number with ours, unless they are already equal.
    if '1shot' in config_train_name:  # Replace shot number
        config_train_name = config_train_name.replace('1shot', str(shot) + 'shot')
    elif f'{shot}shot' in config_train_name:  # Leave shot as-is
        assert False
        # pass
    else:
        assert False

    # Modify or add seed number in name.
    if 'seed' in config_train_name:
        # Replace current seed number
        # Example: name_seedX with name_seed{SEED}, where SEED is our current seed
        assert len(re.findall('seed.', config_train_name)) == 1
        config_train_name = re.sub('seed.', 'seed{}'.format(seed), config_train_name)
    else:  # Append seed string
        config_train_name = config_train_name + f'_seed{seed}'

    return config_train_name