# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_mapper.ipynb.

# %% auto 0
__all__ = ['Callback', 'init_cb', 'msg_cb', 'block_counter_cb', 'msg_bs_cb', 'CallbackHandler', 'Mapper']

# %% ../nbs/01_mapper.ipynb 2
from .imports import *
from .utils import *

# %% ../nbs/01_mapper.ipynb 4
class Callback:
    "Base class for callbacks for `Mapper`."

    def __init__(self, verbose=False, **kwargs):
        store_attr(**locals_to_params(locals()))

    def before_init(self, **kwargs):
        pass

    def after_init(self, **kwargs):
        pass

    def before_batch(self, **kwargs):
        pass

    def before_batch_rows(self, **kwargs):
        pass

    def after_batch_rows(self, **kwargs):
        pass

    def after_batch(self, **kwargs):
        pass


def init_cb(cb, **kwargs):
    if type(cb).__name__ == "type":
        cb = cb(**kwargs)
    return cb


class msg_cb(Callback):
    "A `Callback` to print out the `target` attribute of the `Mapper`."

    def __init__(self, target=[""], verbose=False):
        super().__init__(**locals_to_params(locals()))

    def printer(self, cls, **kwargs):
        for t in self.target:
            if hasattr(cls, t):
                msg.info(f"MSG CB: {getattr(cls, t)}", spaced=True)

    def before_init(self, cls, **kwargs):
        self.printer(cls, **kwargs)

    def before_batch(self, cls, **kwargs):
        self.printer(cls, **kwargs)

    def before_batch_rows(self, cls, **kwargs):
        self.printer(cls, **kwargs)

    def after_batch_rows(self, cls, **kwargs):
        self.printer(cls, **kwargs)

    def after_batch(self, cls, **kwargs):
        self.printer(cls, **kwargs)


class block_counter_cb(Callback):
    "A `Callback` to count the number of blocks."

    def before_init(self, cls, **kwargs):
        cls.block_counter = 0
        if self.verbose:
            msg.info(f"BLOCK COUNTER: {cls.block_counter}", spaced=True)

    def after_batch(self, cls, **kwargs):
        cls.block_counter += 1
        if self.verbose:
            msg.info(f"BLOCK COUNTER: {cls.block_counter}", spaced=True)


class msg_bs_cb(Callback):
    "A `Callback` to print the batch size."

    def before_batch(self, df, **kwargs):
        msg.info(f"DF BATCH SIZE: {len(df)}", spaced=True)


class CallbackHandler(Callback):
    def __init__(
        self, cbs=[], verbose=False, default=[block_counter_cb, msg_bs_cb], **kwargs
    ):
        cb_types = [type(cb).__name__ for cb in cbs]
        default_cbs = [cb for cb in default if cb.__name__ not in cb_types]
        self.cbs = default_cbs + cbs
        self.init_cbs(verbose=verbose, **kwargs)

    def init_cbs(self, **kwargs):
        self.cbs = [init_cb(cb, **kwargs) for cb in self.cbs]

    def before_init(self, **kwargs):
        [cb.before_init(**kwargs) for cb in self.cbs]

    def after_init(self, **kwargs):
        [cb.after_init(**kwargs) for cb in self.cbs]

    def before_batch(self, **kwargs):
        [cb.before_batch(**kwargs) for cb in self.cbs]

    def before_batch_rows(self, df, **kwargs):
        [cb.before_batch_rows(**kwargs) for _ in range(len(df)) for cb in self.cbs]

    def after_batch_rows(self, df, **kwargs):
        [cb.after_batch_rows(**kwargs) for _ in range(len(df)) for cb in self.cbs]

    def after_batch(self, **kwargs):
        [cb.after_batch(**kwargs) for cb in self.cbs]


class Mapper:
    """
    A class to map a function to a dataframe.
    The function should expect a dataframe as input and return a dataframe as output.
    """

    def __init__(
        self,
        udf=noop,  # A function that takes a dataframe as input and returns a dataframe as output.
        udf_kwargs={},  # The keyword arguments to pass to the `udf`.
        cbs=[],  # A list of `Callback`s to run before and after the mapping. It will add `block_counter_cb` and `msg_bs_cb` by default.
        verbose=False,  # Whether to print the batch size.
        **kwargs,
    ):
        self.cb_handler = CallbackHandler(cbs=cbs, verbose=verbose, **kwargs)
        self.cb_handler.before_init(cls=self, **kwargs)
        udf = partial(udf, **udf_kwargs)
        store_attr(**locals_to_params(locals()))
        self.cb_handler.after_init(cls=self, **kwargs)

    def map(self, df):
        return df.apply(self.udf, axis=1, result_type="expand")

    def __call__(self, df):
        self.cb_handler.before_batch(cls=self, df=df)
        self.cb_handler.before_batch_rows(cls=self, df=df)

        df = self.map(df)

        self.cb_handler.after_batch_rows(cls=self, df=df)
        self.cb_handler.after_batch(cls=self, df=df)
        return df

