from os.path import expanduser, join
from dataManager.gFunDataset import SimpleGfunDataset


def load_from_pickle(path, dataset_name, nrows):
    import pickle

    filepath = join(path, f"{dataset_name}_{nrows}.pkl")

    with open(filepath, "rb") as f:
        loaded = pickle.load(f)
    print(f"- Loaded dataset from {filepath}")
    loaded.show_dimension()
    return loaded


def get_dataset(datasetp_path, args):
    dataset = SimpleGfunDataset(
        dataset_name="rai",
        datadir=datasetp_path,
        textual=True,
        visual=False,
        multilabel=False,
        set_tr_langs=args.tr_langs,
        set_te_langs=args.te_langs
    )
    return dataset