import os

user = os.environ.get("USER")
results_folder = os.path.join("/BASE/FOLDER/TO/STORE/WEIGHTS/AND/LOGS", "EfficientCVBench")
# PATH: /netscratch/<user>/slurm
slurm_output_folder = os.path.join("/FOLDER/FOR/SLURM/TO/WRITE/LOGS/TO", "slurm")


_ds_paths = {
    "cifar": "/PATH/TO/CIFAT",
    "tinyimagenet": "/PATH/TO/TINYIMAGENET",
    "stanford_cars": "/PATH/TO/CARS/SUPERFOLDER",
    "oxford_pet": "/PATH/TO/PET/SUPERFOLDER",
    "flowers102": "/PATH/TO/FLOWERS/SUPERFOLDER",
    "food101": "/PATH/TO/FOOD/SUPERFOLDER",
    "aircraft": "/PATH/TO/AIRCRAFT/SUPERFOLDER",
    "fornet": "/PATH/TO/FORNET",
    "counteranimal": "/PATH/TO/CounterAnimal/LAION-final",
}


def ds_path(dataset, args=None):
    """Get the (base) path for any dataset.

    Args:
    -----
        dataset (str): The dataset I'm looking for.
        args (DotDict, optional): Run args. If args.custom_dataset_path is set, this one is always returned.

    Returns:
    --------
        str: Path to the dataset root folder.
    """
    if args is not None and "custom_dataset_path" in args and args.custom_dataset_path is not None:
        return args.custom_dataset_path
    return _ds_paths[dataset]