from qtpy.QtWidgets import QVBoxLayout, QWidget, QGroupBox 
from epicure.laptrack_centroids import LaptrackCentroids
from epicure.laptrack_overlaps import LaptrackOverlaps
from laptrack.data_conversion import convert_split_merge_df_to_napari_graph
from napari.utils import progress
import time
import pandas as pd
import numpy as np
from skimage.measure import regionprops, regionprops_table
from multiprocessing.pool import ThreadPool as Pool
import epicure.Utils as ut
import epicure.epiwidgets as wid

class Tracking(QWidget):
    """
        Handles tracking of cells, track operations with the Tracks layer
    """
    def __init__(self, napari_viewer, epic):
        super().__init__()
        self.viewer = napari_viewer
        self.epicure = epic
        self.graph = None      ## init 
        self.tracklayer = None      ## track layer with information (centroids, labels, tree..)
        self.track_data = None ## keep the updated data, and update the layer only from time to time (slow to do)
        self.tracklayer_name = "Tracks"  ## name of the layer containing tracks
        self.nframes = self.epicure.nframes
        self.properties = ["label", "centroid"]

        layout = QVBoxLayout()
        
        ## Add update track button 
        self.track_update = wid.add_button( "Update tracks display", self.update_track_layer, "Update the Track layer with the changements made since the last update" )
        layout.addWidget(self.track_update)
        

        ## Method specific
        track_method, self.track_choice = wid.list_line( "Tracking method", "Choose the tracking method to use and display its parameter", func=None )
        layout.addWidget(self.track_choice)
        
        self.track_choice.addItem("Laptrack-Centroids")
        self.create_laptrack_centroids()
        layout.addWidget(self.gLapCentroids)
        
        self.track_choice.addItem("Laptrack-Overlaps")
        self.create_laptrack_overlap()
        layout.addWidget(self.gLapOverlap)
        
        self.track_go = wid.add_button( "Track", self.do_tracking, "Launch the tracking with the current parameter. Can take time" )
        layout.addWidget(self.track_go)
        self.setLayout(layout)

        ## General tracking options
        frame_line, self.frame_range, self.range_group = wid.checkgroup_help( "Track only some frames", False, "Option to track only a given range of frames", None ) 
        self.frame_range.clicked.connect( self.show_frame_range )
        range_layout = QVBoxLayout()
        ntrack, self.start_frame = wid.ranged_value_line( "Track from frame:", 0, self.nframes-1, 1, 0, "Set first frame to begin tracking" )
        range_layout.addLayout(ntrack)
        
        entrack, self.end_frame = wid.ranged_value_line( "Until frame:", 1, self.nframes-1, 1, self.nframes-1, "Set the last frame unitl which to track" )
        range_layout.addLayout(entrack)
        self.start_frame.valueChanged.connect( self.changed_start )
        self.end_frame.valueChanged.connect( self.changed_end )
        
        self.range_group.setLayout( range_layout )
        layout.addWidget( self.frame_range )
        layout.addWidget( self.range_group )
        
        self.show_frame_range()
        self.show_trackoptions()
        self.track_choice.currentIndexChanged.connect(self.show_trackoptions)
        

    def show_frame_range( self ):
        """ Show/Hide frame range options """
        self.range_group.setVisible( self.frame_range.isChecked() )
        
    #### settings

    def get_current_settings( self ):
        """ Get current settings to save as preferences """
        settings = {}
        settings["Track method"] = self.track_choice.currentText() 
        settings["Add feat"] = self.check_penalties.isChecked()
        settings["Max distance"] = self.max_dist.text()
        settings["Splitting cost"] = self.splitting_cost.text()
        settings["Merging cutoff"] = self.merging_cost.text()
        settings["Min IOU"] = self.min_iou.text()
        settings["Over split"] = self.split_cost.text()
        settings["Over merge"] = self.merg_cost.text()
        return settings

    def apply_settings( self, settings ):
        """ Set the parameters/current display from the prefered settings """
        for setty, val in settings.items():
            if setty == "Track method":
                self.track_choice.setCurrentText( val )
            if setty == "Add feat":
                self.check_penalties.setChecked( val )
            if setty == "Max distance":
                self.max_dist.setText( val )
            if setty == "Splitting cost":
                self.splitting_cost.setText( val )
            if setty == "Merging cutoff":
                self.merging_cost.setText( val )
            if setty == "Min IOU":
                self.min_iou.setText( val )
            if setty == "Over split":
                self.split_cost.setText( val )
            if setty == "Over merge":
                self.merg_cost.setText( val )
            
    ##########################################
    #### Tracks layer and function

    def reset( self ):
        """ Reset Tracks layer and data """
        self.graph = None
        self.track_data = None
        ut.remove_layer( self.viewer, "Tracks" )

    def init_tracks(self):
        """ Add a track layer with the new tracks """
        track_table, track_prop = self.create_tracks()
        
        ## plot tracks
        if len(track_table) > 0:
            self.viewer.add_tracks(
                track_table,
                graph=self.graph, 
                name=self.tracklayer_name,
                properties = track_prop,)
            self.viewer.layers[self.tracklayer_name].visible=True
            self.viewer.layers[self.tracklayer_name].color_by="track_id"
            ut.set_active_layer(self.viewer, "Segmentation")
            self.tracklayer = self.viewer.layers[self.tracklayer_name]
            self.track_data = self.tracklayer.data
            #self.track.display_id = True
            self.color_tracks_as_labels()

    def color_tracks_as_labels(self):
        """ Color the tracks the same as the label layer """
        ## must color it manually by getting the Label layer colors for each track_id
        cols = np.zeros((len(self.tracklayer.data[:,0]),4))
        for i, tr in enumerate(self.tracklayer.data[:,0]):
            cols[i] = (self.epicure.seglayer.get_color(tr))
        self.tracklayer._track_colors = cols
        self.tracklayer.events.color_by()
    
    def color_tracks_by_lineage(self):
        """ Color the tracks by their lineage (daughters same colors as parents) """
        ## must color it manually by getting the Label layer colors for each track_id
        cols = np.zeros((len(self.tracklayer.data[:,0]),4))
        for i, tr in enumerate(self.tracklayer.data[:,0]):
            ## find the parent cell,n going up the tree until no more parent
            while tr in self.graph.keys():
                tr = self.graph[tr][0]
            cols[i] = (self.epicure.seglayer.get_color(tr))
        self.tracklayer._track_colors = cols
        self.tracklayer.events.color_by()

    def replace_tracks(self, track_df):
        """ Replace all tracks based on the dataframe """
        track_table, track_prop = self.build_tracks(track_df)
        self.tracklayer.data = track_table
        self.track_data = self.tracklayer.data
        self.tracklayer.properties = track_prop
        self.tracklayer.refresh()
        self.color_tracks_as_labels()

    def reset_tracks(self):
        """ Reset tracks and reload them from labels """
        ut.remove_layer(self.viewer, "Tracks")
        self.init_tracks()

    def update_track_layer(self):
        """ Update the track layer (slow) """
        self.viewer.window._status_bar._toggle_activity_dock(True)
        progress_bar = progress(total=1)
        progress_bar.set_description( "Updating track layer" )
        self.tracklayer.data = self.track_data
        progress_bar.close()
        self.color_tracks_as_labels()
        self.viewer.window._status_bar._toggle_activity_dock(False)

    def measure_track_features(self, track_id):
        """ Measure features (length, total displacement...) of given track """
        features = {}
        track = self.get_track_data( track_id )
        start = int(np.min(track[:,1]))
        end = int(np.max(track[:,1]))
        features["Label"] = track_id
        features["TrackDuration"] = end - start + 1
        features["TrackStart"] = start
        features["TrackEnd"] = end
        if (end-start) == 0:
            ## only one frame
            features["TotalDisplacement"] = None
            features["NetDisplacement"] = None
            features["Straightness"] = None
        else:
            features["TotalDisplacement"] = ut.total_distance( track[:,2:4] )
            features["NetDisplacement"] = ut.net_distance( track[:,2:4] )
            if features["TotalDisplacement"] > 0:
                features["Straightness"] = features["NetDisplacement"]/features["TotalDisplacement"]
            else:
                features["Straightness"] = None
        return features

    def measure_features( self, track_id, features ):
        """ Measure features along all the track """
        mask = self.epicure.get_mask( track_id )
        res = {}
        for feat in features:
            res[feat] = []
        for frame in mask:
            props = regionprops( frame )
            if len(props) > 0:
                if "Area" in features:
                    res["Area"].append( props[0].area )
                if "Hull" in features:
                    res["Hull"].append( props[0].area_convex )
                if "Elongation" in features:
                    res["Elongation"].append( props[0].axis_major_length )
                if "Eccentricity" in features:
                    res["Eccentricity"].append( props[0].eccentricity )
                if "Perimeter" in features:
                    res["Perimeter"].append( props[0].perimeter )
                if "Solidity" in features:
                    res["Solidity"].append( props[0].solidity )
        return res

    def measure_specific_feature( self, track_id, featureName ):
        """ Measure some specific feature """
        if featureName == "Similarity":
            import skimage.metrics as imetrics
            movie = self.epicure.get_label_movie( track_id, extend=1.5 )
            sim_scores = []
            for i in range(0, len(movie)-1):
                score = imetrics.normalized_mutual_information( movie[i], movie[i+1] ) 
                sim_scores.append(score)
            return sim_scores

    def measure_labels(self, segimg):
        """ Get the dataframe of the labels in the segmented image """
        resdf = None
        for iframe, frame in progress(enumerate(segimg)):
            frame_table = self.get_one_frame( frame, iframe )
            if resdf is None:
                resdf = pd.DataFrame(frame_table)
            else:
                resdf = pd.concat([resdf, pd.DataFrame(frame_table)])
        return resdf

    def add_track_frame(self, label, frame, centroid, tree=None, group=None):
        """ Add one frame to the track """
        new_frame = np.array([label, frame, centroid[0], centroid[1]])
        self.track_data = np.vstack((self.track_data, new_frame))
            
    def get_track_list(self):
        """ Get list of unique track_ids """
        return np.unique( self.track_data[:,0] )
    
    def get_tracks_on_frame( self, tframe ):
        """ Return list of tracks present on given frame """
        return np.unique( self.track_data[ self.track_data[:,1]==tframe, 0] ) 

    def has_track(self, label):
        """ Test if track label is present """
        return label in self.track_data[:,0]
    
    def has_tracks(self, labels):
        """ Test if track labels are present """
        return np.isin( labels, self.track_data[:,0] )
    
    def nb_points(self):
        """ Number of points in the tracks """
        return self.track_data.shape[0]

    def nb_tracks(self):
        """ Return number of tracks """
        #return self.track._manager.__len__()
        return len(self.get_track_list())

    def gaped_track(self, track_id):
        """ Check if there is a gap (missing frame) in a track """
        indexes = self.get_track_indexes(track_id)
        if len(indexes) <= 0:
            return False
        track_frames = self.track_data[indexes,1]
        return ((np.max(track_frames)-np.min(track_frames)+1) > len(track_frames) )

    def gap_frames(self, track_id):
        """ Returns the frame(s) at which the gap(s) are """
        track_frames = self.get_track_column( track_id, "frame" )
        gaps = []
        if len( track_frames ) > 0:
            min_frame = int( np.min(track_frames) )
            max_frame = int( np.max(track_frames) )
            gaps = np.setdiff1d( np.arange(min_frame+1, max_frame), track_frames ).tolist()
            if len(gaps) > 0:
                gaps.sort()
        return gaps
            
    def check_gap(self, tracks=None, verbose=None):
        """ Check if there is a track with a gap, flag it if yes """
        if tracks is None:
            tracks = self.get_track_list()
        gaped = []
        for track in tracks:
            if self.gaped_track( track ):
                gaped.append(track)
        if verbose is None:
            verbose = self.epicure.verbose
        if verbose > 0 and len(gaped)>0:
            ut.show_warning("Gap in track(s) "+str(gaped)+"\n"
            +"Consider doing sanity_check in Editing onglet to fix it")
        return gaped

    def get_track_indexes(self, track_id):
        """ Get indexes of track_id tracks position in the arrays """
        if isinstance( track_id,  int ):
            return (np.flatnonzero( self.track_data[:,0] == track_id ) )
        return (np.flatnonzero( np.isin( self.track_data[:,0], track_id ) ) )

    def get_track_indexes_from_frame(self, track_id, frame):
        """ Get indexes of track_id tracks position in the arrays from the given frame """
        if type(track_id) == int:
            return (np.argwhere( (self.track_data[:,0] == track_id)*(self.track_data[:,1]>= frame) )).flatten()
        return (np.argwhere( np.isin( self.track_data[:,0], track_id )*(self.track_data[:,1]>= frame) )).flatten()

    def get_index(self, track_id, frame ):
        """ Get index of track_id at given frame """
        if np.isscalar(track_id):
            track_id = [track_id]
        return np.argwhere( (np.isin(self.track_data[:,0], track_id))*(self.track_data[:,1] == frame) )

    def get_small_tracks(self, max_length=1):
        """ Get tracks smaller than the given threshold """
        labels = []
        lengths = []
        positions = []
        for lab in self.get_track_list():
            indexes = self.get_track_indexes(lab)
            length = len(indexes)
            if length <= max_length:
                pos = self.mean_position( indexes, only_first=False )
                labels.append(lab)
                lengths.append(length)
                positions.append(pos)
        return labels, lengths, positions

    def get_track_data(self, track_id):
        """ Return the data of track track_id """
        indexes = self.get_track_indexes( track_id )
        track = self.track_data[indexes,]
        return track
    
    def get_track_column( self, track_id, column ):
        """ Return the chosen column (frame, x, y, label) of track track_id """
        indexes = self.get_track_indexes( track_id )
        if column == "frame":
            i = 1
        if column == "label":
            i = 0
        if column == "pos":
            i = [2,3]
        track = self.track_data[indexes, i]
        return track

    def get_frame_data( self, track_id, ind ):
        """ Get ind-th data of track track_id """
        track = self.get_track_data( track_id )
        return track[ind]

    def get_position( self, track_id, frame ):
        """ Get position of the track at given frame """
        ind = self.get_index( track_id, frame )
        return [int(self.track_data[ind, 2]), int(self.track_data[ind,3])]

    def mean_position(self, indexes, only_first=False):
        """ Mean positions of tracks at indexes """
        if len(indexes) <= 0:
            return None
        track = self.track_data[indexes,]
        ## keep only the first frame of the tracks
        if only_first:
            min_frame = np.min(track[:,1])
            track = track[track[:,1]==min_frame,]
        return ( int(np.mean(track[:,1])), int(np.mean(track[:,2])), int(np.mean(track[:,3])) )

    def get_first_frame(self, track_id):
        """ Returns first frame where track_id is present """
        track = self.get_track_data( track_id )
        if len(track) <= 0:
            return None
        return int( np.min(track[:,1]) )

    
    def get_last_frame(self, track_id):
        """ Returns last frame where track_id is present """
        track = self.get_track_data( track_id )
        return int(np.max(track[:,1]))
    
    def get_extreme_frames(self, track_id):
        """ Returns the first and last frames where track_id is present """
        track = self.get_track_data( track_id )
        if track.shape[0] > 0:
            return (int(np.min(track[:,1])), int(np.max(track[:,1])) )
        return None, None

    def get_mean_position(self, track_id, only_first=False):
        """ Get mean position of the track """
        indexes = self.get_track_indexes( track_id )
        return self.mean_position( indexes, only_first )

    def update_centroid(self, track_id, frame, ind=None, cx=None, cy=None):
        """ Update track at given frame """
        if ind is None:
            ind = self.get_index( track_id, frame )
        if cx is None:
            prop = ut.getPropLabel( self.epicure.seg[frame], track_id )
            self.track_data[ind, 2:4] = prop.centroid[1]
        else:
            self.track_data[ind, 2] = cx
            self.track_data[ind, 3] = cy

    def swap_frame_id(self, tid, otid, frame):
        """ Swap the ids of two tracks at frame """
        ind = int(self.get_index(tid, frame))
        oind = int(self.get_index(otid, frame))
        ## check if one of the label is an extreme of a track and potentially in the graph
        for track_index in [tid, otid]:
            min_frame, max_frame = self.get_extreme_frames( track_index )
            if (min_frame == frame) or (max_frame == frame):
                self.update_graph( track_index, frame )
        self.track_data[[ind, oind],0] = [otid, tid]

    def update_track_on_frame(self, track_ids, frame):
        """ Update (add or modify) tracks at given frame """
        frame_table = regionprops_table( np.where(np.isin(self.epicure.seg[frame], track_ids), self.epicure.seg[frame], 0), properties=self.properties )
        for x, y, tid in zip(frame_table["centroid-0"], frame_table["centroid-1"], frame_table["label"]):
            index = self.get_index(tid, frame)
            if len(index) > 0:
                self.update_centroid( tid, frame, index, int(x), int(y) )
            else:
                cur_cell = np.array( [tid, frame, int(x), int(y)] )
                cur_cell = np.expand_dims(cur_cell, axis=0)
                self.track_data = np.append(self.track_data, cur_cell, axis=0)
    
    def add_one_frame(self, track_ids, frame, refresh=True):
        """ Add one frame from track """
        for tid in track_ids:
            frame_table = regionprops_table(np.uint8(self.epicure.seg[frame]==tid), properties=self.properties)
            cur_cell = np.array( [tid, frame, int(frame_table["centroid-0"]), int(frame_table["centroid-1"])], dtype=np.uint32 )
            cur_cell = np.expand_dims(cur_cell, axis=0)
            self.track_data = np.append(self.track_data, cur_cell, axis=0)

    def remove_one_frame( self, track_id, frame, handle_gaps=False, refresh=True ):
        """ 
        Remove one frame from track(s) 
        """
        inds = self.get_index( track_id, frame )
        if np.isscalar(track_id):
            track_id = [track_id]
        check_for_gaps = False
        for tid in track_id:
            ## removed frame is in the extremity of a track, can be in the graph
            first_frame, last_frame = self.get_extreme_frames( tid )
            if first_frame is None:
                continue
            if (first_frame == frame) or (last_frame == frame):
                self.update_graph( tid, frame )
            else:
                check_for_gaps = True
        self.track_data = np.delete( self.track_data, inds, axis=0 )
        ## gaps might have been created in the tracks, for now doesn't allow it so split the tracks
        if handle_gaps and check_for_gaps:
            gaped = self.check_gap( track_id, verbose=0 )
            if len(gaped) > 0:
                self.epicure.fix_gaps( gaped )
        
    def get_current_value(self, track_id, frame):
        ind = self.get_index( track_id, frame )
        centx = int(self.track_data[ind, 2])
        centy = int(self.track_data[ind, 3])
        return self.epicure.seg[frame, centx, centy]

    def update_graph(self, track_id, frame):
        """ Update graph if deleted label was linked at that frame, assume keys are unique """
        start_time = time.time()
        if self.graph is not None:
            ## handles current node is last of his branch
            parents = self.last_in_graph( track_id, frame )
            if len(parents) > 0:
                current_label = self.get_current_value( track_id, frame )
                if current_label == 0:
                    for parent in parents:
                        del self.graph[parent]
                else:
                    for parent in parents:
                        self.update_child( parent, track_id, current_label )
            ## handles when current track is first frame of a division
            if self.first_in_graph( track_id, frame ):
                current_label = self.get_current_value( track_id, frame )
                if current_label == 0:
                    del self.graph[track_id]
                else:
                    self.update_key( track_id, current_label ) 
        if self.epicure.verbose > 1:
            ut.show_duration( start_time, header="Graph updated in " )

    def update_child(self, parent, prev_key, new_key):
        """ Change the value of a key in the graph """
        vals = self.graph[parent]
        self.graph[parent] = []
        for val in vals:
            if val == prev_key:
                self.graph[parent].append(new_key)
            else:
                self.graph[parent].append(val)

    def update_key(self, prev_key, new_key):
        """ Change the value of a key in the graph """
        self.graph[new_key] = self.graph.pop(prev_key)

    def is_parent( self, cur_id ):
        """ Return if the current id is in the graph (as a parent, so in values) """
        if self.graph is None:
            return False
        return any( cur_id in vals for vals in self.graph.values() )

    def add_division( self, childa, childb, parent ):
        """ Add info of a division to the graph of divisions/merges """
        if self.graph is None:
            self.graph = {}
        self.graph[childa] = parent
        self.graph[childb] = parent

    def remove_division( self, parent ):
        """ Remove a division event from the graph """
        keys = list(self.graph.keys())
        for key in keys:
            vals = self.graph[key]
            if isinstance(vals, list):
                if vals[0] == parent:
                    del self.graph[key]
            else:
                if vals == parent:
                    del self.graph[key]

    def last_in_graph(self, track_id, frame):
        """ Check if given label and frame is the last of a branch, in the graph """
        parents = []
        for key, vals in self.graph.items():
            if not isinstance(vals, list):
                vals = [vals]
            for val in vals:
                if val == track_id:
                    last_frame = self.get_last_frame( val )
                    if last_frame == frame:
                        parents.append(key)
        return parents

    def first_in_graph(self, track_id, frame):
        """ Check if the given label and frame is the first in the branch so the node in the graph """
        if track_id in self.graph.keys():
            first_frame = self.get_first_frame(track_id)
            if first_frame == frame:
                return True
        return False

    def remove_tracks(self, track_ids):
        """ Remove track with given id """
        inds = self.get_track_indexes(track_ids)
        self.track_data = np.delete(self.track_data, inds, axis=0)
        self.remove_ids_from_graph( track_ids )
    
    def remove_ids_from_graph( self, track_ids ):
        """ Remove all ids from the graph """
        track_ids_set = set( track_ids )
        self.graph = {
            key: vals for key, vals in self.graph.items()
            if (key not in track_ids_set) and ( not any( val in track_ids_set for val in (vals if isinstance(vals, list) else [vals])) )
        }
        """
        for tid in track_ids:
            self.graph.pop( tid, None )
        
        keys = self.graph.keys()
        for key in keys:
            kgr = self.graph[ key ]
            if not isinstance( kgr, list ):
                kgr = [kgr]
            for val in kgr:
                if val in track_ids:
                    del self.graph[key]
                    continue
                """
       
    def build_tracks(self, track_df):
        """ Create tracks from dataframe (after tracking) """
        track = track_df[["track_id", "frame", "centroid-0", "centroid-1"]]
        #frame_prop = frame_table[["tree_id", "label", "nframes", "group"]]
        return np.array(track, int), None #dict(frame_prop)

    def create_tracks(self):
        """ Create tracks from labels (without tracking) """
        properties = ['label', 'centroid']
        track_table = np.empty((0,4), int)
        if self.epicure.nframes >= 0:
            for iframe, frame in progress(enumerate(self.epicure.seg)):
                frame_track, frame_prop = self.get_one_frame( frame, iframe )
                track_table = np.vstack((track_table, frame_track))
        
        return track_table, None # track_prop

    def add_track_features(self, labels):
        """ Add features specific to tracks (eg nframes) """
        nframes = np.zeros(len(labels), int)
        if self.epicure.verbose > 2:
            print("REPLACE BY COUNT METHOD")
        for track_id in np.unique(labels):
            cur_track = np.argwhere(labels == track_id)
            nframes[ list(cur_track) ] = len(cur_track)
        return nframes
    
    def get_one_frame(self, seg, frame):
        """ Get the regionprops results into the dataframe """
        frame_table = regionprops_table(seg, properties=self.properties)
        ndata = len(frame_table["label"])
        if frame is not None:
            frame_table["frame"] = np.repeat(frame, ndata)
        frame_table["track_id"] = frame_table["label"]
        #frame_table["tree_id"] = [0]*ndata    ##### @Tochange to read div/merge infos ?
        #frame_table["group"] = [0]*ndata    ##### @Tochange to read div/merge infos ?
        #frame_table["nframes"] = [0]*ndata
        frame_table = pd.DataFrame(frame_table)
        frame_track = frame_table[["track_id", "frame", "centroid-0", "centroid-1"]]
        #frame_prop = frame_table[["tree_id", "label", "nframes", "group"]]
        return np.array(frame_track, int), None #dict(frame_prop)


    ##########################################
    #### Tracking functions

    def changed_start(self, i):
        """ Ensures that end frame > start frame """
        if i > self.end_frame.value():
            self.end_frame.setValue(i+1)

    def changed_end(self, i):
        if i < self.start_frame.value():
            self.start_frame.setValue(i-1)

    def find_parents(self, labels, twoframes):
        """ Find in the first frame the parents of labels from second frame """
        
        if self.track_choice.currentText() == "Laptrack-Centroids":
            return self.laptrack_centroids_twoframes(labels, twoframes)
        
        if self.track_choice.currentText() == "Laptrack-Overlaps":
            return self.laptrack_overlaps_twoframes(labels, twoframes)
        

    def do_tracking(self):
        """ Start the tracking with the selected options """
        if self.frame_range.isChecked():
            start = self.start_frame.value()
            end = self.end_frame.value()
        else:
            start = 0
            end = self.nframes-1
        start_time = time.time()
        self.viewer.window._status_bar._toggle_activity_dock(True)
        self.epicure.inspecting.reset_all_events()
        
        if self.track_choice.currentText() == "Laptrack-Centroids":
            if self.epicure.verbose > 1:
                print("Starting track with Laptrack-Centroids")
            self.laptrack_centroids( start, end )
            self.epicure.tracked = 1
        if self.track_choice.currentText() == "Laptrack-Overlaps":
            if self.epicure.verbose > 1:
                print("Starting track with Laptrack-Centroids")
            self.laptrack_overlaps( start, end )
            self.epicure.tracked = 1
        
        self.epicure.finish_update(contour=2)
        #self.epicure.reset_free_label()
        self.viewer.window._status_bar._toggle_activity_dock(False)
        if self.epicure.verbose > 0:
            ut.show_duration( start_time, header="Tracking done in " )

    def show_trackoptions(self):
        self.gLapCentroids.setVisible(self.track_choice.currentText() == "Laptrack-Centroids")
        self.gLapOverlap.setVisible(self.track_choice.currentText() == "Laptrack-Overlaps")

    def relabel_nonunique_labels(self, track_df):
        """ After tracking, some track can be splitted and get same label, fix that """
        inittids = np.unique(track_df["track_id"])
        labtracks = []
        saved_data = np.copy(self.epicure.seglayer.data)
        mframes = []
        tids = []
        used = np.unique( saved_data )
        all_labels = np.unique(track_df["label"])
        for tid in inittids:
            cdf = track_df[track_df["track_id"]==tid]
            #print(cdf)
            min_frame = np.min( cdf["frame"] )
            #labtrack = int( cdf["label"][cdf["frame"]==min_frame] )
            for lab in np.unique(cdf["label"]):
                labtracks.append(lab)
                mframes.append( min_frame )
                tids.append(tid)
        if len(labtracks) != len(np.unique(labtracks)):
            ## some labels are present several times
            used = used.tolist()
            for lab in all_labels :
                indexes = list(np.where(np.array(labtracks)==lab)[0])
                if len(indexes)>1:
                    minframes = [mframes[indy] for indy in range(len(labtracks)) if labtracks[indy]==lab]
                    indmin = indexes[ np.argmin( minframes ) ]
                    ## for the other(s), change the label
                    newvals = ut.get_free_labels( used, len(indexes) )
                    used = used + newvals
                    for i, ind in enumerate(indexes):
                        if ind != indmin:
                            tid = tids[ind]
                            newval = newvals[i]
                            track_df.loc[ (track_df["track_id"]==tid)  & (track_df["label"]==lab) , "label"] = newval
                            for frame in track_df["frame"][(track_df["track_id"]==tid) & (track_df["label"]==newval)]:
                                mask = (saved_data[frame]==lab)
                                self.epicure.seglayer.data[frame][mask] = newval
        

    def relabel_trackids(self, track_df, splitdf, mergedf):
        """ Change the trackids to take the first label of each track """
        replace_map = dict()
        new_trackids = track_df['track_id'].copy() 
        new_splitdf = splitdf.copy()
        new_mergedf = mergedf.copy()
        for tid in np.unique(track_df['track_id']):
            ctrack_df = track_df[track_df['track_id']==tid]
            newval = int( ctrack_df.loc[ ctrack_df["frame"]==np.min(ctrack_df["frame"]), "label" ].iloc[0] )
            ## have to replace if different
            if tid != newval:
                new_trackids.loc[ track_df['track_id']==tid ] = newval
                if not new_splitdf.empty:
                    new_splitdf.loc[ splitdf["parent_track_id"]==tid, "parent_track_id" ] = newval
                    new_splitdf.loc[ splitdf["child_track_id"]==tid, "child_track_id" ] = newval
                if not new_mergedf.empty:
                    new_mergedf.loc[ mergedf["parent_track_id"]==tid, "parent_track_id" ] = newval
                    new_mergedf.loc[ mergedf["child_track_id"]==tid, "child_track_id" ] = newval
        return new_trackids, new_splitdf, new_mergedf

    def change_labels(self, track_df):
        """ Change the labels at each frame according to tracks """
        frames = track_df["frame"]
        ## change the other ones
        for frame in np.unique(frames):
            self.change_frame_labels(frame, track_df)

    def change_frame_labels(self, frame, track_df):
        """ Change the labels at given frame according to tracks """
        frame_df = track_df[track_df["frame"]==frame]
        #coordinates = frame_df[['centroid-0', 'centroid-1']].astype(int).values
        track_ids = frame_df['track_id'].astype(int).values
        old_labels = frame_df["label"].astype(int).values
        seglayer = np.copy(self.epicure.seglayer.data[frame])
        for i, lab in enumerate(old_labels):
            mask = (seglayer==lab)
            self.epicure.seglayer.data[frame][mask] = track_ids[i]

    def label_to_dataframe( self, labimg, frame ):
        """ from label, get dataframe of centroids with properties """
        df = pd.DataFrame(regionprops_table(labimg, properties=self.region_properties))
        df["frame"] = frame
        return df
    
    def labels_to_centroids(self, start_frame, end_frame, locked=True):
        """ Get centroids of each cell in dataframe """
        regionprops = []
        #if self.epicure.process_parallel:
        #    with Pool(self.epicure.nparallel) as pool:
        #        regionprops = pool.map( self.label_frame_todf, range(start_frame, end_frame+1) )
        #    regionprops_df = pd.concat(regionprops)
        #else:
        for frame in range(start_frame, end_frame+1):
            df = self.label_frame_todf(frame)
            regionprops.append(df)
        regionprops_df = pd.concat(regionprops)
        return regionprops_df
    
    def labels_ready(self, start_frame, end_frame, locked=True):
        """ Get labels of unlocked cells to track """
        res_labels = []
        if self.epicure.process_parallel:
            with Pool(self.epicure.nparallel) as pool:
                labels = pool.map( self.current_label_frame, range(start_frame, end_frame+1) )
            res_labels = np.array(labels)
        else:
            for frame in range(start_frame, end_frame+1):
                labels = self.current_label_frame(frame)
                res_labels.append(labels)
            res_labels = np.array(res_labels)
        return res_labels
    
    def label_frame_todf( self, frame ):
        """ For current frame, get label frame image then dataframe of centroids """
        clabel = self.current_label_frame(frame)
        return self.label_to_dataframe( clabel, frame )
    
    def current_label_frame( self, frame ):
        """ For current frame, get label frame image """
        clabel = None
        #if self.track_only_in_roi.isChecked():
        #    clabel = self.epicure.only_current_roi(frame)
        if clabel is None:
            clabel = self.epicure.seg[frame]
        return clabel

    def after_tracking(self, track_df, split_df, merge_df, progress_bar, indprogress):
        """ Steps after tracking: get/show the graph from the track_df """
        graph = None
        progress_bar.set_description( "Update labels and tracks" )
        ## shift all by 1 so that doesn't start at 0
        if len(split_df) > 0:
            split_df[["parent_track_id"]] += 1
            split_df[["child_track_id"]] += 1
        if len(merge_df) > 0:
            merge_df[["parent_track_id"]] += 1
            merge_df[["child_track_id"]] += 1
        track_df[["track_id"]] += 1
       
        ## relabel if some track have the same label
        self.relabel_nonunique_labels(track_df)
        ## relabel track ids so that they are equal to the first label of the track
        newtids, split_df, merge_df = self.relabel_trackids(track_df, split_df, merge_df)
        track_df["track_id"] = newtids
        self.change_labels(track_df)

        # create graph of division/merging
        self.graph = convert_split_merge_df_to_napari_graph(split_df, merge_df)

        progress_bar.update(indprogress+1)
        
        ## update display if active
        self.replace_tracks(track_df)

        progress_bar.update(indprogress+2)
        ## update the list of events, or others 
        self.epicure.updates_after_tracking()
        progress_bar.update(indprogress+3)
        return graph


    
############ Laptrack centroids option
    
    def create_laptrack_centroids(self):
        """ GUI of the laptrack option """
        self.gLapCentroids = QGroupBox("Laptrack-Centroids")
        glap_layout = QVBoxLayout()
        mdist, self.max_dist = wid.value_line( "Max distance", "15.0", "Maximal distance between two labels in consecutive frames to link them (in pixels)" )
        glap_layout.addLayout(mdist)
        ## splitting ~ cell division
        scost, self.splitting_cost = wid.value_line( "Splitting cutoff", "1", "Weight to split a track in two (increasing it favors division)" )
        glap_layout.addLayout(scost)
        ## merging ~ error ?
        mcost, self.merging_cost = wid.value_line( "Merging cutoff", "1", "Weight to merge to labels together" )
        glap_layout.addLayout(mcost)

        add_feat, self.check_penalties, self.bpenalties = wid.checkgroup_help( "Add features cost", True, "Add cell features in the tracking calculation", None )
        self.create_penalties()
        glap_layout.addWidget(self.check_penalties)
        glap_layout.addWidget(self.bpenalties)
        self.gLapCentroids.setLayout(glap_layout)

    def show_penalties(self):
        self.bpenalties.setVisible(not self.bpenalties.isVisible())

    def create_penalties(self):
        pen_layout = QVBoxLayout()
        areaCost, self.area_cost = wid.value_line( "Area difference", "2", "Weight of the difference of area between two labels to link them (0 to ignore)" )
        pen_layout.addLayout(areaCost)
        solidCost, self.solidity_cost = wid.value_line( "Solidity difference", "0", "Weight of the difference of solidity between two labels to link them (0 to ignore)" )
        pen_layout.addLayout(solidCost)
        self.bpenalties.setLayout(pen_layout)

    def laptrack_centroids_twoframes(self, labels, twoframes):
        """ Perform tracking of two frames with current parameters """
        laptrack = LaptrackCentroids(self, self.epicure)
        laptrack.max_distance = float(self.max_dist.text())
        self.region_properties = ["label", "centroid"]
        if self.check_penalties.isChecked():
            self.region_properties.append("area")
            self.region_properties.append("solidity")
            laptrack.penal_area = float(self.area_cost.text())
            laptrack.penal_solidity = float(self.solidity_cost.text())
        laptrack.set_region_properties(with_extra=self.check_penalties.isChecked())
            
        df = self.twoframes_centroid(twoframes)
        if set(np.unique(df["label"])) == set(labels):
            ## no other labels
            return [None]*len(labels) 
        laptrack.splitting_cost = False ## disable splitting option
        laptrack.merging_cost = False ## disable merging option
        parent_labels = laptrack.twoframes_track(df, labels)
        return parent_labels
    
    def twoframes_centroid(self, img):
        """ Get centroids of first frame only """
        df0 = self.label_to_dataframe( img[0], 0 )
        df1 = self.label_to_dataframe( img[1], 1 )
        return pd.concat([df0, df1])
    
    def laptrack_centroids(self, start, end):
        """ Perform track with laptrack option and chosen parameters """
        ## Laptrack tracker
        laptrack = LaptrackCentroids(self, self.epicure)
        laptrack.max_distance = float(self.max_dist.text())
        laptrack.splitting_cost = float(self.splitting_cost.text())
        laptrack.merging_cost = float(self.merging_cost.text())
        self.region_properties = ["label", "centroid"]
        if self.check_penalties.isChecked():
            self.region_properties.append("area")
            self.region_properties.append("solidity")
            laptrack.penal_area = float(self.area_cost.text())
            laptrack.penal_solidity = float(self.solidity_cost.text())
        laptrack.set_region_properties(with_extra=self.check_penalties.isChecked())

        progress_bar = progress(total=7)
        progress_bar.set_description( "Prepare tracking" )
        if self.epicure.verbose > 1:
            print("Convert labels to centroids: use track info ?")
        df = self.labels_to_centroids(start, end, locked=True)
        progress_bar.update(1)
        if self.epicure.verbose > 1:
            print("GO tracking")
        progress_bar.set_description( "Do tracking with LapTrack Centroids" )
        track_df, split_df, merge_df = laptrack.track_centroids(df)
        progress_bar.update(2)
        if self.epicure.verbose > 1:
            print("After tracking, update everything")
        self.after_tracking(track_df, split_df, merge_df, progress_bar, 2)
        progress_bar.update(6)
        progress_bar.close()
    
############ Laptrack overlap option

    def create_laptrack_overlap(self):
        """ GUI of the laptrack overlap option """
        self.gLapOverlap = QGroupBox("Laptrack-Overlaps")
        glap_layout = QVBoxLayout()
        miou, self.min_iou = wid.value_line( "Min IOU", "0.1", "Minimum Intersection Over Union score to link to labels together" )
        glap_layout.addLayout(miou)
        
        scost, self.split_cost = wid.value_line( "Splitting cost", "0.2", "Weight of linking a parent label with two labels (increasing it for more divisions)" )
        glap_layout.addLayout(scost)
        
        mcost, self.merg_cost = wid.value_line( "Merging cost", "0.2", "Weight of merging two parent labels into one" )
        glap_layout.addLayout(mcost)

        self.gLapOverlap.setLayout(glap_layout)

    def laptrack_overlaps(self, start, end):
        """ Perform track with laptrack overlap option and chosen parameters """
        ## Laptrack tracker
        laptrack = LaptrackOverlaps(self, self.epicure)
        miniou = float(self.min_iou.text())
        if miniou >= 1.0:
            miniou = 1.0
        laptrack.cost_cutoff = 1.0 - miniou
        laptrack.splitting_cost = float(self.split_cost.text())
        laptrack.merging_cost = float(self.merg_cost.text())
        self.region_properties = ["label", "centroid"]

        progress_bar = progress(total=6)
        progress_bar.set_description( "Prepare tracking" )
        labels = self.labels_ready(start, end, locked=True)
        progress_bar.update(1)
        progress_bar.set_description( "Do tracking with LapTrack Overlaps" )
        track_df, split_df, merge_df = laptrack.track_overlaps(labels)
        progress_bar.update(2)
        
        ## get dataframe of coordinates to create the graph 
        df = self.labels_to_centroids(start, end, locked=True)
        progress_bar.update(3)
        coordinate_df = df.set_index(["frame", "label"])
        tdf = track_df.set_index(["frame", "label"])
        track_df2 = pd.merge( tdf, coordinate_df, right_index=True, left_index=True).reset_index()
        self.after_tracking(track_df2, split_df, merge_df, progress_bar, 3)
        progress_bar.update(6)
        progress_bar.close()
    
    def laptrack_overlaps_twoframes(self, labels, twoframes):
        """ Perform tracking of two frames with current parameters """
        laptrack = LaptrackOverlaps(self, self.epicure)
        miniou = float(self.min_iou.text())
        if miniou >= 1.0:
            miniou = 1.0
        laptrack.cost_cutoff = 1.0 - miniou
        self.region_properties = ["label", "centroid"]

        laptrack.splitting_cost = False ## disable splitting option
        laptrack.merging_cost = False ## disable merging option
        parent_labels = laptrack.twoframes_track(twoframes, labels)
        return parent_labels


