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

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

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

# %% ../nbs/00_chains.ipynb 3
def chainfn_input(data, tfm, tfm_kwargs={}, input_variables=["df"]):
    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()}
    return tfm_kwargs

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

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 transform_chain(
    transform,
    transform_kwargs={},
    input_variables=["df"],
    output_variables=["df"],
):
    return TransformChain(
        input_variables=input_variables,
        output_variables=output_variables,
        transform=partial(
            chain_fn,
            tfm=transform,
            tfm_kwargs=transform_kwargs,
            input_variables=input_variables,
            output_variables=output_variables,
        ),
    )


def ray_chain_fn(data, chain, block_size=1500, num_cpus=8, num_gpus=1):
    ray_data = data[chain.input_keys[0]]
    # if not is_df(df):
    # res = chain.run(df)
    if path_or_str(ray_data) or not is_iter(ray_data):
        res = chain.run(ray_data)
        # print(f'\n\nRES: {res}\n\n')
        return {chain.output_keys[0]: res}
    if block_size is None or len(ray_data) <= block_size:
        res = chain.run(ray_data)
        # print(f'\n\nRES: {res}\n\n')
        return {chain.output_keys[0]: res}
    ray.init(ignore_reinit_error=True)
    num_blocks = int(np.ceil(len(ray_data) / block_size))
    msg.info(f"Running chain on {num_blocks} blocks.", spaced=True)
    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()["GPU"], num_gpus)
        num_gpus /= num_blocks
        num_cpus = None
    if not is_df(ray_data):
        ds = rd.from_items(ray_data).repartition(num_blocks)
    else:
        ds = rd.from_pandas(ray_data).repartition(num_blocks)
    res = ds.map_batches(
        lambda x: chain.run(x),
        batch_size=block_size,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
        batch_format="pandas",
    ).to_pandas()
    return {chain.output_keys[0]: res}


def ray_chain(chain, block_size=1500, num_cpus=8, num_gpus=1):
    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,
    )


def noop_chain():
    return transform_chain(noop)
