#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Created on Tue May 12 23:47:52 2020
# @author: Ajit Johnson Nirmal
"""
!!! abstract "Short Description"
    `sm.pl.gate_finder`: The function opens the OME-TIFF image inside Napari and overlays points to help with the 
    identifying manual gates for each marker. Use the `sm.pp.rescale` function to apply the identified gates to your data.

## Function
"""

try:
    import napari
except:
    pass

import pandas as pd
import tifffile as tiff
import numpy as np

import dask.array as da
from dask.cache import Cache
import zarr
import os

cache = Cache(2e9)  # Leverage two gigabytes of memory
cache.register()

def gate_finder (image_path, adata, marker_of_interest, from_gate = 6, to_gate = 8, increment = 0.1,
                 markers=None, channel_names = 'default', flip_y=True,
                 x_coordinate='X_centroid',y_coordinate='Y_centroid',
                 point_size=10,imageid='imageid',subset=None,seg_mask=None,**kwargs):
    """
Parameters:
    image_path : string  
        Location to the image file.

    adata : Ann Data Object  

    marker_of_interest : string  
        Marker for which gate is to be defined e.g. 'CD45'.

    from_gate : int, optional  
        Start value gate of interest.

    to_gate : int, optional  
        End value of the gate of interest.
        
    flip_y : bool, optional  
        Flip the Y-axis if needed. Some algorithms output the XY with the Y-coordinates flipped.
        If the image overlays do not align to the cells, try again by setting this to `False`.

    increment : float, optional  
        Increments between the start and end values.

    markers : string, optional  
        Additional markers to be included in the plot for evaluation.

    channel_names : list, optional  
        List of channels in the image in the exact order as image. The default is `adata.uns['all_markers']`

    x_coordinate : string, optional  
        X axis coordinate column name in AnnData object.

    y_coordinate : string, optional  
        Y axis coordinate column name in AnnData object.

    point_size : int, optional  
        point size in the napari plot.

    imageid : string, optional  
        Column name of the column containing the image id.

    subset : string, optional  
        imageid of a single image to be subsetted for analyis.

    seg_mask : string, optional  
        Location to the segmentation mask file.

    **kwargs  
        Other arguments that can be passed to napari viewer.

Example:
```python
    image_path = '/Users/aj/Desktop/ptcl_tma/image.ome.tif'
    sm.pl.gate_finder (image_path, adata, marker_of_interest='CD45',
                 from_gate = 6, to_gate = 8, increment = 0.1,
                 markers=['DNA10'], channel_names = 'default',
                 x_coordinate='X_position',y_coordinate='Y_position',point_size=10,
                 subset= '77', seg_mask=None)
```
    """

    # If no raw data is available make a copy
    if adata.raw is None:
        adata.raw = adata

    # Copy of the raw data if it exisits
    if adata.raw is not None:
        adata.X = adata.raw.X

    # Make a copy of the data with the marker of interest
    data = pd.DataFrame(np.log1p(adata.X), columns = adata.var.index, index= adata.obs.index)[[marker_of_interest]]

    # Generate a dataframe with various gates
    def gate (g, d):
        dd = d.values
        dd = np.where(dd < g, np.nan, dd)
        np.warnings.filterwarnings('ignore')
        dd = np.where(dd > g, 1, dd)
        dd = pd.DataFrame(dd, index = d.index, columns = ['gate-' + str(g)])
        return dd

    # Identify the list of increments
    inc = list(np.arange (from_gate, to_gate, increment))
    inc = [round(num,3) for num in inc]

    # Apply the function
    r_gate = lambda x: gate(g=x, d=data) # Create lamda function
    gated_data = list(map(r_gate, inc)) # Apply function
    # Concat all the results into a single dataframe
    gates = pd.concat(gated_data, axis=1)
    
    # Plot only the Image that is requested
    if subset is not None:
        adata = adata[adata.obs[imageid] == subset]
        
    # Recover the channel names from adata
    if channel_names == 'default':
        channel_names = adata.uns['all_markers']
    else:
        channel_names = channel_names

    # if markers is a string convert to list
    if isinstance(markers, str):
        markers = [markers]
        
    # Index of the marker of interest and corresponding names
    if markers is not None:
        markers.extend([marker_of_interest])
        idx = np.where(np.isin(channel_names,markers))[0]
        channel_names = [channel_names[i] for i in idx]
    else:
        idx = list(range(len(channel_names)))
        channel_names = channel_names


    # Load the segmentation mask
    if seg_mask is not None:
        seg_m = tiff.imread(seg_mask)
        if (len(seg_m.shape) > 2) and (seg_m.shape[0] > 1):
            seg_m = seg_m[0]
        

    ##########################################################################
    # Visulaisation using Napari
    
    # load OME TIFF
    if os.path.isfile(image_path) is True: 
        # Load the image    
        image = tiff.TiffFile(image_path, is_ome=False)
        z = zarr.open(image.aszarr(), mode='r') # convert image to Zarr array
        # Identify the number of pyramids and number of channels
        n_levels = len(image.series[0].levels) # pyramid  
        # If and if not pyramids are available
        if n_levels > 1:
            pyramid = [da.from_zarr(z[i]) for i in range(n_levels)]
            multiscale = True
        else:
            pyramid = da.from_zarr(z)
            multiscale = False   
        # subset channels of interest
        if markers is not None:
            if n_levels > 1:
                for i in range(n_levels-1):
                    pyramid[i] = pyramid[i][idx, :, :]
                n_channels = pyramid[0].shape[0] # identify the number of channels
            else:
                pyramid = pyramid[idx, :, :]
                n_channels = pyramid.shape[0] # identify the number of channels
        else:
            if n_levels > 1:
                n_channels = pyramid[0].shape[0]
            else:
                n_channels = pyramid.shape[0]\
        # check if channel names have been passed to all channels
        if channel_names is not None:
            assert n_channels == len(channel_names), (
                f'number of channel names ({len(channel_names)}) must '
                f'match number of channels ({n_channels})'
            )

        # Load the viewer
        viewer = napari.view_image(
        pyramid,
        channel_axis = 0,
        multiscale=multiscale,
        name = None if channel_names is None else channel_names,
        visible = False, **kwargs)


    # Operations on the ZARR image
    # check the format of image
    if os.path.isfile(image_path) is False: 
        #print(image_path)
        viewer = napari.Viewer()
        viewer.open(image_path, multiscale=True,
                    visible=False,
                    name = None if channel_names is None else channel_names)

    # Add the seg mask
    if seg_mask is not None:
        viewer.add_labels(seg_m, name='segmentation mask', visible=False)

    # subset the gates to include only the image of interest
    gates = gates.loc[adata.obs.index,]

    # Add gating layer
    def add_phenotype_layer (adata, gates, phenotype_layer,x,y,viewer,point_size):
        cells = gates[gates[phenotype_layer] == 1].index
        coordinates = adata[cells]
        # Flip Y axis if needed
        if flip_y is True:
            coordinates = pd.DataFrame({'y': coordinates.obs[y],'x': coordinates.obs[x]})
        else:  
            coordinates = pd.DataFrame({'x': coordinates.obs[x],'y': coordinates.obs[y]})
        #points = coordinates.values.tolist()
        points = coordinates.values
        #import time
        #start = time.time()
        viewer.add_points(points, size=point_size, face_color='white',visible=False,name=phenotype_layer)
        #stop = time.time()
        #print(stop-start)
        

    # Run the function on all gating layer
    for i in gates.columns:
        add_phenotype_layer (adata=adata, gates=gates, 
                             phenotype_layer=i, x=x_coordinate, y=y_coordinate, 
                             viewer=viewer, point_size=point_size)
