import logging
import codecs
import os
import numpy as np
import pandas as pd
import ekorpkit.utils.batch.batcher as batcher
from collections import OrderedDict
from functools import partial, reduce
from ekorpkit.utils import print_status
from ekorpkit.utils.batch import decorator_apply
from ekorpkit.utils.func import check_max_len, check_min_len, elapsed_timer
from ekorpkit.io.file import save_dataframe as save_dataframe_
from ekorpkit import eKonf
from ekorpkit.ekonf import apply_pipe
from tqdm.auto import tqdm


log = logging.getLogger(__name__)


def apply(
    func,
    series,
    description=None,
    verbose=False,
    use_batcher=True,
    minibatch_size=None,
    num_workers=None,
    **kwargs,
):
    batcher_instance = batcher.batcher_instance
    if use_batcher and batcher_instance is not None:
        if batcher_instance is not None:
            batcher_minibatch_size = batcher_instance.minibatch_size
        else:
            batcher_minibatch_size = 1000
        if minibatch_size is None:
            minibatch_size = batcher_minibatch_size
        if num_workers is not None:
            batcher_instance.procs = num_workers
        if batcher_instance.procs > 1:
            batcher_instance.minibatch_size = min(
                int(len(series) / batcher_instance.procs) + 1, minibatch_size
            )
            log.info(
                f"Using batcher with minibatch size: {batcher_instance.minibatch_size}"
            )
            results = decorator_apply(func, batcher_instance, description=description)(
                series
            )
            if batcher_instance is not None:
                batcher_instance.minibatch_size = batcher_minibatch_size
            return results

    if batcher_instance is None:
        log.warning("Warning: batcher not initialized")
    tqdm.pandas(desc=description)
    return series.progress_apply(func)


def apply_pipeline(df, pipeline, pipeline_args, update_args={}, verbose=True):
    pipeline = eKonf.to_dict(pipeline)
    pipeline_args = eKonf.to_dict(pipeline_args)

    pipeline_targets = []
    pipes = OrderedDict()
    if isinstance(pipeline, list):
        for pipe in pipeline:
            pipes[pipe] = pipe
    elif isinstance(pipeline, str):
        pipes[pipeline] = pipeline
    else:
        pipes = pipeline

    if pipes is None or len(pipes) == 0:
        log.warning("No pipeline specified")
        return df

    log.info(f"Applying pipeline: {pipes}")
    for pipe, pipe_arg_name in pipes.items():
        args = pipeline_args.get(pipe_arg_name, {}).copy()
        if pipe != pipe_arg_name:
            args_override = pipeline_args.get(pipe, {})
            args.update(args_override)
        if args and isinstance(args, dict):
            args.update(update_args)
            pipeline_targets.append(args)

    return reduce(apply_pipe, pipeline_targets, df)


def split_column(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    source = args.get("source")
    target = args.get("target")
    _split = args.get(eKonf.Keys.SPLIT)
    if source is None:
        log.warning("No source specified")
        return df
    log.info(f"Split column: {args}")
    df[target] = df[source].str.split(**_split)

    if verbose:
        print(df.tail())
    return df


def eval_columns(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    expressions = args.get("expressions", None)
    engine = args.get("engine", None)
    eval_at = args.get("eval_at", "dataframe")
    if expressions is None:
        log.warning("No expressions specified")
        return df
    log.info(f"Eval columns: {args}")
    if eval_at == "dataframe":
        if isinstance(expressions, list):
            for expr in expressions:
                df.eval(expr, engine=engine, inplace=True)
        else:
            for col, expr in expressions.items():
                df[col] = df.eval(expr, engine=engine)
    elif eval_at == "python":
        if isinstance(expressions, dict):
            for col, expr in expressions.items():
                df[col] = eval(expr)
        else:
            log.warning("Expressions must be a dict to eval at python")
    else:
        if isinstance(expressions, list):
            for expr in expressions:
                pd.eval(expr, engine=engine, inplace=True, target=df)
        else:
            for col, expr in expressions.items():
                df[col] = pd.eval(expr, engine=engine)

    if verbose:
        print(df.tail())
    return df


def combine_columns(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    columns = args.get("columns", None)
    if columns is None:
        log.warning("No columns specified")
        return df
    log.info(f"Combining columns: {args}")
    for col in columns:
        df[col].fillna("", inplace=True)
        df[col] = df[col].astype(str)
    separator = codecs.decode(args["separator"], "unicode_escape")
    df[args["into"]] = df[columns].agg(separator.join, axis=1)
    return df


def drop(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    labels = args.get("labels", None)
    axis = args.get("axis", 1)
    columns = args.get("columns", None)
    index = args.get("index", None)
    level = args.get("level", None)
    errors = args.get("errors", "ignore")
    if verbose:
        log.info(f"droping: {args}")
    df = df.drop(
        columns=columns,
        axis=axis,
        labels=labels,
        index=index,
        level=level,
        errors=errors,
    )
    if verbose:
        print(df.head())
    return df


def melt(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    id_vars = args.get("id_vars", None)
    value_vars = args.get("value_vars", None)
    var_name = args.get("var_name", None)
    value_name = args.get("value_name", None)
    col_level = args.get("col_level", None)
    ignore_index = args.get("ignore_index", True)
    if id_vars is None:
        if verbose:
            log.warning("No id_vars specified")
        return df
    if verbose:
        log.info(f"Melting columns: {args}")
    id_vars = list(id_vars)
    if value_vars:
        if isinstance(value_vars, str):
            value_vars = eval(value_vars)
        else:
            value_vars = list(value_vars)
    df = df.melt(
        id_vars=id_vars,
        value_vars=value_vars,
        var_name=var_name,
        value_name=value_name,
        col_level=col_level,
        ignore_index=ignore_index,
    )
    if verbose:
        print(df.head())
    return df


def plot(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    plot_cfg = args.get("visualize", {}).get("plot", {})
    if "_target_" not in plot_cfg:
        log.warning("No target specified")
        return df
    if verbose:
        log.info(f"Plotting: {plot_cfg}")
    subset = args.get("subset", {})
    col = subset.get("column", None)
    values = subset.get("values", None)
    titles = subset.get("titles", None)
    output_file = args.get("output_file", None)
    output_dir = args.get("output_dir", ".")
    if subset and col:
        title_prefix = plot_cfg["figure"].get("title", None)
        title_prefix = title_prefix + " - " if title_prefix else ""
        if isinstance(values, str):
            values = [values]
        if titles and isinstance(titles, str):
            titles = [titles]
        for i, val in enumerate(values):
            if verbose:
                log.info(f"Plotting subset: {col} == {val}")
            df_sub = df[df[col] == val]
            if output_file and "{}" in output_file:
                output_path = output_file.replace("{}", str(val))
                output_path = f"{output_dir}/{output_path}"
                plot_cfg["savefig"]["fname"] = output_path
            if titles and len(titles) > i:
                plot_cfg["figure"]["title"] = title_prefix + titles[i]
            if verbose:
                log.info(f"Plotting: {plot_cfg}")
                # print(df_sub.head())
            eKonf.instantiate(plot_cfg, df=df_sub)
    else:
        eKonf.instantiate(plot_cfg, df=df)

    return df


def pivot(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    index = args.get("index", None)
    columns = args.get("columns", None)
    values = args.get("values", None)
    fillna = args.get("fillna", None)
    reset_index = args.get("reset_index", True)
    if index is None:
        log.warning("No index specified")
        return df
    if columns is None:
        log.warning("No columns specified")
        return df
    if values is None:
        log.warning("No values specified")
        return df
    if verbose:
        log.info(f"Pivoting columns: {args}")
    if len(values) == 1:
        values = values[0]
    df = df.pivot(index=index, columns=columns, values=values)
    # df = df.set_index(index)[values].unstack()
    if fillna is not None:
        df.fillna(fillna, inplace=True)
    if reset_index:
        if verbose:
            log.info(f"Resetting index, nlevels of columns: {df.columns.nlevels}")
        df.reset_index(inplace=True)
        if df.columns.nlevels > 1:
            df.columns = ["_".join(a).strip("_") for a in df.columns.to_flat_index()]
    if verbose:
        print(df.head())
    return df


def split_sampling(df, args):
    from sklearn.model_selection import train_test_split

    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    stratify_on = args.get("stratify_on", None)
    random_state = args.get("random_state", 123)
    groupby = args.get("groupby", stratify_on)
    unique_key = args.get("unique_key", "id")
    test_size = args.get("test_size", 0.1)
    dev_size = args.get("dev_size", None)
    if verbose:
        log.info(f"Split sampling: {args}")

    train, dev, test = None, None, None
    if stratify_on is None:
        train, test = train_test_split(
            df, test_size=test_size, random_state=random_state
        )
        if dev_size:
            train, dev = train_test_split(
                train, test_size=dev_size, random_state=random_state
            )
    else:
        train, test = train_test_split(
            df, test_size=test_size, random_state=random_state, stratify=df[stratify_on]
        )
        if dev_size:
            train, dev = train_test_split(
                train,
                test_size=dev_size,
                random_state=random_state,
                stratify=train[stratify_on],
            )
    train.reset_index(drop=True, inplace=True)
    test.reset_index(drop=True, inplace=True)
    if dev_size:
        dev.reset_index(drop=True, inplace=True)
    if verbose:
        log.info(f"Total rows: {len(df)}")
        log.info(f"Train: {len(train)}")
        log.info(f"Test: {len(test)}")
        if dev_size:
            log.info(f"Dev: {len(dev)}")

        grp_all = (
            df.groupby(groupby)[unique_key]
            .count()
            .rename("population")
            .transform(lambda x: x / x.sum() * 100)
        )
        grp_train = (
            train.groupby(groupby)[unique_key]
            .count()
            .rename("train")
            .transform(lambda x: x / x.sum() * 100)
        )
        grp_test = (
            test.groupby(groupby)[unique_key]
            .count()
            .rename("test")
            .transform(lambda x: x / x.sum() * 100)
        )
        grp_dists = pd.concat([grp_all, grp_train, grp_test], axis=1)
        if dev_size:
            grp_dev = (
                dev.groupby(groupby)[unique_key]
                .count()
                .rename("dev")
                .transform(lambda x: x / x.sum() * 100)
            )
            grp_dists = pd.concat([grp_dists, grp_dev], axis=1)
        print(grp_dists)

    output_dir = args.get("output_dir", ".")
    train_file = args.get("train_file", None)
    test_file = args.get("test_file", None)
    dev_file = args.get("dev_file", None)
    if train_file and train is not None:
        save_dataframe_(
            train, output_dir=output_dir, output_file=train_file, verbose=verbose
        )
    if test_file and test is not None:
        save_dataframe_(
            test, output_dir=output_dir, output_file=test_file, verbose=verbose
        )
    if dev_file and dev is not None:
        save_dataframe_(
            dev, output_dir=output_dir, output_file=dev_file, verbose=verbose
        )

    return df


def top_values(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    groupby = args.get("groupby", None)
    value_var = args.get("value_var", None)
    value_label = args.get("value_label", None)
    value_name = args.get("value_name", None)
    value_separator = args.get("value_separator", ", ")
    top_n = args.get("top_n", 5)
    columns = args.get("columns", None)
    use_batcher = args.get("use_batcher", True)
    minibatch_size = args.get("minibatch_size", None)

    if verbose:
        log.info(f"Split sampling: {args}")

    def label(row):
        return f"{row[value_label]}[{round(row[value_var]*100,0):.0f}%]"

    if value_label is not None and value_name is not None:
        # df[value_name] = df.apply(label, axis=1)
        df[value_name] = apply(
            label,
            df,
            description="labeling",
            verbose=verbose,
            use_batcher=use_batcher,
            minibatch_size=minibatch_size,
        )

    top_n_grp = (
        df.sort_values(by=value_var, ascending=False)
        .groupby(groupby)
        .head(top_n)
        .reset_index(drop=True)
    )
    top_n_grp = top_n_grp.sort_values(by=groupby)
    top_n_grp = top_n_grp.groupby(groupby, as_index=False).agg(
        {value_name: value_separator.join}
    )
    top_n_grp = top_n_grp[groupby + [value_name]]
    df.drop(columns=value_name, inplace=True)
    df = pd.merge(df, top_n_grp, on=groupby, how="left")
    if columns:
        df = df[columns:]
    if verbose:
        print(df.head())

    return df


def subset(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    random_state = args.get("random_state", 123)
    head_n = args.get("head_n", None)
    tail_n = args.get("tail_n", None)
    sample_n = args.get("sample_n", None)
    sample_frac = args.get("sample_frac", None)
    if not (head_n or tail_n or sample_n or sample_frac):
        log.error("Must specify one of head_n, tail_n, sample_n, sample_frac")
        return df

    if verbose:
        log.info(f"Subsetting: {args}")
    dfs = []
    if head_n:
        dfs.append(df.head(head_n))
    if tail_n:
        dfs.append(df.tail(tail_n))
    if sample_n:
        dfs.append(df.sample(sample_n, random_state=random_state))
    if sample_frac:
        dfs.append(df.sample(frac=sample_frac, random_state=random_state))
    if verbose:
        log.info(f"Total rows: {len(df)}")

    if len(dfs) > 0:
        df = pd.concat(dfs)

    if verbose:
        log.info(f"Subset rows: {len(df)}")

    return df


def sampling(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    random_state = args.get("random_state", 123)
    groupby = args.get("groupby", None)
    sample_size_per_group = args.get("sample_size_per_group", None)
    value_var = args.get("value_var", None)
    columns = args.get("columns", None)

    if verbose:
        log.info(f"Split sampling: {args}")

    if groupby is None:
        if sample_size_per_group < 1:
            df_sample = df.sample(frac=sample_size_per_group, random_state=random_state)
        else:
            df_sample = df.sample(n=sample_size_per_group, random_state=random_state)
    else:
        if sample_size_per_group < 1:
            df_sample = df.groupby(groupby, group_keys=False).apply(
                lambda x: x.sample(
                    frac=sample_size_per_group, random_state=random_state
                )
            )
        else:
            df_sample = df.groupby(groupby, group_keys=False).apply(
                lambda x: x.sample(
                    n=min(len(x), sample_size_per_group), random_state=random_state
                )
            )
    df_sample.reset_index(drop=True, inplace=True)
    if columns:
        df_sample = df_sample[columns:]
    if verbose:
        log.info(f"Total rows: {len(df)}")
        log.info(f"Sample: {len(df_sample)}")
        print(df_sample.head())

    if groupby is not None and verbose:
        grp_all = df.groupby(groupby)[value_var].count().rename("population")
        grp_sample = df_sample.groupby(groupby)[value_var].count().rename("sample")
        grp_dists = pd.concat([grp_all, grp_sample], axis=1)
        print(grp_dists)

    save_dataframe(df_sample, args)

    return df


def aggregate_columns(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    onto_column = args.get("onto", None)
    aggregations = args.get("aggregations", None)
    reset_index = args.get("reset_index", False)
    if onto_column is None and aggregations is None:
        if verbose:
            log.warning("No columns or aggregations are specified")
        return df
    groupby_cloumns = args["groupby"]
    if groupby_cloumns is None:
        if verbose:
            log.warning("No groupby specified")
        return df
    separator = codecs.decode(args["separator"], "unicode_escape")

    num_docs = df.shape[0]
    if verbose:
        log.info(f"Aggregating columns: {args}")
    if aggregations:
        df = df.groupby(groupby_cloumns, as_index=False).agg(aggregations)
    else:
        df[onto_column].fillna("", inplace=True)
        df = df.groupby(groupby_cloumns, as_index=False).agg(
            {onto_column: separator.join}
        )
    if reset_index:
        if verbose:
            log.info(f"Resetting index, nlevels of columns: {df.columns.nlevels}")
        df.reset_index(inplace=True)
        if df.columns.nlevels > 1:
            df.columns = ["_".join(a).strip("_") for a in df.columns.to_flat_index()]
    n_docs = df.shape[0]
    if verbose:
        print(df.tail())
        log.info(f"{num_docs} documents aggregated into {n_docs} documents")
    return df


def explode_splits(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    separator = codecs.decode(args["separator"], "unicode_escape")
    id_key = args.get("id_key", "id")
    split_key = args.get("split_key", "seg_id")
    if verbose:
        log.info(f"Exploding column: {args}")

    num_docs = df.shape[0]
    df[apply_to] = df[apply_to].str.split(separator)
    df = df.explode(apply_to)
    df[split_key] = df.groupby(id_key).cumcount()
    n_docs = df.shape[0]
    if verbose:
        log.info(f"{num_docs} documents exploded into {n_docs} documents")
    return df


def rename_columns(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    new_names = args.get("new_names", None)
    if new_names is None:
        if verbose:
            log.warning("No columns specified")
        return df
    if verbose:
        log.info(f"Renaming columns: {args}")
    if new_names is not None:
        df.rename(columns=new_names, inplace=True)
    if verbose:
        print(df.head())
    return df


def reset_index(df, args):
    verbose = args.get("verbose", False)
    index_column_name = args.get("index_column_name", None)
    if index_column_name is None:
        index_column_name = "index"
    drop_index = args.get("drop_index", False)
    if verbose:
        log.info(f"Resetting index: {args}")
    df = df.reset_index(drop=drop_index)
    if not drop_index and index_column_name != "index":
        df.rename(columns={"index": index_column_name}, inplace=True)
    if verbose:
        print(df.head())
    return df


def normalize(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        log.warning("No columns specified")
        return df
    normalizer = args.get("preprocessor", {}).get("normalizer", None)
    if normalizer is None:
        log.warning("No normalizer specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    log.info("instantiating normalizer")
    normalizer = eKonf.instantiate(normalizer)
    for key in apply_to:
        with elapsed_timer(format_time=True) as elapsed:
            df[key] = apply(
                normalizer.normalize,
                df[key],
                description=f"Normalizing column: {key}",
                verbose=verbose,
            )
            log.info(" >> elapsed time to normalize: {}".format(elapsed()))
    return df


def fillna(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    fill_with = args.get("fill_with", "")
    if verbose:
        log.info(f"Filling missing values: {args}")
    for key in apply_to:
        df[key].fillna(fill_with, inplace=True)
    return df


def predict(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    num_workers = args.get("num_workers", 1)
    use_batcher = args.get("use_batcher", True)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    columns = args["columns"]

    to_predict = args.get("to_predict", None)
    model = args.get("model", None)
    if model is None:
        log.warning("No model specified")
        return df

    if isinstance(apply_to, list):
        apply_to = apply_to[0]

    log.info("instantiating model")

    _target_ = model[eKonf.Keys.TARGET]
    model = eKonf.instantiate(model)

    if "sentiment.analyser" in _target_:
        key = apply_to
        with elapsed_timer(format_time=True) as elapsed:
            predictions = apply(
                model.predict,
                df[key],
                description=f"Predicting [{key}]",
                verbose=verbose,
                use_batcher=use_batcher,
                num_workers=num_workers,
            )
            pred_df = pd.DataFrame(predictions.tolist(), index=predictions.index)
            if columns:
                columns += pred_df.columns.tolist()
                args["columns"] = columns
            df = df.join(pred_df)
            log.info(" >> elapsed time to predict: {}".format(elapsed()))
    else:
        df = model.predict(df, to_predict)
        if columns:
            columns.append(model._to_predict["predicted"])
            args["columns"] = columns

    save_dataframe(df, args)

    return df


def segment(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    use_batcher = args.get("use_batcher", True)
    minibatch_size = args.get("minibatch_size", None)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    segmenter = args.get("preprocessor", {}).get("segmenter", None)
    if segmenter is None:
        log.warning("No segmenter specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    log.info("instantiating segmenter")
    segmenter = eKonf.instantiate(segmenter)
    for key in apply_to:
        with elapsed_timer(format_time=True) as elapsed:
            df[key] = apply(
                segmenter.segment_article,
                df[key],
                description=f"Splitting column: {key}",
                verbose=verbose,
                use_batcher=use_batcher,
                minibatch_size=minibatch_size,
            )
            log.info(" >> elapsed time to segment: {}".format(elapsed()))
    return df


def tokenize(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    use_batcher = args.get("use_batcher", True)
    minibatch_size = args.get("minibatch_size", None)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        log.warning("No columns specified")
        return df
    tokenizer = args.get("preprocessor", {}).get("tokenizer", None)
    if tokenizer is None:
        log.warning("No tokenizer specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    log.info("instantiating tokenizer")
    tokenizer = eKonf.instantiate(tokenizer)
    for key in apply_to:
        with elapsed_timer(format_time=True) as elapsed:
            df[key] = apply(
                tokenizer.tokenize_article,
                df[key],
                description=f"Tokenizing column: {key}",
                verbose=verbose,
                use_batcher=use_batcher,
                minibatch_size=minibatch_size,
            )
            log.info(" >> elapsed time to segment: {}".format(elapsed()))
    return df


def extract_tokens(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    use_batcher = args.get("use_batcher", True)
    nouns_only = args.get("nouns_only", True)
    filter_stopwords_only = args.get("filter_stopwords_only", False)
    minibatch_size = args.get("minibatch_size", None)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        log.warning("No columns specified")
        return df
    tokenizer = args.get("preprocessor", {}).get("tokenizer", None)
    if tokenizer is None:
        log.warning("No tokenizer specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    log.info("instantiating tokenizer")
    tokenizer = eKonf.instantiate(tokenizer)
    if filter_stopwords_only:
        extract_func = tokenizer.filter_article_stopwords
    elif nouns_only:
        extract_func = tokenizer.extract_nouns
    else:
        extract_func = tokenizer.extract_tokens
    log.info(f"extract_func: {extract_func}")

    for key in apply_to:
        with elapsed_timer(format_time=True) as elapsed:
            df[key] = apply(
                extract_func,
                df[key],
                description=f"Extracting column: {key}",
                verbose=verbose,
                use_batcher=use_batcher,
                minibatch_size=minibatch_size,
            )
            log.info(" >> elapsed time to extract tokens: {}".format(elapsed()))
    return df


def chunk(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    use_batcher = args.get("use_batcher", True)
    minibatch_size = args.get("minibatch_size", None)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        log.warning("No columns specified")
        return df
    segmenter = args.get("preprocessor", {}).get("segmenter", None)
    if segmenter is None:
        log.warning("No segmenter specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    log.info("instantiating segmenter")
    segmenter = eKonf.instantiate(segmenter)
    for key in apply_to:
        with elapsed_timer(format_time=True) as elapsed:
            df[key] = apply(
                segmenter.chunk_article,
                df[key],
                description=f"Chunking column: {key}",
                verbose=verbose,
                use_batcher=use_batcher,
                minibatch_size=minibatch_size,
            )
            log.info(" >> elapsed time to segment: {}".format(elapsed()))
    return df


def general_function(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", None)
    method = args.get(eKonf.Keys.METHOD, None)
    if verbose:
        log.info(f"Running dataframe function: {args}")

    with elapsed_timer(format_time=True) as elapsed:
        if apply_to is None:
            df = getattr(df, method[eKonf.Keys.NAME])(**method[eKonf.Keys.PARMS])
        else:
            if isinstance(apply_to, str):
                apply_to = [apply_to]
            for key in apply_to:
                log.info(f"processing column: {key}")
                df[key] = getattr(df[key], method[eKonf.Keys.NAME])(
                    **method[eKonf.Keys.PARMS]
                )

        log.info(" >> elapsed time to replace: {}".format(elapsed()))
        if verbose:
            print(df.head())
    return df


def replace_regex(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    patterns = args.get("patterns", {})
    if patterns is None:
        if verbose:
            log.warning("No patterns specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    if verbose:
        log.info(f"Replacing regex: {args}")
    for key in apply_to:
        with elapsed_timer(format_time=True) as elapsed:
            for pat, repl in patterns.items():
                df[key] = df[key].str.replace(pat, repl, regex=True).str.strip()
            if verbose:
                log.info(" >> elapsed time to replace regex: {}".format(elapsed()))
    return df


def remove_startswith(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    startswith = args.get("startswith", {})
    if startswith is None:
        if verbose:
            log.warning("No startswith text specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    if verbose:
        log.info(f"Remove startswith: {args}")
    for key in apply_to:
        with elapsed_timer(format_time=True) as elapsed:
            for starting_text in startswith:
                log.info(f"Remove text starting with {starting_text} from [{key}]")
                idx = df[key].str.lower().str.startswith(starting_text, na=False)
                start_pos = len(starting_text)
                df.loc[idx, key] = df.loc[idx, key].str[start_pos:].str.strip()
            if verbose:
                log.info(
                    " >> elapsed time to remove startswith: {}\n".format(elapsed())
                )
    return df


def filter_length(df, args, **kwargs):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", "text")
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    min_length = args.get("min_length", None)
    max_length = args.get("max_length", None)
    if min_length is None and max_length is None:
        if verbose:
            log.warning("No length specified")
        return df
    len_func = args[eKonf.Keys.FUNC].get("len_bytes", None)
    len_func = eKonf.instantiate(len_func)
    _check_max_len = partial(check_max_len, max_len=max_length, len_func=len_func)
    _check_min_len = partial(check_min_len, min_len=min_length, len_func=len_func)

    if verbose:
        log.info(f"Filtering by length: {args}")
    for key in apply_to:
        with elapsed_timer(format_time=True) as elapsed:
            if min_length and min_length > 0:
                n_docs = df.shape[0]
                idx = apply(
                    _check_min_len,
                    df[key].astype(str),
                    verbose=verbose,
                    description=f"min length: {min_length}",
                )
                df = df[idx]
                if verbose:
                    log.info(
                        f"{(n_docs-df.shape[0])} of {n_docs} documents removed due to length is less than {min_length}"
                    )
            if max_length and max_length > 0:
                n_docs = df.shape[0]
                idx = apply(
                    _check_max_len,
                    df[key].astype(str),
                    verbose=verbose,
                    description=f"max length: {max_length}",
                )
                df = df[idx]
                if verbose:
                    log.info(
                        f"{(n_docs-df.shape[0])} of {n_docs} documents removed due to length is greater than {max_length}"
                    )
            if verbose:
                log.info(" >> elapsed time to filter length: {}".format(elapsed()))
    return df


def filter_query(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    query = args.get("query", None)
    if query is None:
        if verbose:
            log.warning("No query specified")
        return df
    if isinstance(query, str):
        query = [query]

    if verbose:
        log.info(f"Filtering by qeury: {args}")
    with elapsed_timer(format_time=True) as elapsed:
        for qry in query:
            if verbose:
                log.info(f"Preprocessing query: {qry}")
            n_docs = df.shape[0]
            df = df.query(qry, engine="python")
            if verbose:
                log.info(f"filtered {df.shape[0]} out of {n_docs} documents by {qry}")
        if verbose:
            log.info(" >> elapsed time to filter query: {}".format(elapsed()))
    return df


def drop_duplicates(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", None)
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    if verbose:
        log.info(f"Dropping duplicates: {args}")
    with elapsed_timer(format_time=True) as elapsed:
        num_docs = df.shape[0]
        df = df.drop_duplicates(subset=apply_to)
        n_docs = df.shape[0]
        if verbose:
            log.info(
                f"{n_docs} documents after dropping {(num_docs-n_docs)} duplicates from [{apply_to}]"
            )
            log.info(" >> elapsed time to drop duplicates: {}".format(elapsed()))
    return df


def save_samples(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", "text")
    sample_length_to_print = args.get("sample_length_to_print", 1000)
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    num_samples_to_save = args.get("num_samples_to_save", None)
    smaple_file_prefix = args.get("sample_file_prefix", "sample")
    if verbose:
        log.info(f"Saving samples: {args}")

    sample_separator = "-" * 100 + "\n"
    df_sample = df.sample(num_samples_to_save)[apply_to]
    sample_text = ""
    print_text = ""
    for i, row in df_sample.iterrows():
        for key in apply_to:
            stext = row[key]
            if len(stext) > sample_length_to_print:
                ptext = stext[:sample_length_to_print] + "..."
            else:
                ptext = stext
            sample_text += key + ": \n" + stext + "\n\n"
            print_text += key + ": \n" + ptext + "\n\n"
        sample_text += sample_separator
        print_text += sample_separator
    sample_text = sample_text.strip()
    print_text = print_text.strip()
    sample_file = smaple_file_prefix + ".txt"
    open(sample_file, "w", encoding="utf-8").write(sample_text)

    if verbose:
        print(sample_separator)
        print(print_text)
        log.info(f"Saved {num_samples_to_save} samples to {sample_file}")

    return df


def stdout_samples(df, args):
    from contextlib import redirect_stdout

    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", "text")
    sample_length_to_print = args.get("sample_length_to_print", 1000)
    if apply_to is None:
        if verbose:
            log.warning("No columns specified")
        return df
    if isinstance(apply_to, str):
        apply_to = [apply_to]
    num_samples = args.get("num_samples", None)
    output_dir = args.get("output_dir", None)
    if output_dir is None:
        output_dir = "."
    output_file = args.get("output_file", None)
    head = args.get("head", None)
    tail = args.get("tail", None)
    if head is None:
        head = 5
    if tail is None:
        tail = 5

    if verbose:
        log.info(f"Print samples: {args}")

    sample_separator = "-" * 100 + "\n"
    df_sample = df.sample(num_samples)[apply_to]

    print_text = ""
    for i, row in df_sample.iterrows():
        for key in apply_to:
            ptext = row[key]
            print_text += key + ": \n" + ptext + "\n\n"
        print_text += sample_separator
    print_text = print_text.strip()

    if output_file is not None:
        output_file = os.path.join(output_dir, output_file)
        with open(output_file, "w", encoding="utf-8") as f:
            with redirect_stdout(f):
                print(print_text)
                print(sample_separator)
                print(df.head(head))
                print(sample_separator)
                print(df.tail(tail))
    print(print_text)
    print(sample_separator)
    print(df.head(head))
    print(sample_separator)
    print(df.tail(tail))

    if verbose:
        log.info(f"Saved {num_samples} samples to {output_file}")

    return df


def save_as_text(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    apply_to = args.get("apply_to", "text")
    corpus_name = args.get("corpus_name", eKonf.Keys.CORPUS)
    output_dir = args.get("output_dir", ".")
    output_file = args.get("output_file", None)
    doc_separator = args.get("doc_separator", "\n\n")
    if isinstance(apply_to, list):
        apply_to = apply_to[0]

    os.makedirs(os.path.abspath(output_dir), exist_ok=True)
    if output_file:
        filename = output_file
    else:
        filename = f"{corpus_name}.txt"
    output_file_path = f"{output_dir}/{filename}"

    with elapsed_timer(format_time=True) as elapsed:
        n_loaded = len(df.index)
        doc_separator = str(doc_separator).encode("utf-8").decode("unicode_escape")
        with open(output_file_path, "w") as fo:
            fo.write(doc_separator.join(df[apply_to].dropna().tolist()))
        log.info(f"Corpus is exported to {output_file_path}")
        n_sampled = len(df.index)
        status = [[" x ", corpus_name, n_loaded, n_sampled, elapsed(), filename]]

    if verbose:
        print_status(status)
    return df


def split_dataframe(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    num_splits = args.get("num_splits", 1)
    if num_splits <= 1:
        return df
    if verbose:
        log.info(f"Splitting dataframe into {num_splits} splits")
    return np.array_split(df, num_splits)


def concat_dataframes(dataframes, args):
    from ekorpkit.io.file import concat_dataframes as concat_dataframes_

    args = eKonf.to_dict(args)

    return concat_dataframes_(data=dataframes, **args)


def merge_dataframe(df=None, args=None):
    if args is None:
        raise ValueError("args must be specified")
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    how = args.get("how", "inner")
    merge_on = args.get("merge_on", None)
    left_on = args.get("left_on", None)
    right_on = args.get("right_on", None)
    if merge_on is None and (right_on is None or left_on is None):
        raise ValueError("merge_on or (left_on and right_on) must be specified")
    if isinstance(merge_on, str):
        merge_on = [merge_on]
    if isinstance(left_on, str):
        left_on = [left_on]
    if isinstance(right_on, str):
        right_on = [right_on]

    df_to_merge = load_dataframe(None, args)
    if merge_on:
        df = df.merge(df_to_merge, how=how, on=merge_on)
    else:
        df = df.merge(df_to_merge, how=how, left_on=left_on, right_on=right_on)
    if verbose:
        print(df.tail())
    return df


def save_metadata(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    filepath = args.get("filepath", None)
    filetype = args.get("filetype", None)
    column_info = args.get("column_info", None)
    split_name = args.get("split_name", None)

    if verbose:
        log.info(f"Saving metadata: {args}")

    meta_info = column_info.get("meta", None)
    meta_columns = []
    if isinstance(meta_info, dict):
        meta_columns = list(meta_info.keys())
        if eKonf.Keys.SPLIT in meta_columns and eKonf.Keys.SPLIT not in df.columns:
            df[eKonf.Keys.SPLIT] = split_name
        df_meta = df[meta_columns]
        save_dataframe_(df_meta, filepath, filetype, verbose)

    data_info = column_info.get("data", None)
    if isinstance(data_info, dict):
        data_keys = list(data_info.keys())
        # if eKonf.Keys.SPLIT in data_columns and eKonf.Keys.SPLIT not in df.columns:
        #     df[eKonf.Keys.SPLIT] = split_name
        data_columns = [
            col for col in df.columns if col not in meta_columns or col in data_keys
        ]
        df = df[data_columns]

    return df


def save_dataframe(df, args):
    args = eKonf.to_dict(args)

    if df is None:
        log.warning("Dataframe is None")
        return df

    save_dataframe_(df, **args)
    return df


def load_dataframe(df=None, args=None):
    from ekorpkit.io.file import get_filepaths, load_dataframe as _load_dataframe

    if args is None:
        raise ValueError("args must be specified")

    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    concatenate = args.get("concatenate", True)
    filepath = args.get("filepath", None)
    data_dir = args.get("data_dir", None)
    data_file = args.get("data_file", None)
    dtype = args.get("dtype", None)
    if isinstance(dtype, list):
        dtype = {k: "str" for k in dtype}
    parse_dates = args.get("parse_dates", False)
    columns = args.get("columns", None)

    if filepath:
        filepaths = get_filepaths(filepath)
    else:
        filepaths = get_filepaths(data_file, data_dir)
    log.info(f"Loading {len(filepaths)} dataframes from {filepaths}")
    if len(filepaths) == 1:
        df = _load_dataframe(
            filepaths[0],
            verbose=verbose,
            dtype=dtype,
            parse_dates=parse_dates,
            columns=columns,
        )
        return df
    else:
        if concatenate:
            df = pd.concat(
                [
                    _load_dataframe(
                        f,
                        verbose=verbose,
                        dtype=dtype,
                        parse_dates=parse_dates,
                        columns=columns,
                    )
                    for f in filepaths
                ]
            )
            return df
        else:
            df_dict = {}

            for f in filepaths:
                df = _load_dataframe(
                    f,
                    verbose=verbose,
                    dtype=dtype,
                    parse_dates=parse_dates,
                    columns=columns,
                )
                df_name = os.path.basename(f)
                df_dict[df_name] = df
            return df_dict


def summary_stats(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    output_dir = args.get("output_dir", ".")
    output_file = args.get("output_file", None)
    stat_args = args.get("info", {}).get("stats", None)
    if stat_args is None:
        log.warning("No stats specified")
        return df
    if output_file is None:
        output_file = "summary_stats.csv"
    os.makedirs(os.path.abspath(output_dir), exist_ok=True)
    info_path = f"{output_dir}/{output_file}"
    stat_fn = eKonf.instantiate(stat_args)
    stats = stat_fn(df)
    eKonf.save(stats, f=info_path)
    log.info(f"Saving summary stats: {info_path}")
    if verbose:
        eKonf.pprint(stats)

    return df


def save_as_json(df, args):
    args = eKonf.to_dict(args)
    verbose = args.get("verbose", False)
    corpus_name = args.get("corpus_name", eKonf.Keys.CORPUS)
    output_dir = args.get("output_dir", ".")
    output_file = args.get("output_file", None)
    force_ascii = args.get("force_ascii", False)

    os.makedirs(os.path.abspath(output_dir), exist_ok=True)
    if output_file:
        filename = output_file
    else:
        filename = f"{corpus_name}.json"
    output_file_path = f"{output_dir}/{filename}"

    df.to_json(output_file_path, orient="records", lines=True, force_ascii=force_ascii)
    log.info(f"Corpus is exported to {output_file_path}")
    return df


def pipeline(data=None, **cfg):
    args = eKonf.to_dict(cfg)
    verbose = args.get("verbose", False)

    if eKonf.is_instantiatable(data):
        data = eKonf.instantiate(data)

    df = data.data
    if df is None:
        raise ValueError("No dataframe to process")

    process_pipeline = args.get(eKonf.Keys.PIPELINE, [])
    if process_pipeline is None:
        process_pipeline = []

    if len(process_pipeline) > 0:
        df = apply_pipeline(df, process_pipeline, args)
    else:
        log.warning("No pipeline specified")

    return df
