import joblib

from tqdm.auto import tqdm

from preprocessing_utils import eic_text_preprocessing

from TrialMatchAI.src.Preporcessor.preprocess_clinical_notes import tokenize_clinical_note

import pandas as pd

import os

memory = joblib.Memory(".")

def ParallelExecutor(use_bar="tqdm", **joblib_args):

    """Utility for tqdm progress bar in joblib.Parallel"""

    all_bar_funcs = {

        "tqdm": lambda args: lambda x: tqdm(x, **args),

        "False": lambda args: iter,

        "None": lambda args: iter,

    }

    def aprun(bar=use_bar, **tq_args):

        def tmp(op_iter):

            if str(bar) in all_bar_funcs.keys():

                bar_func = all_bar_funcs[str(bar)](tq_args)

            else:

                raise ValueError("Value %s not supported as bar type" % bar)

            # Pass n_jobs from joblib_args

            return joblib.Parallel(n_jobs=joblib_args.get("n_jobs", 10))(bar_func(op_iter))

        return tmp

    return aprun

class Preprocessor:

    def __init__(self, id_list, n_jobs):

        self.id_list = id_list

        self.n_jobs = n_jobs

    def preprocess_clinical_trials_text(self):

        parallel_runner = ParallelExecutor(n_jobs=self.n_jobs)(total=len(self.id_list))

        X = parallel_runner(

            joblib.delayed(eic_text_preprocessing)(

            [_id]

            )

            for _id in self.id_list

        )     

        return pd.concat(X).reset_index(drop=True)

    def preprocess_patient_clinical_notes(self):

        parallel_runner = ParallelExecutor(n_jobs=self.n_jobs)(total=len(self.id_list))

        X = parallel_runner(

            joblib.delayed(tokenize_clinical_note)(

            [_id]

            )

            for _id in self.id_list

        )     

        return pd.concat(X).reset_index(drop=True)