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

# %% auto 0
__all__ = ['chainfn_input', 'chainfn_output', 'chain_fn_args', 'chain_fn', 'transform_chain', 'ray_chain_fn', 'ray_chain',
           'noop_chain', 'docs_to_json_chain', 'add_str_to_docs_chain']

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

# %% ../nbs/01_chains.ipynb 4
def chainfn_input(data, tfm, tfm_kwargs={}, input_variables=["df"]):
    try:
        for k in input_variables:
            tfm_kwargs[k] = data.get(k, None)
        fn_args = inspect.signature(tfm).parameters
        pos_args = [k for k in fn_args if fn_args[k].default == inspect._empty]
        for pk, ik in zip(pos_args, input_variables):
            if pk != "kwargs":
                tfm_kwargs[pk] = tfm_kwargs.pop(ik)
        if "kwargs" not in fn_args:
            tfm_kwargs = {k: v for k, v in tfm_kwargs.items() if k in fn_args.keys()}
    except Exception as e:
        msg.fail(f"Error in chainfn_input: {e}", spaced=True)
    return tfm_kwargs


# def chain_fn(data, tfm, tfm_kwargs={}, input_variables=["df"], output_variables=["df"]):
#     tfm_kwargs = chainfn_input(data, tfm, tfm_kwargs, input_variables)
#     # print(f'\n\nTFM: {tfm}\n\n')
#     # print(f'\n\nTFM KWARGS: {tfm_kwargs.keys()}\n\n')
#     fn_res = tfm(**tfm_kwargs)
#     return chainfn_output(fn_res, output_variables)


def chainfn_output(fn_res, output_variables=["df"]):
    if not list_or_tuple(fn_res):
        fn_res = [fn_res]
    if len(fn_res) > len(output_variables):
        fn_res = [fn_res]
    return {k: r for k, r in zip(output_variables, fn_res)}


def chain_fn_args(data, tfm, tfm_kwargs={}, vars_kwargs_mapping={}):
    try:
        tfm_kwargs = {**data, **tfm_kwargs}
        for k, v in vars_kwargs_mapping.items():
            tfm_kwargs[v] = tfm_kwargs.pop(k)
        fn_args = inspect.signature(tfm).parameters
        if "kwargs" not in fn_args:
            tfm_kwargs = {k: v for k, v in tfm_kwargs.items() if k in fn_args.keys()}
    except Exception as e:
        msg.fail(f"Error in chain_fn_args: {e}", spaced=True)
    return tfm_kwargs


def chain_fn(data, tfm, tfm_kwargs={}, vars_kwargs_mapping={}, output_variables=["df"]):
    tfm_kwargs = chain_fn_args(data, tfm, tfm_kwargs, vars_kwargs_mapping)
    # print(f"\n\nTFM: {tfm}\n\n")
    # print(f"\n\nTFM KWARGS: {tfm_kwargs.keys()}\n\n")
    fn_res = tfm(**tfm_kwargs)
    fn_res = {**data, **chainfn_output(fn_res, output_variables)}
    # print(f"\n\nFN RES: {fn_res}\n\n")
    return fn_res


def transform_chain(
    transform,
    transform_kwargs={},
    vars_kwargs_mapping={},  # For mapping `input_variables` to `transform_kwargs`
    input_variables=["df"],
    output_variables=["df"],
    verbose=False,
):
    "Return a `transform` function wrapped in a `TransformChain`."
    return TransformChain(
        input_variables=input_variables,
        output_variables=output_variables,
        transform=partial(
            chain_fn,
            tfm=transform,
            tfm_kwargs=transform_kwargs,
            vars_kwargs_mapping=vars_kwargs_mapping,
            output_variables=output_variables,
        ),
        verbose=verbose,
    )


def ray_chain_fn(data, chain, block_size=1500, num_cpus=8, num_gpus=1):
    """
    This function is used to run a given chain function on a Ray cluster.
    It splits the input data into blocks, and runs the chain function on each block in parallel.
    The results are then combined and returned.
    """

    try:
        # If the block size is None, run the chain function on the entire dataset.
        if block_size is None:
            return chain(data)

        # Convert the data into a dictionary of lists.
        ray_data = {k: v for k, v in data.items()}
        for k, v in ray_data.items():
            if not is_list(v):
                ray_data[k] = [v]

        # Get the maximum length of any list in the data dictionary.
        max_len = max([len(v) for v in ray_data.values()])

        # If the maximum length is 1 or less than the block size, run the chain function on the entire dataset.
        if max_len == 1 or max_len <= block_size:
            return chain(data)

        # If the maximum length is greater than the block size, split the data into blocks.
        for k, v in ray_data.items():
            if len(v) < max_len:
                ray_data[k] = v * max_len

        # Convert the data dictionary into a Pandas DataFrame.
        df = pd.DataFrame(ray_data)

        # Initialize Ray.
        ray.init(ignore_reinit_error=True)

        # Calculate the number of blocks to split the data into.
        num_blocks = int(np.ceil(len(df) / block_size))
        msg.info(f"Running chain on {num_blocks} blocks.", spaced=True)

        # Calculate the number of CPUs and GPUs to use per block.
        num_cpus = min(ray.available_resources()["CPU"], num_cpus)
        num_cpus /= num_blocks
        if num_gpus is not None:
            num_gpus = min(ray.available_resources().get("GPU", 0), num_gpus)
            if num_gpus > 0:
                num_gpus /= num_blocks
                num_cpus = None
            else:
                num_gpus = None
        # Convert the Pandas DataFrame into a Ray dataset and run the chain function on each block.
        ds = rd.from_pandas(df).repartition(num_blocks)
        res = (
            ds.map_batches(
                lambda x: chain(x),
                batch_size=block_size,
                num_cpus=num_cpus,
                num_gpus=num_gpus,
            )
            .to_pandas()
            .to_dict(orient="list")
        )

        return res
    except Exception as e:
        msg.fail(f"Error in ray_chain_fn: {e}", spaced=True)
        return chain(data)


def ray_chain(chain, block_size=1500, num_cpus=8, num_gpus=1, verbose=False):
    """
    This function is used to run a given chain function on a Ray cluster if necessary.
    """
    tfm = partial(
        ray_chain_fn,
        chain=chain,
        block_size=block_size,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
    )
    input_variables = chain.input_keys
    output_variables = chain.output_keys
    return TransformChain(
        input_variables=input_variables,
        output_variables=output_variables,
        transform=tfm,
        verbose=verbose,
    )


def noop_chain():
    return transform_chain(noop)


def docs_to_json_chain(
    json_folder,
    data={},
    data_key="data",
    with_metadata=True,
    with_content=True,
    indent=None,
    input_variables=["docs"],
    output_variables=["docs"],
    verbose=False,
):
    return transform_chain(
        docs_to_json,
        transform_kwargs=dict(
            json_folder=json_folder,
            data=data,
            data_key=data_key,
            with_metadata=with_metadata,
            with_content=with_content,
            indent=indent,
        ),
        vars_kwargs_mapping={input_variables[0]: "docs"},
        input_variables=input_variables,
        output_variables=output_variables,
        verbose=verbose,
    )
    
def add_str_to_docs_chain(
    str,
    input_variables=["docs"],
    output_variables=["cat_docs"],
    verbose=False,
):
    "Chain that adds a string to the docs"
    return transform_chain(
        add_str_to_docs,
        transform_kwargs=dict(str=str),
        input_variables=input_variables,
        output_variables=output_variables,
        vars_kwargs_mapping={input_variables[0]: "docs"},
        verbose=verbose,
    )
