from os import cpu_count
from bfio import BioReader, BioWriter, OmeXml
import argparse, logging
import numpy as np
from pathlib import Path
from cellpose import dynamics, utils
import torch
from concurrent.futures import ThreadPoolExecutor, wait, Future
import typing

""" Plugin Constants """
TILE_SIZE = 2048    # Largest chunk of an image to process
TILE_OVERLAP = 256  # Amount of overlap between tiles
NITER = 200         # Number of iterations to run flow dynamics

# Use a gpu if it's available
USE_GPU = torch.cuda.is_available()
if USE_GPU:
    DEV = torch.device("cuda")
else:
    DEV = torch.device("cpu")

# Initialize the logger
logging.basicConfig(format='%(asctime)s - %(name)-8s - %(levelname)-8s - %(message)s',
                    datefmt='%d-%b-%y %H:%M:%S')
logger = logging.getLogger("main")
logger.setLevel(logging.INFO)

def overlap(previous_values: np.ndarray,
            current_values: np.ndarray,
            tile: np.ndarray
            ) -> typing.Tuple[np.ndarray,list,list]:
    """Resolve label values between tiles
    
    This function takes a row/column from the previous tile and a row/column
    from the current tile and finds labels that that likely match. If labels
    in the current tile should be replaced with labels from the previous tile,
    the pixels in the current tile are removed from ``tile`` and the label value
    and pixel coordinates of the label are stored in ``labels`` and ``indices``
    respectively.

    Args:
        previous_values (np.ndarray): Previous tile edge values
        current_values (np.ndarray): Current tile edge values
        tile (np.ndarray): Current tile pixel values, flattened

    Returns:
        typing.Tuple[np.ndarray,np.ndarray,np.ndarray]: Returns the modified
            tile with overlapping labels removed, a list of new labels, and a
            list of indices associated with the new labels.
    """
    
    # Get a list of unique values in the previous and current tiles
    previous_labels = np.unique(previous_values)
    if previous_values[0] == 0:
        previous_labels = previous_values[1:]
    
    current_labels = np.unique(current_values)
    if current_labels[0] == 0:
        current_labels = current_labels[1:]
    
    # Initialize outputs
    indices = []
    labels = []
    
    if previous_labels.size != 0 and current_labels.size != 0:
    
        # Find overlapping indices
        for label in current_labels:
            
            new_labels,counts = np.unique(previous_values[current_values==label],return_counts=True)
            
            if new_labels.size == 0:
                continue
            
            if new_labels[0] == 0:
                new_labels = new_labels[1:]
                counts = counts[1:]
                
            if new_labels.size == 0:
                continue
                
            # Get the most frequently occuring overlapping label
            labels.append(new_labels[np.argmax(counts)])
            
            # Add indices to output, remove pixel values from the tile
            indices.append(np.argwhere(tile==label))
            tile[indices[-1]] = 0
            
    return tile, labels, indices

def mask_thread(coords: typing.Tuple[int,int,int],
                file_path: Path,
                bw: BioWriter,
                cellprob_threshold: float,
                flow_threshold: float,
                dependency1: typing.Optional[Future],
                dependency2: typing.Optional[Future]
                ) -> typing.Tuple[np.ndarray,np.ndarray,np.uint32]:
    """[summary]

    Args:
        coords: x,y,z starting coordinates of the tile to process
        file_path: Vector field file path
        bw: Output file
        cellprob_threshold: Cell probability threshold
        flow_threshold: Flow field threshold
        dependency1: Tile future to the left of the current tile
        dependency2: Tile future above the current tile

    Returns:
        typing.Tuple[np.ndarray,np.ndarray,list]: Returns the right column of
            the processed tile, bottom row of the processed tile, and largest
            label value.
    """   
    
    # Calculate indice for the tile
    x,y,z = coords
    with BioReader(file_path) as br:
        
        image_X = br.X
        image_Y = br.Y
        
        x_min = max([0, x - TILE_OVERLAP])
        x_max = min([image_X, x + TILE_SIZE + TILE_OVERLAP])

        y_min = max([0, y - TILE_OVERLAP])
        y_max = min([image_Y, y + TILE_SIZE + TILE_OVERLAP])

        tile = br[y_min:y_max, x_min:x_max, z:z + 1, :3, 0]

    logger.debug('Calculating flows and masks  for tile [{}:{},{}:{},{}:{}]'.format(y, y_max, x,
                 x_max, z, z + 1))

    # Get flows and probabilities
    cellprob = tile[:,:,0,0].squeeze()
    dP = tile[:,:,0,1:].squeeze().transpose(2,0,1)

    # Compute flows for the tile
    p = dynamics.follow_flows(-1 * dP * (cellprob > cellprob_threshold) / 5., 
                                niter=NITER, interp=True,
                                use_gpu=USE_GPU,device=DEV)
    mask = dynamics.get_masks(p, iscell=(cellprob>cellprob_threshold),
                              flows=dP, threshold=flow_threshold,
                              use_gpu=USE_GPU,device=DEV)
    mask = utils.fill_holes_and_remove_small_masks(mask, min_size=15)

    # reshape mask based on tile
    x_overlap = x - x_min
    x_min = x
    x_max = min([image_X, x + TILE_SIZE])

    y_overlap = y - y_min
    y_min = y
    y_max = min([image_Y, y + TILE_SIZE])

    mask = mask[y_overlap:y_max - y_min + y_overlap,
                x_overlap:x_max - x_min + x_overlap,
                np.newaxis, np.newaxis, np.newaxis].astype(np.uint32)
    
    """ Fix tile conflicts if image is large enough to require tiling """
    # Get previously processed tiles if they exist
    dependency1 = None if dependency1 is None else dependency1.result()
    dependency2 = None if dependency2 is None else dependency2.result()
    
    # Get offset to make labels consistent between tiles
    offset = 0 if dependency1 is None else dependency1[2]
    
    current_x = mask[:,0].squeeze()
    current_y = mask[0,:].squeeze()
    shape = mask.shape
    mask = mask.reshape(-1)
    
    # Resolve label conflicts along the left border
    if x > 0:
        
        mask, labels_x, indices_x = overlap(dependency1[0].squeeze(),current_x,mask)
        
    if y > 0:
        
        mask, labels_y, indices_y = overlap(dependency2[1].squeeze(),current_y,mask)
    
    _, image = np.unique(mask, return_inverse=True)
    image = image.astype(np.uint32)
    image[image>0] = image[image>0] + offset
    
    if x > 0:
        for label,ind in zip(labels_x,indices_x):
            if ind.size==0:
                continue
            image[ind] = label
            
    if y > 0:
        for label,ind in zip(labels_y,indices_y):
            if ind.size==0:
                continue
            image[ind] = label

    image = image.reshape(shape)
    bw[y_min:y_max, x_min:x_max, z:z + 1, 0, 0] = image
    
    return image[:,-1],image[-1,:],image.max()

def close_thread(dependency: Future,
                 bw: BioWriter):
    """ Close an image once the final tile is written

    Args:
        dependency (Future): The final tile thread
        bw (BioWriter): The BioWriter to clsoe

    Returns:
        Returns True when completed
    """    
    
    dependency.result()
    
    bw.close()
    
    return True

def main(inpDir: Path,
         cellprob_threshold: float,
         flow_threshold: float,
         outDir: Path
         ) -> None:
    
    # Get the list of files in path
    files = [p for p in Path(inpDir).iterdir() if p.name.endswith('_flow.ome.zarr')]
    
    num_threads = max([cpu_count()//2,1])
    logger.info(f'Processing tiles with {num_threads} threads using {DEV}')
    
    if len(files) == 0:
        logger.critical('No flow files detected.')
        quit()
    
    processes = []
    with ThreadPoolExecutor(num_threads) as executor:

        # Loop through files in inpDir image collection and process
        for ind,fpath in enumerate(files):
        
            br = BioReader(fpath)
                
            threads = np.empty((br.shape[:3]),dtype=object)
                            
            logger.debug(
                'Processing image ({}/{}): {}'.format(ind, len(files),
                                                      fpath))

            # TODO: Hard coding to ome.tif for now, this should be changed later.
            path = Path(outDir).joinpath(fpath.name.replace('_flow.ome.zarr','.ome.tif'))

            bw = BioWriter(file_path=Path(path), metadata=br.metadata)
            bw.dtype=np.dtype(np.uint32)
            bw.C = 1
            bw.channel_names = ['label']

            for z in range(0, br.Z, 1):
                
                y_ind = None
                dependency1 = None

                for y in range(0, br.Y, TILE_SIZE):

                    for x in range(0, br.X, TILE_SIZE):
                        
                        dependency2 = None if y_ind is None else threads[y_ind,x//TILE_SIZE,z]

                        processes.append(executor.submit(mask_thread,
                                                         (x,y,z),
                                                         fpath,bw,
                                                         cellprob_threshold,flow_threshold,
                                                         dependency1,dependency2))
                        dependency1 = processes[-1]
                        threads[y//TILE_SIZE,x//TILE_SIZE,z] = dependency1
                    
                    y_ind = y//TILE_SIZE
            
            executor.submit(close_thread,dependency1,bw)
                        
        done, not_done = wait(processes, 0)

        logger.info(f'Percent complete: {100 * len(done) / len(processes):6.3f}%')

        while len(not_done) > 0:
            for r in done:
                r.result()
            done, not_done = wait(processes, 15)
            logger.info(f'Percent complete: {100 * len(done) / len(processes):6.3f}%')

if __name__ == '__main__':
    
    ''' Argument parsing '''
    logger.info("Parsing arguments...")
    parser = argparse.ArgumentParser(prog='main', description='Cellpose parameters')
    
    # Input arguments
    parser.add_argument('--inpDir', dest='inpDir', type=str,
                        help='Input image collection to be processed by this plugin', required=True)
    parser.add_argument('--flowThreshold', required=False,
                        default=0.8, type=float, help='flow error threshold, 0 turns off this optional QC step')
    parser.add_argument('--cellprobThreshold', required=False,
                        default=0.0, type=float, help='cell probability threshold, centered at 0.0')
    
    # Output arguments
    parser.add_argument('--outDir', dest='outDir', type=str,
                        help='Output collection', required=True)
    
    # Parse the arguments
    args = parser.parse_args()
    inpDir = Path(args.inpDir)
    logger.info('inpDir = {}'.format(inpDir))
    outDir = args.outDir
    logger.info('outDir = {}'.format(outDir))
    cellprob_threshold = args.cellprobThreshold
    logger.info('cellprobThreshold = {}'.format(cellprob_threshold))
    flow_threshold= args.flowThreshold
    logger.info('flowThreshold = {}'.format(flow_threshold))
    
    main(inpDir,
         cellprob_threshold,
         flow_threshold,
         outDir)