import pathlib
import json


def get_train_val_paths(all_paths, path_to_train_val_pkl):
    """"
    Split a list of all paths to patients' data into train & validation parts using patients' IDs.

    Parameters
    ----------
    all_paths: list
        An output of `get_paths_to_patient_files`.
    path_to_train_val_pkl: str
        A path to a pkl file storing train & validation IDs.

    Returns
    -------
    (list, list)
        Two lists of paths to train & validation data samples.
    """
    path_to_train_val_pkl = pathlib.Path(path_to_train_val_pkl)
    with open(path_to_train_val_pkl) as f:
        train_val_split = json.load(f)

    train_paths = [path for path in all_paths
                   if any(patient_id + '_ct.nii.gz' in str(path[0]) for patient_id in train_val_split['train'])]

    val_paths = [path for path in all_paths
                 if any(patient_id + '_ct.nii.gz' in str(path[0]) for patient_id in train_val_split['val'])]

    return train_paths, val_paths