import shutil, ctypes, time, os, zarr, pprint, psutil, dask, copy
import numpy as np, os, glob, tempfile
from eubi_bridge.base.data_manager import BatchManager
from multiprocessing.pool import ThreadPool

from dask import array as da
from distributed import LocalCluster, Client
from dask_jobqueue import SLURMCluster
from pathlib import Path
from typing import Union

# from eubi_bridge.ngff.multiscales import Pyramid
# from eubi_bridge.ngff import defaults
from eubi_bridge.ebridge_base import BridgeBase, downscale
from eubi_bridge.utils.convenience import sensitive_glob, take_filepaths, is_zarr_group
from eubi_bridge.utils.metadata_utils import print_printable, get_printables

import logging, warnings

logging.getLogger('distributed.diskutils').setLevel(logging.CRITICAL)

def verify_filepaths_for_cluster(
                             filepaths
                             ):
    print(f"Verifying file extensions for distributed setup.")
    formats = ['lif', 'czi', 'lsm',
               'nd2', 'ome.tiff', 'ome.tif',
               'tiff', 'tif', 'zarr',
               'png', 'jpg', 'jpeg']
    for filepath in filepaths:
        verified = any(list(map(lambda path, ext: path.endswith(ext), [filepath] * len(formats), formats)))
        if not verified:
            root, ext = os.path.splitext(filepath)
            warnings.warn(f"Distributed execution is not supported for the {ext} format")
            warnings.warn(f"Falling back on multithreading.")
            break
    if verified:
        print(f"File extensions were verified for distributed setup.")
    return verified

class EuBIBridge:
    """
    EuBIBridge is a conversion tool for bioimage datasets, allowing for both unary and aggregative conversion of image
    data collections to OME-Zarr format.

    Attributes:
        config_gr (zarr.Group): Configuration settings stored in a Zarr group.
        config (dict): Dictionary representation of configuration settings for cluster, conversion, and downscaling.
        dask_config (dict): Dictionary representation of configuration settings for dask.distributed.
        root_defaults (dict): Installation defaults of configuration settings for cluster, conversion, and downscaling.
        root_dask_defaults (dict): Installation defaults of configuration settings for dask.distributed.
    """
    def __init__(self,
                 configpath = f"{os.path.expanduser('~')}/.eubi_bridge",
               ):
        """
        Initializes the EuBIBridge class and loads or sets up default configuration.

        Args:
            configpath (str, optional): Path to store configuration settings. Defaults to the home directory.
        """

        root_dask_defaults = {'distributed.adaptive.interval': '1s', 'distributed.adaptive.maximum': '.inf',
             'distributed.adaptive.minimum': 0, 'distributed.adaptive.target-duration': '5s',
             'distributed.adaptive.wait-count': 3, 'distributed.admin.event-loop': 'tornado',
             'distributed.admin.large-graph-warning-threshold': '10MB',
             'distributed.admin.log-format': '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
             'distributed.admin.log-length': 10000, 'distributed.admin.low-level-log-length': 1000,
             'distributed.admin.max-error-length': 10000, 'distributed.admin.pdb-on-err': False,
             'distributed.admin.system-monitor.disk': True, 'distributed.admin.system-monitor.gil.enabled': True,
             'distributed.admin.system-monitor.gil.interval': '1ms', 'distributed.admin.system-monitor.host-cpu': False,
             'distributed.admin.system-monitor.interval': '500ms', 'distributed.admin.system-monitor.log-length': 7200,
             'distributed.admin.tick.cycle': '1s', 'distributed.admin.tick.interval': '20ms',
             'distributed.admin.tick.limit': '3s', 'distributed.client.heartbeat': '5s',
             'distributed.client.preload': [], 'distributed.client.preload-argv': [],
             'distributed.client.scheduler-info-interval': '2s', 'distributed.client.security-loader': None,
             'distributed.comm.compression': False, 'distributed.comm.default-scheme': 'tcp',
             'distributed.comm.offload': '10MiB', 'distributed.comm.require-encryption': None,
             'distributed.comm.retry.count': 0, 'distributed.comm.retry.delay.max': '20s',
             'distributed.comm.retry.delay.min': '1s', 'distributed.comm.shard': '64MiB',
             'distributed.comm.socket-backlog': 2048, 'distributed.comm.timeouts.connect': '30s',
             'distributed.comm.timeouts.tcp': '30s', 'distributed.comm.tls.ca-file': None,
             'distributed.comm.tls.ciphers': None, 'distributed.comm.tls.client.cert': None,
             'distributed.comm.tls.client.key': None, 'distributed.comm.tls.max-version': None,
             'distributed.comm.tls.min-version': 1.2, 'distributed.comm.tls.scheduler.cert': None,
             'distributed.comm.tls.scheduler.key': None, 'distributed.comm.tls.worker.cert': None,
             'distributed.comm.tls.worker.key': None, 'distributed.comm.ucx.create-cuda-context': None,
             'distributed.comm.ucx.cuda-copy': None, 'distributed.comm.ucx.environment': {},
             'distributed.comm.ucx.infiniband': None, 'distributed.comm.ucx.nvlink': None,
             'distributed.comm.ucx.rdmacm': None, 'distributed.comm.ucx.tcp': None,
             'distributed.comm.websockets.shard': '8MiB', 'distributed.comm.zstd.level': 3,
             'distributed.comm.zstd.threads': 0, 'distributed.dashboard.export-tool': False,
             'distributed.dashboard.graph-max-items': 5000,
             'distributed.dashboard.link': '{scheme}://{host}:{port}/status',
             'distributed.dashboard.prometheus.namespace': 'dask', 'distributed.deploy.cluster-repr-interval': '500ms',
             'distributed.deploy.lost-worker-timeout': '15s',
             'distributed.diagnostics.computations.ignore-files': [r'runpy\.py', 'pytest', r'py\.test',
                                                                   r'pytest-script\.py', '_pytest', 'pycharm',
                                                                   'vscode_pytest', r'get_output_via_markers\.py'],
             'distributed.diagnostics.computations.ignore-modules': ['asyncio', 'functools', 'threading', 'datashader',
                                                                     'dask', 'debugpy', 'distributed', 'coiled', 'cudf',
                                                                     'cuml', 'matplotlib', 'pluggy', 'prefect',
                                                                     'rechunker', 'xarray', 'xgboost', 'xdist',
                                                                     '__channelexec__', 'execnet'],
             'distributed.diagnostics.computations.max-history': 100, 'distributed.diagnostics.computations.nframes': 0,
             'distributed.diagnostics.cudf': False, 'distributed.diagnostics.erred-tasks.max-history': 100,
             'distributed.diagnostics.nvml': True, 'distributed.nanny.environ': {},
             'distributed.nanny.pre-spawn-environ.MALLOC_TRIM_THRESHOLD_': 65536,
             'distributed.nanny.pre-spawn-environ.MKL_NUM_THREADS': 1,
             'distributed.nanny.pre-spawn-environ.OMP_NUM_THREADS': 1,
             'distributed.nanny.pre-spawn-environ.OPENBLAS_NUM_THREADS': 1, 'distributed.nanny.preload': [],
             'distributed.nanny.preload-argv': [], 'distributed.p2p.comm.buffer': '1 GiB',
             'distributed.p2p.comm.concurrency': 10, 'distributed.p2p.comm.message-bytes-limit': '2 MiB',
             'distributed.p2p.comm.retry.count': 10, 'distributed.p2p.comm.retry.delay.max': '30s',
             'distributed.p2p.comm.retry.delay.min': '1s', 'distributed.p2p.storage.buffer': '100 MiB',
             'distributed.p2p.storage.disk': True, 'distributed.p2p.threads': None, 'distributed.rmm.pool-size': None,
             'distributed.scheduler.active-memory-manager.interval': '2s',
             'distributed.scheduler.active-memory-manager.measure': 'optimistic',
             'distributed.scheduler.active-memory-manager.policies': [
                 {'class': 'distributed.active_memory_manager.ReduceReplicas'}],
             'distributed.scheduler.active-memory-manager.start': True, 'distributed.scheduler.allowed-failures': 3,
             'distributed.scheduler.allowed-imports': ['dask', 'distributed'],
             'distributed.scheduler.bandwidth': '100000000', 'distributed.scheduler.blocked-handlers': [],
             'distributed.scheduler.contact-address': None,
             'distributed.scheduler.dashboard.bokeh-application.allow_websocket_origin': ['*'],
             'distributed.scheduler.dashboard.bokeh-application.check_unused_sessions_milliseconds': 500,
             'distributed.scheduler.dashboard.bokeh-application.keep_alive_milliseconds': 500,
             'distributed.scheduler.dashboard.status.task-stream-length': 1000,
             'distributed.scheduler.dashboard.tasks.task-stream-length': 100000,
             'distributed.scheduler.dashboard.tls.ca-file': None, 'distributed.scheduler.dashboard.tls.cert': None,
             'distributed.scheduler.dashboard.tls.key': None, 'distributed.scheduler.default-data-size': '1kiB',
             'distributed.scheduler.default-task-durations.rechunk-split': '1us',
             'distributed.scheduler.default-task-durations.split-shuffle': '1us',
             'distributed.scheduler.default-task-durations.split-stage': '1us',
             'distributed.scheduler.default-task-durations.split-taskshuffle': '1us',
             'distributed.scheduler.events-cleanup-delay': '1h',
             'distributed.scheduler.http.routes': ['distributed.http.scheduler.prometheus',
                                                   'distributed.http.scheduler.info', 'distributed.http.scheduler.json',
                                                   'distributed.http.health', 'distributed.http.proxy',
                                                   'distributed.http.statics'],
             'distributed.scheduler.idle-timeout': None, 'distributed.scheduler.locks.lease-timeout': '30s',
             'distributed.scheduler.locks.lease-validation-interval': '10s',
             'distributed.scheduler.no-workers-timeout': None, 'distributed.scheduler.preload': [],
             'distributed.scheduler.preload-argv': [], 'distributed.scheduler.rootish-taskgroup': 5,
             'distributed.scheduler.rootish-taskgroup-dependencies': 5,
             'distributed.scheduler.unknown-task-duration': '500ms', 'distributed.scheduler.validate': False,
             'distributed.scheduler.work-stealing': True, 'distributed.scheduler.work-stealing-interval': '1s',
             'distributed.scheduler.worker-saturation': 1.1, 'distributed.scheduler.worker-ttl': '5 minutes',
             'distributed.version': 2, 'distributed.worker.blocked-handlers': [],
             'distributed.worker.connections.incoming': 10, 'distributed.worker.connections.outgoing': 50,
             'distributed.worker.daemon': True,
             'distributed.worker.http.routes': ['distributed.http.worker.prometheus', 'distributed.http.health',
                                                'distributed.http.statics'],
             'distributed.worker.lifetime.duration': None, 'distributed.worker.lifetime.restart': False,
             'distributed.worker.lifetime.stagger': '0 seconds', 'distributed.worker.memory.max-spill': False,
             'distributed.worker.memory.monitor-interval': '100ms', 'distributed.worker.memory.pause': 0.8,
             'distributed.worker.memory.rebalance.measure': 'optimistic',
             'distributed.worker.memory.rebalance.recipient-max': 0.6,
             'distributed.worker.memory.rebalance.sender-min': 0.3,
             'distributed.worker.memory.rebalance.sender-recipient-gap': 0.1,
             'distributed.worker.memory.recent-to-old-time': '30s', 'distributed.worker.memory.spill': 0.7,
             'distributed.worker.memory.spill-compression': 'auto', 'distributed.worker.memory.target': 0.6,
             'distributed.worker.memory.terminate': 0.95, 'distributed.worker.memory.transfer': 0.1,
             'distributed.worker.multiprocessing-method': 'spawn', 'distributed.worker.preload': [],
             'distributed.worker.preload-argv': [], 'distributed.worker.profile.cycle': '1000ms',
             'distributed.worker.profile.enabled': True, 'distributed.worker.profile.interval': '10ms',
             'distributed.worker.profile.low-level': False, 'distributed.worker.resources': {},
             'distributed.worker.transfer.message-bytes-limit': '50MB', 'distributed.worker.use-file-locking': True,
             'distributed.worker.validate': False
        }

        defaults = dict(
            cluster = dict(
                n_jobs=4,
                threads_per_worker=1,
                memory_limit='auto',
                temp_dir='auto',
                no_worker_restart=False,
                verbose=False,
                no_distributed=False,
                on_slurm = False,
            ),
            readers = dict(
                as_mosaic=False,
                view_index=0,
                phase_index=0,
                illumination_index=0,
                scene_index=0,
                rotation_index=0,
                mosaic_tile_index=0,
                sample_index=0
            ),                
            conversion = dict(
                zarr_format = 2,
                time_chunk = 1,
                channel_chunk = 1,
                z_chunk = 96,
                y_chunk=96,
                x_chunk=96,
                time_shard_coef = 1,
                channel_shard_coef = 1,
                z_shard_coef = 3,
                y_shard_coef = 3,
                x_shard_coef = 3,
                time_range = None,
                channel_range = None,
                z_range = None,
                y_range = None,
                x_range = None,
                dimension_order='tczyx',
                compressor = 'blosc',
                compressor_params = {},
                overwrite = False,
                use_tensorstore=False,
                use_gpu=False,
                rechunk_method='tasks',
                trim_memory=False,
                metadata_reader = 'bfio',
                save_omexml = True,
                squeeze = False
            ),
            downscale = dict(
                time_scale_factor = 1,
                channel_scale_factor = 1,
                z_scale_factor = 2,
                y_scale_factor = 2,
                x_scale_factor = 2,
                n_layers=3,
                downscale_method='simple',
            )
        )

        self.root_defaults = defaults
        self.root_dask_defaults = root_dask_defaults
        config_gr = zarr.open_group(configpath, mode = 'a')
        config = config_gr.attrs
        for key in defaults.keys():
            if key not in config.keys():
                config[key] = {}
                for subkey in defaults[key].keys():
                    if subkey not in config[key].keys():
                        config[key][subkey] = defaults[key][subkey]
            config_gr.attrs[key] = config[key]
        self.config = dict(config_gr.attrs)
        ###
        if not 'dask_config' in config_gr.keys():
            config_gr.create_group('dask_config')
        dask_config = config_gr['dask_config'].attrs
        for key in root_dask_defaults.keys():
            if key not in dask_config.keys():
                dask_config[key] = root_dask_defaults[key]
        config_gr['dask_config'].attrs.update(dict(dask_config))
        self.dask_config = dict(config_gr['dask_config'].attrs)
        self.config_gr = config_gr
        ###
        self._dask_temp_dir = None
        self.client = None

    def reset_config(self):
        """
        Resets the cluster, conversion and downscale parameters to the installation defaults.
        """
        self.config_gr.attrs.update(self.root_defaults)
        self.config = dict(self.config_gr.attrs)

    def reset_dask_config(self):
        """
        Resets the dask configuration parameters to the installation defaults.
        """
        self.config_gr['dask_config'].attrs.update(self.root_dask_defaults)
        self.dask_config = dict(self.config_gr['dask_config'].attrs)

    def show_config(self):
        """
        Displays the current cluster, conversion, and downscale parameters.
        """
        pprint.pprint(self.config)

    def show_dask_config(self):
        """
        Displays the current dask.distributed parameters.
        """
        pprint.pprint(self.dask_config)

    def show_root_defaults(self):
        """
        Displays the installation defaults for cluster, conversion, and downscale parameters.
        """
        pprint.pprint(self.root_defaults)

    def show_root_dask_defaults(self):
        """
        Displays the installation defaults for dask.distributed.
        """
        pprint.pprint(self.root_dask_defaults)

    def _collect_params(self, param_type, **kwargs):
        """
        Gathers parameters from the configuration, allowing for overrides.

        Args:
            param_type (str): The type of parameters to collect (e.g., 'cluster', 'conversion', 'downscale').
            **kwargs: Parameter values that may override defaults.

        Returns:
            dict: Collected parameters.
        """
        params = {}
        for key in self.config[param_type].keys():
            if key in kwargs.keys():
                params[key] = kwargs[key]
            else:
                params[key] = self.config[param_type][key]
        return params

    def configure_cluster(self,
                          memory_limit: str = 'default',
                          n_jobs: int = 'default',
                          no_worker_restart: bool = 'default',
                          on_slurm: bool = 'default',
                          temp_dir: str = 'default',
                          threads_per_worker: int = 'default',
                          no_distributed: bool = 'default',
                          verbose: bool = 'default'
                         ):
        """
        Updates cluster configuration settings. To update the current default value for a parameter, provide that parameter with a value other than 'default'.

        The following parameters can be configured:
            - memory_limit (str, optional): Memory limit per worker.
            - n_jobs (int, optional): Number of parallel jobs.
            - no_worker_restart (bool, optional): Whether to prevent worker restarts.
            - on_slurm (bool, optional): Whether running on a SLURM cluster.
            - temp_dir (str, optional): Temporary directory for Dask workers.
            - threads_per_worker (int, optional): Number of threads per worker.
            - verbose (bool, optional): Enables detailed logging.

        Args:
            memory_limit (str, optional): Memory limit per worker.
            n_jobs (int, optional): Number of parallel jobs.
            no_worker_restart (bool, optional): Whether to prevent worker restarts.
            on_slurm (bool, optional): Whether running on a SLURM cluster.
            temp_dir (str, optional): Temporary directory for Dask workers.
            threads_per_worker (int, optional): Number of threads per worker.
            verbose (bool, optional): Enables detailed logging.

        Returns:
            None
        """

        params = {
            'memory_limit': memory_limit,
            'n_jobs': n_jobs,
            'no_worker_restart': no_worker_restart,
            'on_slurm': on_slurm,
            'temp_dir': temp_dir,
            'threads_per_worker': threads_per_worker,
            'no_distributed': no_distributed,
            'verbose': verbose
        }

        for key in params:
            if key in self.config['cluster'].keys():
                if params[key] != 'default':
                    self.config['cluster'][key] = params[key]
        self.config_gr.attrs['cluster'] = self.config['cluster']

    def configure_readers(self,
                          as_mosaic: bool = 'default',
                          view_index: int = 'default',
                          phase_index: int = 'default',
                          illumination_index: int = 'default',
                          scene_index: int = 'default',
                          rotation_index: int = 'default',
                          mosaic_tile_index: int = 'default',
                          sample_index: int = 'default',
                         ):
        """
        Updates reader configuration settings. To update the current default value for a parameter, provide that parameter with a value other than 'default'.

        Returns:
            None
        """

        params = {
            'as_mosaic': as_mosaic,
            'view_index': view_index,
            'phase_index': phase_index,
            'illumination_index': illumination_index,
            'scene_index': scene_index,
            'rotation_index': rotation_index,
            'mosaic_tile_index': mosaic_tile_index,
            'sample_index': sample_index
        }

        for key in params:
            if key in self.config['readers'].keys():
                if params[key] != 'default':
                    self.config['readers'][key] = params[key]
        self.config_gr.attrs['readers'] = self.config['readers']


    def configure_conversion(self,
                             zarr_format: int = 'default',
                             compressor: str = 'default',
                             compressor_params: dict = 'default',
                             time_chunk: int = 'default',
                             channel_chunk: int = 'default',
                             z_chunk: int = 'default',
                             y_chunk: int = 'default',
                             x_chunk: int = 'default',
                             time_shard_coef: int = 'default',
                             channel_shard_coef: int = 'default',
                             z_shard_coef: int = 'default',
                             y_shard_coef: int = 'default',
                             x_shard_coef: int = 'default',
                             time_range: int = 'default',
                             channel_range: int = 'default',
                             z_range: int = 'default',
                             y_range: int = 'default',
                             x_range: int = 'default',
                             dimension_order: str = 'default',
                             overwrite: bool = 'default',
                             rechunk_method: str = 'default',
                             # rechunkers_max_mem: str = 'default',
                             trim_memory: bool = 'default',
                             use_tensorstore: bool = 'default',
                             use_gpu: bool = 'default',
                             metadata_reader: str = 'default',
                             save_omexml: bool = 'default',
                             squeeze: bool = 'default'
                             ):
        """
        Updates conversion configuration settings. To update the current default value for a parameter, provide that parameter with a value other than 'default'.

        The following parameters can be configured:
            - compressor (str, optional): Compression algorithm.
            - compressor_params (dict, optional): Parameters for the compressor.
            - output_chunks (list, optional): Chunk size for output.
            - overwrite (bool, optional): Whether to overwrite existing data.
            - rechunk_method (str, optional): Method used for rechunking.
            - trim_memory (bool, optional): Whether to trim memory usage.
            - use_tensorstore (bool, optional): Whether to use TensorStore for writing.
            - save_omexml (bool, optional): Whether to create a METADATA.ome.xml file.
        Args:
            compressor (str, optional): Compression algorithm.
            compressor_params (dict, optional): Parameters for the compressor.
            output_chunks (list, optional): Chunk size for output.
            overwrite (bool, optional): Whether to overwrite existing data.
            rechunk_method (str, optional): Method used for rechunking.
            trim_memory (bool, optional): Whether to trim memory usage.
            use_tensorstore (bool, optional): Whether to use TensorStore for storage.
            save_omexml (bool, optional): Whether to create a METADATA.ome.xml file.

        Returns:
            None
        """

        params = {
            'zarr_format': zarr_format,
            'compressor': compressor,
            'compressor_params': compressor_params,
            "time_chunk": time_chunk,
            "channel_chunk": channel_chunk,
            "z_chunk": z_chunk,
            "y_chunk": y_chunk,
            "x_chunk": x_chunk,
            "time_shard_coef": time_shard_coef,
            "channel_shard_coef": channel_shard_coef,
            "z_shard_coef": z_shard_coef,
            "y_shard_coef": y_shard_coef,
            "x_shard_coef": x_shard_coef,
            "time_range": time_range,
            "channel_range": channel_range,
            "z_range": z_range,
            "y_range": y_range,
            "x_range": x_range,
            "dimension_order": dimension_order,
            'overwrite': overwrite,
            'rechunk_method': rechunk_method,
            # 'rechunkers_max_mem': rechunkers_max_mem,
            'trim_memory': trim_memory,
            'use_tensorstore': use_tensorstore,
            'use_gpu': use_gpu,
            'metadata_reader': metadata_reader,
            'save_omexml': save_omexml,
            'squeeze': squeeze
        }

        for key in params:
            if key in self.config['conversion'].keys():
                if params[key] != 'default':
                    self.config['conversion'][key] = params[key]
        self.config_gr.attrs['conversion'] = self.config['conversion']

    def configure_downscale(self,
                            downscale_method: str = 'default',
                            n_layers: int = 'default',
                            time_scale_factor: int = 'default',
                            channel_scale_factor: int = 'default',
                            z_scale_factor: int = 'default',
                            y_scale_factor: int = 'default',
                            x_scale_factor: int = 'default',
                            ):
        """
        Updates downscaling configuration settings. To update the current default value for a parameter, provide that parameter with a value other than 'default'.

        The following parameters can be configured:
            - downscale_method (str, optional): Downscaling algorithm.
            - n_layers (int, optional): Number of downscaling layers.
            - scale_factor (list, optional): Scaling factors for each dimension.

        Args:
            downscale_method (str, optional): Downscaling algorithm.
            n_layers (int, optional): Number of downscaling layers.
            scale_factor (list, optional): Scaling factors for each dimension.

        Returns:
            None
        """

        params = {
            'downscale_method': downscale_method,
            'n_layers': n_layers,
            'time_scale_factor': time_scale_factor,
            "channel_scale_factor": channel_scale_factor,
            "z_scale_factor": z_scale_factor,
            "y_scale_factor": y_scale_factor,
            "x_scale_factor": x_scale_factor,
        }

        for key in params:
            if key in self.config['downscale'].keys():
                if params[key] != 'default':
                    self.config['downscale'][key] = params[key]
        self.config_gr.attrs['downscale'] = self.config['downscale']

    def _set_dask_temp_dir(self, temp_dir = 'auto'):
        if self._dask_temp_dir is not None:
            self._dask_temp_dir.cleanup()
        if temp_dir in ('auto', None):
            temp_dir = tempfile.TemporaryDirectory()
        else:
            os.makedirs(temp_dir, exist_ok=True)
            temp_dir = tempfile.TemporaryDirectory(dir=temp_dir)
        self._dask_temp_dir = temp_dir
        return self

    def _start_cluster(self,
                    n_jobs: int = 4,
                    threads_per_worker: int = 1,
                    memory_limit: str = 'auto',
                    temp_dir='auto',
                    verbose = False,
                    on_slurm = False,
                    no_distributed = False,
                    config_kwargs = {},
                    **kwargs
                    ):

        config_dict = copy.deepcopy(self.dask_config)
        config_dict.update(**config_kwargs)

        # worker_options = {
        #     "memory_target_fraction": 0.8,
        #     "memory_spill_fraction": 0.9,
        #     "memory_pause_fraction": 0.95,
        #     # "memory_terminate_fraction": 0.98
        # }
        scheduler_options = {
            "allowed_failures": 100,
            "idle_timeout": "1h",
            "worker_ttl": "1d"  # Set to a large value, e.g., 1 day
        }

        self._set_dask_temp_dir(temp_dir)

        dask.config.set(config_dict) # use dictionary notation here.

        if no_distributed:
            config_dict.update(scheduler = 'threads',
                               pool = ThreadPool(n_jobs)
                               )
            dask.config.set(config_dict)
            print(f"Process running locally via multithreading.")
        else:
            if memory_limit == 'auto':
                reserve_fraction = kwargs.get('reserve_memory_fraction', 0.1)
                min_per_worker = kwargs.get('min_memory_per_worker', 1 * 1024 ** 3)

                total_mem = psutil.virtual_memory().total
                reserved_mem = total_mem * reserve_fraction
                available_mem = max(0, total_mem - reserved_mem)
                mem_per_worker = max(available_mem / n_jobs, min_per_worker)
                mem_gb = mem_per_worker / (1 * 1024 ** 3)
                memory_limit = f"{mem_gb} GB"
                print(f"{memory_limit} memory has been allocated per worker.")

            if on_slurm:
                print(f"Process running on Slurm.")
                cluster = SLURMCluster(
                                        cores=threads_per_worker,
                                        processes=1,
                                        nanny=False,
                                        scheduler_options=scheduler_options,
                                        n_workers=n_jobs,
                                        memory=memory_limit,
                                        local_directory=f"{self._dask_temp_dir.name}",
                                        # **worker_options
                                        )
            else:
                print(f"Process running on local cluster.")
                cluster = LocalCluster(
                                       n_workers=n_jobs,
                                       threads_per_worker=threads_per_worker,
                                       nanny = False,
                                       scheduler_kwargs=scheduler_options,
                                       memory_limit = memory_limit,
                                       local_directory=f"{self._dask_temp_dir.name}",
                                       # **worker_options
                                       )
            cluster.scale(n_jobs)
            self.client = Client(cluster)
            if verbose:
                print(self.client.cluster)
        return self

    def to_zarr(self,
                input_path: Union[Path, str],
                output_path: Union[Path, str],
                includes=None,
                excludes=None,
                time_tag: Union[str, tuple] = None,
                channel_tag: Union[str, tuple] = None,
                z_tag: Union[str, tuple] = None,
                y_tag: Union[str, tuple] = None,
                x_tag: Union[str, tuple] = None,
                concatenation_axes: Union[int, tuple, str] = None,
                **kwargs
                ):
        """
        Converts image data to OME-Zarr format and optionally applies downscaling.

        Args:
            input_path (Union[Path, str]): Path to input file or directory.
            output_path (Union[Path, str]): Directory, in which the output OME-Zarrs will be written.
            includes (str, optional): Filename patterns to filter for.
            excludes (str, optional): Filename patterns to filter against.
            time_tag (Union[str, tuple], optional): Time dimension tag.
            channel_tag (Union[str, tuple], optional): Channel dimension tag.
            z_tag (Union[str, tuple], optional): Z dimension tag.
            y_tag (Union[str, tuple], optional): Y dimension tag.
            x_tag (Union[str, tuple], optional): X dimension tag.
            concatenation_axes (Union[int, tuple, str], optional): Axes, along which the images will be concatenated.
            **kwargs: Additional configuration overrides.

        Raises:
            Exception: If no files are found in the input path.

        Prints:
            Process logs including conversion and downscaling time.

        Returns:
            None
        """
        t0 = time.time()
        # Get parameters:
        self.cluster_params = self._collect_params('cluster', **kwargs)
        self.readers_params = self._collect_params('readers', **kwargs)
        self.conversion_params = self._collect_params('conversion', **kwargs)
        self.downscale_params = self._collect_params('downscale', **kwargs)

        if self.conversion_params['use_gpu'] and self.conversion_params['use_tensorstore']:
            raise ValueError("Tensorstore is not supported for GPU arrays.")

        print(f"Base conversion initiated.")
        ###### Handle input data and metadata
        paths = take_filepaths(input_path, includes = includes, excludes = excludes)

        filepaths = sorted(list(paths))

        ###### Start the cluster
        verified_for_cluster = verify_filepaths_for_cluster(filepaths) ### Ensure non-bioformats conversion. If bioformats is needed, fall back on local conversion.
        if not verified_for_cluster:
            self.cluster_params['no_distributed'] = True

        cluster_is_true = not self.cluster_params['no_distributed']

        if cluster_is_true:
            chunks_yx = None
        else:
            chunks_yx = tuple([self.conversion_params['y_chunk'],  # chunks already when reading
                               self.conversion_params['x_chunk']]
                              )

        self._start_cluster(**self.cluster_params)


        series = self.readers_params['scene_index']

        ###### Read and concatenate
        base = BridgeBase(input_path,
                        excludes=excludes,
                        includes=includes,
                        series=series,
                        zarr_format = self.conversion_params['zarr_format'],
                        verbose = self.cluster_params['verbose']
                        )

        base.read_dataset(cluster_is_true,
                          chunks_yx=chunks_yx,
                          readers_params = self.readers_params
                          )

        base.digest(
            time_tag = time_tag,
            channel_tag = channel_tag,
            z_tag = z_tag,
            y_tag = y_tag,
            x_tag = x_tag,
            axes_of_concatenation = concatenation_axes,
            metadata_reader = self.conversion_params['metadata_reader'],
            **kwargs
            )
        print(f"Metadata was extracted")
        verbose = base._verbose

        if 'region_shape' in kwargs.keys():
            self.conversion_params['region_shape'] = kwargs.get('region_shape')
        if verbose:
            print(f"Cluster params:")
            pprint.pprint(self.cluster_params)
            print(f"Readers params:")
            pprint.pprint(self.readers_params)
            print(f"Conversion params:")
            pprint.pprint(self.conversion_params)
            print(f"Downscale params:")
            pprint.pprint(self.downscale_params)

        temp_dir = base._dask_temp_dir
        self.conversion_params['temp_dir'] = temp_dir
        self.downscale_params['temp_dir'] = temp_dir

        if self.client is not None:
            base.client = self.client
        base.set_dask_temp_dir(self._dask_temp_dir)

        ###### Write
        self.base_results = base.write_arrays(output_path,
                                           compute=True,
                                           verbose=verbose,
                                           **self.conversion_params
                                           )
        ###### Downscale
        print(f"Base conversion finished.")
        t1 = time.time()
        print(f"Elapsed for base conversion: {(t1 - t0) / 60} min.")
        n_layers = self.downscale_params['n_layers']
        if n_layers > 1:
            print(f"Downscaling initiated.")
            _ = downscale(
                      self.base_results,
                      **self.downscale_params,
                      zarr_format = self.conversion_params['zarr_format'],
                      rechunk_method = self.conversion_params['rechunk_method'],
                      use_tensorstore = self.conversion_params['use_tensorstore'],
                      verbose = verbose
                      ) # TODO: add to_cupy parameter here.

            print(f"Downscaling finished.")

        ###### Shutdown and clean up
        if self.client is not None:
            self.client.shutdown()
            self.client.close()

        if isinstance(self._dask_temp_dir, tempfile.TemporaryDirectory):
            shutil.rmtree(self._dask_temp_dir.name)
        else:
            shutil.rmtree(self._dask_temp_dir)

        t1 = time.time()
        print(f"Elapsed for conversion + downscaling: {(t1 - t0) / 60} min.")


    def show_pixel_meta(self,
                input_path: Union[Path, str],
                includes=None,
                excludes=None,
                series: int = None, # self.readers_params['scene_index'],
                **kwargs
                ):
        """
        Print pixel-level metadata for all datasets in the 'input_path'.

        Args:
            input_path (Union[Path, str]): Path to input file or directory.
            output_path (Union[Path, str]): Directory, in which the output OME-Zarrs will be written.
            includes (str, optional): Filename patterns to filter for.
            excludes (str, optional): Filename patterns to filter against.
            **kwargs: Additional configuration overrides.

        Raises:
            Exception: If no files are found in the input path.

        Prints:
            Process logs including conversion and downscaling time.

        Returns:
            None
        """

        # Get parameters:
        self.cluster_params = self._collect_params('cluster', **kwargs)
        self.readers_params = self._collect_params('readers', **kwargs)
        self.conversion_params = self._collect_params('conversion', **kwargs)

        paths = take_filepaths(input_path, includes = includes, excludes = excludes)

        filepaths = sorted(list(paths))

        ###### Start the cluster
        verified_for_cluster = verify_filepaths_for_cluster(filepaths)
        if not verified_for_cluster:
            self.cluster_params['no_distributed'] = True

        self._start_cluster(**self.cluster_params)

        series = self.readers_params['scene_index']

        ###### Read and digest
        base = BridgeBase(input_path,
                        excludes=excludes,
                        includes=includes,
                        series=series
                        )

        base.read_dataset(verified_for_cluster,
                          readers_params=self.readers_params
                          )

        base.digest()

        temp_dir = base._dask_temp_dir
        self.conversion_params['temp_dir'] = temp_dir

        if self.client is not None:
            base.client = self.client
        base.set_dask_temp_dir(self._dask_temp_dir)

        printables = {
            path: get_printables(
                manager.axes,
                manager.shapedict,
                manager.scaledict,
                manager.unitdict
            )
            for path, manager in base.batchdata.managers.items()
        }
        for path, printable in printables.items():
            print('---------')
            print(f"")
            print(f"Metadata for '{path}':")
            print_printable(printable)

        ###### Shutdown and clean up
        if self.client is not None:
            self.client.shutdown()
            self.client.close()

        if isinstance(self._dask_temp_dir, tempfile.TemporaryDirectory):
            shutil.rmtree(self._dask_temp_dir.name)
        else:
            shutil.rmtree(self._dask_temp_dir)

    def update_pixel_meta(self,
                          input_path: Union[Path, str],
                          includes=None,
                          excludes=None,
                          time_scale: (int, float) = None,
                          z_scale: (int, float) = None,
                          y_scale: (int, float) = None,
                          x_scale: (int, float) = None,
                          time_unit: str = None,
                          z_unit: str = None,
                          y_unit: str = None,
                          x_unit: str = None,
                          **kwargs
                          ):
        """
        Updates pixel metadata for image files located at the specified input path.

        Args:
            input_path (Union[Path, str]): Path to input file or directory.
            includes (optional): Filename patterns to include.
            excludes (optional): Filename patterns to exclude.
            series (int, optional): Series index to process.
            time_scale, z_scale, y_scale, x_scale ((int, float), optional): Scaling factors for the respective dimensions.
            time_unit, z_unit, y_unit, x_unit (str, optional): Units for the respective dimensions.
            **kwargs: Additional parameters for cluster and conversion configuration.

        Returns:
            None
        """

        # Collect cluster and conversion parameters
        self.cluster_params = self._collect_params('cluster', **kwargs)
        self.readers_params = self._collect_params('readers', **kwargs)
        self.conversion_params = self._collect_params('conversion', **kwargs)

        # Collect file paths based on inclusion and exclusion patterns
        paths = take_filepaths(input_path, includes=includes, excludes=excludes)

        filepaths = sorted(list(paths))

        # Verify file paths for cluster compatibility
        verified_for_cluster = verify_filepaths_for_cluster(filepaths)
        if not verified_for_cluster:
            self.cluster_params['no_distributed'] = True

        cluster_is_true = not self.cluster_params['no_distributed']

        # Start the processing cluster
        self._start_cluster(**self.cluster_params)

        series = self.readers_params['scene_index']

        base = BridgeBase(input_path,
                          excludes=excludes,
                          includes=includes,
                          series=series
                          )

        # Read and digest the dataset
        base.read_dataset(verified_for_cluster,
                          readers_params=self.readers_params
                          )

        # Prepare pixel metadata arguments
        pixel_meta_kwargs_ = dict(time_scale=time_scale,
                                  z_scale=z_scale,
                                  y_scale=y_scale,
                                  x_scale=x_scale,
                                  time_unit=time_unit,
                                  z_unit=z_unit,
                                  y_unit=y_unit,
                                  x_unit=x_unit)
        pixel_meta_kwargs = {key: val for key, val in pixel_meta_kwargs_.items() if val is not None}

        base.digest(**pixel_meta_kwargs)
        print(f"Metadata was extracted")

        if self.client is not None:
            base.client = self.client
        base.set_dask_temp_dir(self._dask_temp_dir)

        # Update metadata for each dataset manager
        for path, manager in base.batchdata.managers.items():
            if is_zarr_group(manager.path):
                manager.sync_pyramid(self.conversion_params['save_omexml'])
            else:
                print(f"Cannot update metadata for non-zarr path: {path}")

        # Shutdown the cluster and clean up temporary directories
        if self.client is not None:
            self.client.shutdown()
            self.client.close()

        if isinstance(self._dask_temp_dir, tempfile.TemporaryDirectory):
            shutil.rmtree(self._dask_temp_dir.name)
        else:
            shutil.rmtree(self._dask_temp_dir)



