from concurrent.futures import ProcessPoolExecutor

import numpy as np
from tqdm import tqdm

from ..logger.logger import Logger
from ..utils import get_class_name


def load(folder: str, n_samples: int, start_index: int = 0, n_processes: int = 8) -> np.ndarray:
    Logger.debug(
        f'{get_class_name(load)} - '
        f'folder: {folder}, '
        f'n_samples: {n_samples}, '
        f'start_index: {start_index}, '
        f'n_processes: {n_processes}'
    )

    with ProcessPoolExecutor(max_workers=n_processes) as executor:
        data = np.array(list(tqdm(executor.map(
            np.load,
            [f'{folder}/{start_index + i}.npy' for i in range(n_samples)]
        ), total=n_samples)))

    return data
