"""Pixel Quality Metrics"""

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

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

# %% ../../nbs/CLI/pqm.ipynb 5
import logging
import zarr
import time
import numpy as np

import dask
from dask import array as da
from dask import delayed
from dask.distributed import Client, LocalCluster, progress
from ..utils_ import is_cuda_available, get_array_module
if is_cuda_available():
    import cupy as cp
    from dask_cuda import LocalCUDACluster
    from rmm.allocators.cupy import rmm_cupy_allocator
import moraine as mr
from .logging import mc_logger
from . import dask_from_zarr, dask_to_zarr

# %% ../../nbs/CLI/pqm.ipynb 6
@mc_logger
def temp_coh(
    intf:str, # complex interferograms/coherence metrix, dtype complex64, shape 2D(pc) or 3D(ras)
    rslc:str, # complex rslc/phase history, dtype complex64, shape 2D(pc) or 3D(ras)
    t_coh:str=None, # output, temporal coherence
    image_pairs:np.ndarray=None, # image pairs
    chunks:int|tuple[int,int]=None, # ras/pc chunk size, same as intf by default
    cuda:bool=False, # if use cuda for processing, false by default
    processes=None, # use process for dask worker over thread, the default is False for cpu, only applied if cuda==False
    n_workers=None, # number of dask worker, the default is 1 for cpu, number of GPU for cuda
    threads_per_worker=None, # number of threads per dask worker, the default is 2 for cpu, only applied if cuda==False
    rmm_pool_size=0.9, # set the rmm pool size, only applied when cuda==True
    **dask_cluster_arg, # other dask local/cudalocal cluster args
):
    '''temporal coherence.
    '''
    intf_path = intf
    rslc_path = rslc
    t_coh_path = t_coh

    logger = logging.getLogger(__name__)
    intf_zarr = zarr.open(intf_path,mode='r'); logger.zarr_info(intf_path,intf_zarr)
    rslc_zarr = zarr.open(rslc_path,mode='r'); logger.zarr_info(rslc_path,rslc_zarr)
    nimage = rslc_zarr.shape[-1]

    if chunks is None: chunks = intf_zarr.chunks[:-1]
    if cuda:
        Cluster = LocalCUDACluster; cluster_args= {
            'n_workers':n_workers,
            'rmm_pool_size':rmm_pool_size}
        cluster_args.update(dask_cluster_arg)
        xp = cp
    else:
        if processes is None: processes = False
        if n_workers is None: n_workers = 1
        if threads_per_worker is None: threads_per_worker = 1
        Cluster = LocalCluster; cluster_args = {'processes':processes, 'n_workers':n_workers, 'threads_per_worker':threads_per_worker}
        cluster_args.update(dask_cluster_arg)
        xp = np
        
    if image_pairs is None:
        image_pairs = mr.TempNet.from_bandwidth(nimage).image_pairs
    image_pairs = image_pairs.astype(np.int32)
    
    logger.info('starting dask local cluster.')
    with Cluster(**cluster_args) as cluster, Client(cluster) as client:
        logger.info('dask local cluster started.')
        logger.dask_cluster_info(cluster)
        if cuda: client.run(cp.cuda.set_allocator, rmm_cupy_allocator)

        cpu_intf = dask_from_zarr(intf_path,chunks=(*chunks,intf_zarr.shape[-1]))
        logger.darr_info('intf', cpu_intf)
        
        cpu_rslc = dask_from_zarr(rslc_path,chunks=(*chunks,rslc_zarr.shape[-1]))
        logger.darr_info('rslc', cpu_rslc)

        logger.info(f'Estimate temporal coherence for DS.')
        if cuda:
            intf = cpu_intf.map_blocks(cp.asarray)
            rslc = cpu_rslc.map_blocks(cp.asarray)
        else:
            intf = cpu_intf
            rslc = cpu_rslc

        intf_delayed = intf.to_delayed()
        intf_delayed = np.squeeze(intf_delayed,axis=-1)
        rslc_delayed = rslc.to_delayed()
        rslc_delayed = np.squeeze(rslc_delayed,axis=-1)
        t_coh_delayed = np.empty_like(intf_delayed,dtype=object)
        temp_coh_delayed = delayed(mr.temp_coh,pure=True,nout=1)

        with np.nditer(intf_delayed,flags=['multi_index','refs_ok'], op_flags=['readwrite']) as it:
            for block in it:
                idx = it.multi_index
                t_coh_delayed[idx] = temp_coh_delayed(intf_delayed[idx],rslc_delayed[idx],image_pairs=image_pairs)
                t_coh_delayed[idx] = da.from_delayed(t_coh_delayed[idx],shape=intf.blocks[idx].shape[:-1],meta=xp.array((),dtype=xp.float32))

            t_coh = da.block(t_coh_delayed.tolist())
    
        if cuda:
            cpu_t_coh = t_coh.map_blocks(cp.asnumpy)
        else:
            cpu_t_coh = t_coh
        logger.info(f'got temporal coherence t_coh.')
        logger.darr_info('t_coh', t_coh)

        logger.info('saving t_coh.')
        _cpu_t_coh = cpu_t_coh.to_zarr(t_coh_path,compute=False,overwrite=True)

        logger.info('computing graph setted. doing all the computing.')
        futures = client.persist(_cpu_t_coh)
        progress(futures,notebook=False); time.sleep(0.1)
        da.compute(futures)
        logger.info('computing finished.')
    logger.info('dask cluster closed.')
