"""Basic math routine"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/CLI/math.ipynb.

# %% auto 0
__all__ = ['math']

# %% ../../nbs/CLI/math.ipynb 4
import logging
import zarr
import numpy as np
import numexpr as ne
import time

import dask
from dask import array as da
from dask.distributed import Client, LocalCluster, progress

from .logging import mc_logger

# %% ../../nbs/CLI/math.ipynb 5
def _math(operation=str,
          **data):
    return ne.evaluate(operation,data)

# %% ../../nbs/CLI/math.ipynb 6
@mc_logger
def math(output:str, # path to output
         operation:str, # operation
         **data):
    '''
    Basic math manipulation. Only elementwise operations are supported. Only one output is supported.
    '''
    output_path = output
    logger = logging.getLogger(__name__)
    for name, path in data.items():
        path_zarr = zarr.open(path,mode='r'); logger.zarr_info(name,path_zarr)
    logger.info('starting dask local cluster.')
    with LocalCluster(processes=False,threads_per_worker=2) as cluster, Client(cluster) as client:
        logger.info('dask local cluster started.')
        logger.dask_cluster_info(cluster)
        names = []; darrs = []
        
        for name, path in data.items():
            names.append(name)
            darr = da.from_zarr(path,inline_array=True); logger.darr_info(name,darr)
            darrs.append(darr)
        darr0 = darrs[0]
        darrs = [darr.to_delayed() for darr in darrs]

        output_delayed = np.empty_like(darrs[0],dtype=object)
        math_delayed = dask.delayed(_math,pure=True,nout=1)
        with np.nditer(darrs[0],flags=['multi_index','refs_ok'], op_flags=['readwrite']) as it:
            for block in it:
                idx = it.multi_index
                math_kw = {}
                for name, darr in zip(names, darrs):
                    math_kw[name] = darr[idx]
                output_delayed[idx] = math_delayed(operation,**math_kw)
                output_delayed[idx] = da.from_delayed(output_delayed[idx],shape=darr0.blocks[idx].shape,meta=np.array(()))
        output = da.block(output_delayed.tolist())
        _output = output.to_zarr(output_path,overwrite=True,compute=False)
        logger.info('computing graph setted. doing all the computing.')
        futures = client.persist(_output)
        progress(futures,notebook=False)
        da.compute(futures)
        time.sleep(0.1)
        logger.info('computing finished.')
    logger.info('dask cluster closed.')
