""""Utility functions for meta learning setup"""

from codes.utils.task_family import TaskFamily


def get_input_and_target(
    config, batch_size, task_family, target_function=None, mode=None, task_id=None
):
    """Method to sample the input and target function"""
    if target_function is None:
        target_function = task_family.sample_task(task_id)

    data = task_family.sample_inputs(batch_size=batch_size, mode=mode)

    graphs, queries, indices, world_graph = data
    targets = target_function(indices, mode)
    return graphs, queries, targets, target_function, world_graph


def bootstrap_task_family(config):
    """Method to bootstrap the task family"""

    # return TaskFamily(config=config, mode="train")

    meta_mode_list = ["train", "valid", "test"]
    return [
        TaskFamily(config=config, meta_mode=meta_mode) for meta_mode in meta_mode_list
    ]
