import os
from concurrent.futures import ProcessPoolExecutor

import numpy as np
from tqdm import tqdm

from .stats import get_numpy_stats
from ..logger.logger import Logger
from ..utils import get_class_name


def save(folder: str, data: np.ndarray, start_index: int = 0, n_processes: int = 8) -> None:
    Logger.debug(
        f'{get_class_name(save)} - '
        f'folder: {folder}, '
        f'data: {get_numpy_stats(data)}, '
        f'start_index: {start_index}, '
        f'n_processes: {n_processes}'
    )
    assert len(data.shape) >= 1, f'data must have at least 1 dimension: {data.shape}'
    n_samples: int = data.shape[0]
    if folder != '':
        os.makedirs(folder, exist_ok=True)
    with ProcessPoolExecutor(max_workers=n_processes) as executor:
        list(tqdm(executor.map(
            np.save,
            [f'{folder}/{start_index + i}.npy' for i in range(n_samples)],
            data
        ), total=n_samples))
