# Code copied without modification from https://github.com/louisabraham/tqdm_joblib
# Original work by Louis Abraham
# Licensed under CC BY-SA 4.0: https://creativecommons.org/licenses/by-sa/4.0/

import contextlib

import joblib
from tqdm.autonotebook import tqdm


@contextlib.contextmanager
def tqdm_joblib(*args, **kwargs):
    """Context manager to patch joblib to report into tqdm progress bar
    given as argument"""

    tqdm_object = tqdm(*args, **kwargs)

    class TqdmBatchCompletionCallback(joblib.parallel.BatchCompletionCallBack):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)

        def __call__(self, *args, **kwargs):
            tqdm_object.update(n=self.batch_size)
            return super().__call__(*args, **kwargs)

    old_batch_callback = joblib.parallel.BatchCompletionCallBack
    joblib.parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
    try:
        yield tqdm_object
    finally:
        joblib.parallel.BatchCompletionCallBack = old_batch_callback
        tqdm_object.close()


def ParallelPbar(desc=None, **tqdm_kwargs):
    class Parallel(joblib.Parallel):
        def __call__(self, it):
            it = list(it)
            if self.n_jobs == 1:
                it = tqdm(it, desc=desc, **tqdm_kwargs)
                return super().__call__(it)
            with tqdm_joblib(total=len(it), desc=desc, **tqdm_kwargs):
                return super().__call__(it)

    return Parallel
