import tkinter as tk
from idlelib.tooltip import Hovertip
from tkinter import filedialog, colorchooser
import tkinter.messagebox as messagebox
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
import matplotlib.patches as mpatches
import pandas as pd
import cv2
import numpy as np
import glob
import os
from Trackrefiner import identify_important_columns, calculate_bac_endpoints, extract_bacteria_info, \
    find_bacteria_neighbors, remove_links, create_links, process_bacterial_life_and_family, create_pickle_files, \
    propagate_bacteria_labels
import logging

logging.getLogger('matplotlib').setLevel(logging.WARNING)
logging.getLogger("PIL").setLevel(logging.WARNING)


class TrackingGUI(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("Bacterial Tracking Visualizer and Editor")

        self.geometry("400x400")

        self.frame = tk.Frame(self)
        self.frame.pack(expand=True)

        # CSV file selection
        self.open_file_frame = tk.Frame(self.frame)
        self.open_file_frame.pack(side='top', pady=(0, 0))

        self.btn_open_file = tk.Button(self.open_file_frame, text="Trackrefiner Refined Tracking CSV File",
                                       command=self.load_csv_file)
        self.btn_open_file.pack(side='top', pady=(0, 0))
        # Adding tooltip
        Hovertip(self.btn_open_file, "Path to CSV file generated by Trackrefiner, containing "
                                     "the results of refined bacterial tracking.", hover_delay=500)

        # CSV neighbor file selection
        self.open_neighbor_file_frame = tk.Frame(self.frame)
        self.open_neighbor_file_frame.pack(side='top', pady=(10, 0))

        self.btn_open_neighbor_file = tk.Button(self.open_neighbor_file_frame,
                                                text="Neighbor CSV file",
                                                command=self.load_neighbor_csv_file)
        self.btn_open_neighbor_file.pack(side='top', pady=(0, 0))
        # Adding tooltip
        Hovertip(self.btn_open_neighbor_file, "Path to a CSV file containing bacterial neighbor data, "
                                              "exported by Trackrefiner", hover_delay=500)

        # Raw image folder selection
        self.open_images_frame = tk.Frame(self.frame)
        self.open_images_frame.pack(side='top', pady=(10, 10))

        self.btn_open_images = tk.Button(self.open_images_frame, text="Bacteria Raw Images Directory",
                                         command=self.load_images_dir)
        self.btn_open_images.pack(side='top', pady=(5, 0))
        # Adding tooltip
        Hovertip(self.btn_open_images, "Path to the directory containing raw bacterial microscopic images "
                                       "in TIFF, PNG, or JPG format.", hover_delay=500)

        # umPerPixel entry
        self.um_per_pixel_frame = tk.Frame(self.frame)
        self.um_per_pixel_frame.pack(side='top', pady=(10, 10))

        self.label_pixel_per_micron = tk.Label(self.um_per_pixel_frame, text="Pixel Per Micron:")
        self.label_pixel_per_micron.pack(side='left', pady=(0, 0))
        self.entry_pixel_per_micron = tk.Entry(self.um_per_pixel_frame)
        self.entry_pixel_per_micron.insert(0, "0.144")  # Default value of 0.144
        self.entry_pixel_per_micron.pack(side='left', pady=(0, 0))
        # Adding tooltip
        Hovertip(self.entry_pixel_per_micron, "Conversion factor for pixels to micrometers.", hover_delay=500)

        # Font size entry
        self.font_size_frame = tk.Frame(self.frame)
        self.font_size_frame.pack(side='top', pady=(10, 10))

        self.label_font_size = tk.Label(self.font_size_frame, text="Bacteria Info Font Size:")
        self.label_font_size.pack(side='left', pady=(0, 0))
        self.entry_font_size = tk.Entry(self.font_size_frame)
        self.entry_font_size.insert(0, "12")  # Default value of 12
        self.entry_font_size.pack(side='left', pady=(0, 0))

        # Object color picker
        self.label_color_frame = tk.Frame(self.frame)
        self.label_color_frame.pack(side='top', pady=(10, 0))

        self.label_color = tk.Label(self.label_color_frame, text="Pick Bacteria Color:")
        self.label_color.pack(side='left', pady=(0, 0))
        self.color_button = tk.Button(self.label_color_frame, text="Select Color", command=self.select_color)
        self.color_button.pack(side='left', pady=(0, 0))

        # Show all time steps button
        self.all_button_frame = tk.Frame(self.frame)
        self.all_button_frame.pack(side='top', pady=(10, 0))
        self.show_all_button = tk.Button(self.all_button_frame, text="Show All Time Steps",
                                         command=self.show_all_time_steps)
        self.show_all_button.pack(side='top', pady=(10, 0))

        # Button to show two specific time steps
        self.two_steps_button_frame = tk.Frame(self.frame)
        self.two_steps_button_frame.pack(side='top', pady=(10, 0))
        self.show_two_steps_button = tk.Button(self.two_steps_button_frame, text="Show Two Specific Time Steps",
                                               command=self.show_two_steps)
        self.show_two_steps_button.pack(side='top', pady=(10, 0))

        # Button to correct tracking
        self.refine_bac_tracking_button_frame = tk.Frame(self.frame)
        self.refine_bac_tracking_button_frame.pack(side='top', pady=(10, 0))
        self.refine_bac_tracking_button = tk.Button(self.refine_bac_tracking_button_frame,
                                                    text="Refine Bacterial Tracking", command=self.refine_tracking)
        self.refine_bac_tracking_button.pack(side='top', pady=(10, 0))

        # Variables for file paths
        self.csv_file = None
        self.csv_neighbor_file = None
        self.raw_images_dir = None

        # dataframes
        self.bacteria_df = None
        self.df_neighbor = pd.DataFrame()
        self.raw_images = []

        self.center_coord_cols = None
        self.all_rel_center_coord_cols = None
        self.parent_image_number_col = None
        self.parent_object_number_col = None
        self.label_col = None
        self.original_cols = []
        self.neighbor_list_array = None
        self.change_happened = False
        self.out_dir = None

        # Conversion factors and settings
        self.um_per_pixel = 0.144
        self.object_color = "#56e64e"
        self.font_size = 12

        # List to track opened windows
        self.open_windows = []

        self.protocol("WM_DELETE_WINDOW", self.on_main_window_close)

    def load_csv_file(self):
        self.csv_file = filedialog.askopenfilename(filetypes=[("CSV files", "*.csv")],
                                                   title="Select CSV file output of Trackrefiner")
        print(f"Loaded CSV file: {self.csv_file}")

        try:
            self.bacteria_df = pd.read_csv(self.csv_file)
            (self.center_coord_cols, self.all_rel_center_coord_cols, self.parent_image_number_col,
             self.parent_object_number_col, self.label_col) = identify_important_columns(self.bacteria_df)

            self.original_cols = self.bacteria_df.columns.values
            self.bacteria_df = self.bacteria_df.rename(columns={'cellAge': 'age'})

            self.bacteria_df['index'] = self.bacteria_df.index

            self.bacteria_df['prev_time_step_index'] = self.bacteria_df.groupby('id')["index"].shift(1)

            # now for daughters
            merged_df = \
                self.bacteria_df[
                    ['ImageNumber', 'ObjectNumber', self.parent_image_number_col,
                     self.parent_object_number_col, 'index']].merge(
                    self.bacteria_df[['ImageNumber', 'ObjectNumber', 'index']],
                    left_on=[self.parent_image_number_col, self.parent_object_number_col],
                    right_on=['ImageNumber', 'ObjectNumber'], how='inner', suffixes=('_2', '_1'))

            divisions = \
                merged_df[merged_df.duplicated(subset='index_1', keep=False)][['index_1', 'index_2',
                                                                               'ImageNumber_1', 'ObjectNumber_1',
                                                                               'ImageNumber_2', 'ObjectNumber_2',
                                                                               f"{self.parent_image_number_col}",
                                                                               f'{self.parent_object_number_col}'
                                                                               ]]

            daughters_idx = divisions['index_2'].values
            self.bacteria_df.loc[daughters_idx, 'prev_time_step_index'] = divisions['index_1'].values

            self.bacteria_df['prev_time_step_index'] = self.bacteria_df['prev_time_step_index'].fillna(-1)
            self.bacteria_df['prev_time_step_index'] = self.bacteria_df['prev_time_step_index'].astype('Int64')

            self.bacteria_df['checked'] = True

            self.out_dir = os.path.dirname(self.csv_file)

            if self.df_neighbor is not None and self.bacteria_df is not None and self.neighbor_list_array is None:
                if self.df_neighbor.shape[0] > 0 and self.bacteria_df.shape[0] > 0:
                    self.neighbor_list_array = find_bacteria_neighbors(self.bacteria_df, self.df_neighbor)

        except (FileNotFoundError, pd.errors.EmptyDataError, pd.errors.ParserError) as e:
            messagebox.showerror("Error", "Error loading CSV file containing measured bacterial features.")
            return

    def load_neighbor_csv_file(self):
        self.csv_neighbor_file = filedialog.askopenfilename(filetypes=[("CSV files", "*.csv")],
                                                            title="Select CSV file containing neighboring data of"
                                                                  " bacteria")
        print(f"Loaded CSV file containing neighboring data of bacteria: {self.csv_neighbor_file}")

        try:
            self.df_neighbor = pd.read_csv(self.csv_neighbor_file)

            if 'Relationship' in self.df_neighbor.columns:
                self.df_neighbor = self.df_neighbor.loc[self.df_neighbor['Relationship'] == 'Neighbors']

            if self.df_neighbor is not None and self.bacteria_df is not None and self.neighbor_list_array is None:
                if self.df_neighbor.shape[0] > 0 and self.bacteria_df.shape[0] > 0:
                    self.neighbor_list_array = find_bacteria_neighbors(self.bacteria_df, self.df_neighbor)

        except (FileNotFoundError, pd.errors.EmptyDataError, KeyError) as e:

            messagebox.showerror("Error", "No CSV file found containing neighboring data of bacteria.")
            return

    def load_images_dir(self):
        self.raw_images_dir = filedialog.askdirectory(title="Select Raw Images Directory")
        print(f"Loaded Raw Images Directory: {self.raw_images_dir}")

        extensions = ('tif', 'tiff', 'jpg', 'png')
        images_found = False
        for ext in extensions:
            if not images_found:
                self.raw_images = sorted(glob.glob(self.raw_images_dir + '/*.' + ext))
                if len(self.raw_images) > 0:
                    images_found = True

        if not images_found:
            messagebox.showerror("No Image Directory", "Please select a raw images directory.")
            return

    def select_color(self):
        color_code = colorchooser.askcolor(title="Choose Object Color")
        if color_code:
            self.object_color = color_code[1]  # Get the hexadecimal color code
            print(f"Selected color: {self.object_color}")

    def update_font_size(self):
        try:
            self.font_size = int(self.entry_font_size.get())  # Get font size from input and convert to integer
        except ValueError:
            self.font_size = 12  # Set a default value if input is invalid

    def update_um_per_pixel(self):
        try:
            self.um_per_pixel = float(
                self.entry_pixel_per_micron.get())  # Get font size from input and convert to float
        except ValueError:
            self.um_per_pixel = 0.144  # Set a default value if input is invalid

    def refine_tracking(self):
        # Open a new window to specify time step 1, time step 2, ID, parent ID, and neighbor distance
        new_window = tk.Toplevel(self)
        new_window.title("Bacterial Tracking Refinement")
        new_window.geometry("400x400")

        # Time step entry
        t_frame = tk.Frame(new_window)
        t_frame.pack(side='top', pady=(10, 0))
        target_bac_label_time_steps = tk.Label(t_frame, text="Target Bacteria Time Step(s) (comma separated):")
        target_bac_label_time_steps.pack(side='left', pady=(0, 0))
        target_bac_entry_time_step = tk.Entry(t_frame)
        target_bac_entry_time_step.pack(side='left', pady=(0, 0))
        # Adding tooltip
        Hovertip(target_bac_entry_time_step,
                 "List of target bacteria Time Steps, separated by `,`. "
                 "Fill it if you want to create or remove a link.", hover_delay=500)

        # ID entry
        id_frame = tk.Frame(new_window)
        id_frame.pack(side='top', pady=(20, 0))
        label_id = tk.Label(id_frame, text="Target Bacteria ID(s) (comma separated):")
        label_id.pack(side='left', pady=(0, 0))
        entry_id = tk.Entry(id_frame)
        entry_id.pack(side='left', pady=(0, 0))
        # Adding tooltip
        Hovertip(entry_id,
                 "List of target bacteria IDs, separated by `,`. Fill it if you want to create or remove a link.",
                 hover_delay=500)

        # time steps of source bacteria
        source_time_steps_frame = tk.Frame(new_window)
        source_time_steps_frame.pack(side='top', pady=(20, 0))
        source_bac_label_time_steps = tk.Label(source_time_steps_frame,
                                               text="Source Bacteria Time Step(s) (comma separated):")
        source_bac_label_time_steps.pack(side='left', pady=(0, 0))
        source_bac_entry_time_step = tk.Entry(source_time_steps_frame)
        source_bac_entry_time_step.pack(side='left', pady=(0, 0))
        # Adding tooltip
        Hovertip(source_bac_entry_time_step,
                 "List of source bacteria Time Steps, separated by `,`. Fill it if you want to create a link.",
                 hover_delay=500)

        # Parent ID entry
        parent_id_frame = tk.Frame(new_window)
        parent_id_frame.pack(side='top', pady=(20, 0))
        label_parent_id = tk.Label(parent_id_frame, text="Source Bacteria ID(s) (comma separated):")
        label_parent_id.pack(side='left', pady=(0, 0))
        entry_parent_id = tk.Entry(parent_id_frame)
        entry_parent_id.pack(side='left', pady=(0, 0))
        # Adding tooltip
        Hovertip(entry_parent_id,
                 "List of source bacteria IDs, separated by `,`. Fill it if you want to create a link.",
                 hover_delay=500)

        # Interval Time Entry
        interval_time_frame = tk.Frame(new_window)
        interval_time_frame.pack(side='top', pady=(20, 0))
        interval_time_label = tk.Label(interval_time_frame, text="Interval Time:")
        interval_time_label.pack(side='left', pady=(0, 0))
        entry_interval_time = tk.Entry(interval_time_frame)
        entry_interval_time.pack(side='left', pady=(0, 0))
        # Adding tooltip
        Hovertip(entry_interval_time, "Time interval between frames in minutes", hover_delay=500)

        # Elongation Rate Method Dropdown
        elongation_method_frame = tk.Frame(new_window)
        elongation_method_frame.pack(side='top', pady=(20, 0))
        elongation_method_label = tk.Label(elongation_method_frame, text="Elongation Rate Method:")
        elongation_method_label.pack(side='left', pady=(0, 0))

        elongation_methods = ["Average", "Linear Regression"]
        elongation_method_var = tk.StringVar(value="Average")
        elongation_method_dropdown = tk.OptionMenu(elongation_method_frame, elongation_method_var, *elongation_methods)
        elongation_method_dropdown.pack(side='left', pady=(0, 0))
        # Adding tooltip
        Hovertip(elongation_method_dropdown, "Method to calculate elongation rate: Average or Linear Regression",
                 hover_delay=500)

        # Save Pickle Checkbox
        save_pickle_frame = tk.Frame(new_window)
        save_pickle_frame.pack(side='top', pady=(20, 0))
        save_pickle_var = tk.BooleanVar(value=False)
        save_pickle_checkbox = tk.Checkbutton(save_pickle_frame, text="Save Pickle", variable=save_pickle_var)
        save_pickle_checkbox.pack(side='left', pady=(0, 0))
        # Adding tooltip
        Hovertip(save_pickle_checkbox, "Save results in .pickle format", hover_delay=500)

        # Buttons for "Remove Links" and "Create Links" on the same row
        action_frame = tk.Frame(new_window)
        action_frame.pack(side='top', pady=(20, 0))

        remove_button = tk.Button(action_frame, text="Remove Links",
                                  command=lambda: self.run_remove_links(target_bac_entry_time_step.get(),
                                                                        entry_id.get(), new_window))
        remove_button.pack(side='left', padx=5)

        create_button = tk.Button(action_frame, text="Create Links",
                                  command=lambda: self.run_create_links(target_bac_entry_time_step.get(),
                                                                        entry_id.get(), entry_parent_id.get(),
                                                                        source_bac_entry_time_step.get(), new_window))
        create_button.pack(side='left', padx=5)

        # Save button on a separate row
        save_frame = tk.Frame(new_window)
        save_frame.pack(side='top', pady=(20, 0))
        save_button = tk.Button(save_frame, text="Save Refinement",
                                command=lambda: self.modified_tracking_links(entry_interval_time.get(),
                                                                             elongation_method_var.get(),
                                                                             save_pickle_var.get(), new_window))
        save_button.pack()

    def show_two_steps(self):

        self.update_font_size()  # Update font size before performing actions
        self.update_um_per_pixel()

        # Open a new window to specify time step 1, time step 2, ID, parent ID, and neighbor distance
        new_window = tk.Toplevel(self)
        new_window.title("Trackrefiner Specify Time Steps and Settings")
        new_window.geometry("420x320")

        # Time step 1 entry
        t1_frame = tk.Frame(new_window)
        t1_frame.pack(side='top', pady=(10, 0))

        label_t1 = tk.Label(t1_frame, text="Time Step 1:")
        label_t1.pack(side='left', pady=(0, 0))
        entry_t1 = tk.Entry(t1_frame)
        entry_t1.pack(side='left', pady=(0, 0))
        # Adding tooltip
        Hovertip(entry_t1, "Enter the time step where the source objects are present.", hover_delay=500)

        # Time step 2 entry
        t2_frame = tk.Frame(new_window)
        t2_frame.pack(side='top', pady=(10, 0))
        label_t2 = tk.Label(t2_frame, text="Time Step 2:")
        label_t2.pack(side='left', pady=(0, 0))
        entry_t2 = tk.Entry(t2_frame)
        entry_t2.pack(side='left', pady=(0, 0))
        Hovertip(entry_t2, "Enter the time step where the target objects appear.", hover_delay=500)

        # Object number entry
        obj_num_frame = tk.Frame(new_window)
        obj_num_frame.pack(side='top', pady=(10, 0))
        label_obj_num = tk.Label(obj_num_frame, text="Bacteria Objects Number(s) (comma separated):")
        label_obj_num.pack(side='left', pady=(0, 0))
        entry_obj_num = tk.Entry(obj_num_frame)
        entry_obj_num.pack(side='left', pady=(0, 0))
        Hovertip(entry_obj_num, "Enter one or multiple object numbers of bacteria, separated by commas.",
                 hover_delay=500)

        # ID entry
        id_frame = tk.Frame(new_window)
        id_frame.pack(side='top', pady=(10, 0))
        label_id = tk.Label(id_frame, text="Bacteria ID(s) (comma separated):")
        label_id.pack(side='left', pady=(0, 0))
        entry_id = tk.Entry(id_frame)
        entry_id.pack(side='left', pady=(0, 0))
        Hovertip(entry_id, "Enter one or multiple object IDs, separated by commas.", hover_delay=500)

        # Parent ID entry
        parent_id_frame = tk.Frame(new_window)
        parent_id_frame.pack(side='top', pady=(10, 0))
        label_parent_id = tk.Label(parent_id_frame, text="Bacteria Parent ID(s) (comma separated):")
        label_parent_id.pack(side='left', pady=(0, 0))
        entry_parent_id = tk.Entry(parent_id_frame)
        entry_parent_id.pack(side='left', pady=(0, 0))
        Hovertip(entry_parent_id, "Enter parent bacteria IDs, separated by commas.", hover_delay=500)

        # Neighbor distance entry
        distance_frame = tk.Frame(new_window)
        distance_frame.pack(side='top', pady=(10, 0))
        label_distance = tk.Label(distance_frame, text="Neighbor Distance (unit: um):")
        label_distance.pack(side='left', pady=(0, 0))
        entry_distance = tk.Entry(distance_frame)
        entry_distance.pack(side='left', pady=(0, 0))
        Hovertip(entry_distance, "Specify the maximum distance (in micrometers) to define bacterial neighbors.",
                 hover_delay=500)

        # Dropdown menu for selecting the visualization option
        vis_mode_frame = tk.Frame(new_window)
        vis_mode_frame.pack(side='top', pady=(10, 0))
        label_option = tk.Label(vis_mode_frame, text="Visualization Mode:")
        label_option.pack(side='left', pady=(0, 0))

        selected_option = tk.StringVar(vis_mode_frame)
        selected_option.set("Two Different Slides")  # Default value

        option_menu = tk.OptionMenu(vis_mode_frame, selected_option, "Two Different Slides", "Slide Show")
        option_menu.pack(side='left', pady=(0, 0))

        # Submit button to handle input
        submit_frame = tk.Frame(new_window)
        submit_frame.pack(side='top', pady=(10, 0))
        submit_button = tk.Button(submit_frame, text="Show",
                                  command=lambda: self.handle_two_steps(entry_t1.get(), entry_t2.get(),
                                                                        entry_obj_num.get(), entry_id.get(),
                                                                        entry_parent_id.get(),
                                                                        entry_distance.get(), selected_option.get(),
                                                                        new_window))
        submit_button.pack()

    def run_create_links(self, target_bac_time_steps, target_bac_ids, source_bac_ids, source_bac_time_steps, window):

        print(f"Target objects Time Step: {target_bac_time_steps}, Target Objects IDs: {target_bac_ids}, "
              f"Source Objects Time steps: {source_bac_time_steps}, Source Objects IDs: {source_bac_ids}")

        target_bac_time_steps = target_bac_time_steps.replace(" ", "").split(',') if target_bac_time_steps else None
        target_bac_ids = target_bac_ids.replace(" ", "").split(',') if target_bac_ids else None

        source_bac_time_steps = source_bac_time_steps.replace(" ", "").split(',') if source_bac_time_steps else None
        source_bac_ids = source_bac_ids.replace(" ", "").split(',') if source_bac_ids else None

        if all(v is not None for v in (source_bac_ids, source_bac_time_steps, target_bac_ids, target_bac_time_steps)):

            # Convert to list of integers
            source_bac_ids = [int(x) for x in source_bac_ids if x.isdigit()]
            source_bac_time_steps = [int(x) for x in source_bac_time_steps if x.isdigit()]

            target_bac_ids = [int(x) for x in target_bac_ids if x.isdigit()]
            target_bac_time_steps = [int(x) for x in target_bac_time_steps if x.isdigit()]

            cond1 = len(source_bac_time_steps) == len(target_bac_time_steps)
            cond2 = len(source_bac_time_steps) == len(target_bac_time_steps) == 1

            run_processing = False

            if cond1:
                if cond2:
                    if len(source_bac_ids) == 1 and len(target_bac_ids) >= 1:
                        run_processing = True

                    elif len(source_bac_ids) == len(target_bac_ids):
                        run_processing = True
                else:
                    if len(source_bac_ids) == len(target_bac_ids) == len(source_bac_time_steps) == len(
                            target_bac_time_steps):
                        run_processing = True

            if run_processing:
                self.bacteria_df = create_links(self.bacteria_df, self.df_neighbor, self.neighbor_list_array,
                                                self.parent_image_number_col, self.parent_object_number_col,
                                                self.center_coord_cols, source_bac_ids, source_bac_time_steps,
                                                target_bac_ids, target_bac_time_steps)

                self.change_happened = True

                messagebox.showinfo("Success", "Links created successfully!")

                self.remove_specific_window(['Trackrefiner Specify Time Steps and Settings',
                                             'Trackrefiner Show All Time Steps', 'Trackrefiner Slider Control',
                                             'Time Step'])

    def run_remove_links(self, incorrect_target_bac_time_steps, incorrect_target_bac_ids, window):

        print(f"Time Step: {incorrect_target_bac_time_steps}, Target Objects IDs: {incorrect_target_bac_ids}")

        incorrect_target_bac_time_steps = \
            incorrect_target_bac_time_steps.replace(" ", "").split(',') if incorrect_target_bac_time_steps else None

        incorrect_target_bac_ids = \
            incorrect_target_bac_ids.replace(" ", "").split(',') if incorrect_target_bac_ids else None

        if all(v is not None for v in (incorrect_target_bac_ids, incorrect_target_bac_time_steps)):

            # Convert to list of integers
            incorrect_target_bac_time_steps = [int(x) for x in incorrect_target_bac_time_steps if x.isdigit()]
            incorrect_target_bac_ids = [int(x) for x in incorrect_target_bac_ids if x.isdigit()]

            if ((len(incorrect_target_bac_ids) == len(incorrect_target_bac_time_steps)) or
                    len(incorrect_target_bac_time_steps) == 1):
                self.bacteria_df = remove_links(self.bacteria_df, self.df_neighbor, self.neighbor_list_array,
                                                self.parent_image_number_col, self.parent_object_number_col,
                                                self.center_coord_cols, incorrect_target_bac_ids,
                                                incorrect_target_bac_time_steps)

                self.change_happened = True

                messagebox.showinfo("Success", "Links removed successfully!")
                self.remove_specific_window(['Trackrefiner Specify Time Steps and Settings',
                                             'Trackrefiner Show All Time Steps', 'Trackrefiner Slider Control',
                                             'Time Step'])

    def modified_tracking_links(self, interval_time, elongation_rate_method, save_pickle, window):

        if self.change_happened:

            interval_time = int(interval_time)

            assigning_cell_type = False
            cell_type_array = np.array([])

            # label correction
            trackrefiner_out_df = propagate_bacteria_labels(self.bacteria_df, self.parent_image_number_col,
                                                            self.parent_object_number_col, self.label_col)

            # process the tracking data
            processed_df, processed_df_with_specific_cols = \
                process_bacterial_life_and_family(trackrefiner_out_df, interval_time, elongation_rate_method,
                                                  assigning_cell_type, cell_type_array, self.label_col,
                                                  self.center_coord_cols)

            # final step
            corrected_df = processed_df[self.original_cols]

            if save_pickle:

                folder_name = f'{self.out_dir}/pickle_files'
                while os.path.exists(folder_name):
                    folder_name += "#"

                create_pickle_files(processed_df_with_specific_cols, f'{folder_name}/',
                                    assigning_cell_type)

            trackrefiner_out_file_name = f'{self.csv_file}'
            while os.path.exists(trackrefiner_out_file_name):
                trackrefiner_out_file_name = trackrefiner_out_file_name.rsplit('.csv', 1)[0] + "#.csv"

            corrected_df.to_csv(trackrefiner_out_file_name, index=False)

            messagebox.showinfo("Success", f"saved successfully in {self.out_dir} directory!")

    def handle_two_steps(self, t1, t2, object_numbers, ids, parent_ids, distance, mode, window):

        # Logic to handle the specified input and show the visualizations
        print(f"Time Step 1: {t1}, Time Step 2: {t2}, Object Numbers: {object_numbers}, IDs: {ids}, "
              f"Parent IDs: {parent_ids}, Distance: {distance}, Mode: {mode}, "
              f"\nCSV file containing neighboring data of bacteria: {self.csv_neighbor_file}")
        # window.destroy()
        # Use this data to display images based on user input
        self.display_images(t1, t2, object_numbers, ids, parent_ids, distance, mode)

    def display_images(self, timestep1, timestep2, specified_object_numbers, specified_ids_str,
                       specified_parent_ids_str, neighbor_distance_str, selected_mode):

        # Parse the object Numbers, specified IDs, and Parent IDs

        object_numbers = specified_object_numbers.replace(" ", "").split(',') if specified_object_numbers else []
        object_numbers = [int(x) for x in object_numbers if x.isdigit()]  # Convert to list of integers

        specified_ids = specified_ids_str.replace(" ", "").split(',') if specified_ids_str else []
        specified_ids = [int(x) for x in specified_ids if x.isdigit()]  # Convert to list of integers

        specified_parent_ids = specified_parent_ids_str.replace(" ", "").split(',') if specified_parent_ids_str else []
        specified_parent_ids = [int(x) for x in specified_parent_ids if x.isdigit()]  # Convert to list of integers

        # Parse the neighbor distance if provided
        try:
            neighbor_distance = float(neighbor_distance_str) if neighbor_distance_str else None
        except ValueError:
            messagebox.showerror("Invalid Input", "Please enter a valid number for neighbor distance.")
            return

        # Load the data
        if self.bacteria_df is None:
            messagebox.showerror("No CSV File", "Please select a CSV file.")
            return

        # Get the unique time steps from the dataset
        available_time_steps = self.bacteria_df['ImageNumber'].unique()

        # Check if the time steps are valid
        if int(timestep1) not in available_time_steps or int(timestep2) not in available_time_steps:
            max_time_step = available_time_steps.max()
            messagebox.showerror("Invalid Time Step", f"Specified time step is incorrect. "
                                                      f"The dataset has time steps ranging from 1 to {max_time_step}.")
            return

        if len(self.raw_images) == 0:
            messagebox.showerror("No Image Directory", "Please select a raw images directory.")
            return

        # Filter for timestep 1 and 2
        df_t1 = self.bacteria_df[self.bacteria_df['ImageNumber'] == int(timestep1)].reset_index(drop=True)
        df_t2 = self.bacteria_df[self.bacteria_df['ImageNumber'] == int(timestep2)].reset_index(drop=True)

        # Close any previously opened windows
        self.close_previous_windows()

        # Check the selected mode and display images accordingly
        if selected_mode == "Two Different Slides":
            if specified_ids:
                # Display each time step in a separate window
                self.show_image(df_t1, self.raw_images[int(timestep1) - 1], f"Time Step {timestep1}",
                                object_numbers, specified_ids.copy(), specified_parent_ids, neighbor_distance,
                                self.df_neighbor)
                self.show_image(df_t2, self.raw_images[int(timestep2) - 1], f"Time Step {timestep2}",
                                object_numbers, specified_ids.copy(), specified_parent_ids, neighbor_distance,
                                self.df_neighbor)
            else:
                # Display each time step in a separate window
                self.show_image(df_t1, self.raw_images[int(timestep1) - 1], f"Time Step {timestep1}",
                                object_numbers, specified_ids, specified_parent_ids, neighbor_distance,
                                self.df_neighbor)
                self.show_image(df_t2, self.raw_images[int(timestep2) - 1], f"Time Step {timestep2}",
                                object_numbers, specified_ids, specified_parent_ids, neighbor_distance,
                                self.df_neighbor)

        elif selected_mode == "Slide Show":
            # Display both time steps in a single window with slider control
            self.show_sliding_image(df_t1, df_t2, self.raw_images[int(timestep1) - 1],
                                    self.raw_images[int(timestep2) - 1], object_numbers, specified_ids,
                                    specified_parent_ids, neighbor_distance,
                                    timestep1, timestep2, self.df_neighbor)

    def show_image(self, df_current, image_path, window_title, specified_object_numbers=None, specified_ids=None,
                   specified_parent_ids=None, neighbor_distance=None, df_neighbor=None):

        # Create a new window for the plot
        img_window = tk.Toplevel(self)
        img_window.title(window_title)

        # Keep track of this window
        self.open_windows.append(img_window)

        # Configure the window to resize
        img_window.grid_rowconfigure(0, weight=1)
        img_window.grid_columnconfigure(0, weight=1)

        # Create a frame to hold the canvas and toolbar
        frame = tk.Frame(img_window)
        frame.grid(row=0, column=0, sticky='nsew')

        # Read the image
        img = cv2.imread(image_path)
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # Create a figure
        fig, ax = plt.subplots()

        # Dictionary to store line objects and their corresponding metadata
        line_objects = {}

        # Display the raw image
        ax.imshow(img_rgb)

        # Extract objects and plot as per the logic from your existing utils
        objects_center_coord_x, objects_center_coord_y, objects_major_current, objects_orientation_current = \
            extract_bacteria_info(df_current, self.um_per_pixel, self.center_coord_cols,
                                  "AreaShape_MajorAxisLength", "AreaShape_Orientation")

        if specified_object_numbers or specified_ids or specified_parent_ids:

            selected_bac_by_user = df_current.loc[(df_current['ObjectNumber'].isin(specified_object_numbers)) |
                                                  (df_current['id'].isin(specified_ids)) |
                                                  (df_current['parent_id'].isin(specified_parent_ids))]

            if df_neighbor.shape[0] > 0:
                selected_bac_with_neighbors = \
                    selected_bac_by_user.merge(df_neighbor, left_on=['ImageNumber', 'ObjectNumber'],
                                               right_on=['First Image Number', 'First Object Number'], how='left')
                selected_bac_with_neighbors_with_info = \
                    selected_bac_with_neighbors.merge(df_current,
                                                      left_on=['Second Image Number', 'Second Object Number'],
                                                      right_on=['ImageNumber', 'ObjectNumber'], how='left',
                                                      suffixes=('_1', '_2'))

                neighbors_bac_id = [v for v in selected_bac_with_neighbors_with_info['id_2'].values.tolist() if
                                    str(v) != 'nan']
                specified_ids.extend(neighbors_bac_id)

            # Check if the object is within the neighbor distance from the specified objects
            if neighbor_distance:
                other_bac = df_current.loc[~ df_current['id'].isin(selected_bac_by_user['id'].values)]

                merged_df = selected_bac_by_user.merge(other_bac, how='cross', suffixes=('_1', '_2'))

                dist_cond = (
                        np.sqrt(np.power((merged_df[self.center_coord_cols['x'] + '_1'] -
                                          merged_df[self.center_coord_cols['x'] + '_2']), 2) +
                                np.power((merged_df[self.center_coord_cols['y'] + '_1'] -
                                          merged_df[self.center_coord_cols['y'] + '_2']), 2)) <= neighbor_distance)

                neighbors_bac_id = merged_df[dist_cond]['id_2'].values.tolist()
                specified_ids.extend(neighbors_bac_id)

        # Filter objects by ID, Parent ID, or distance
        for cell_idx in range(df_current.shape[0]):
            cell_obj_num = df_current.iloc[cell_idx]['ObjectNumber']
            cell_id = df_current.iloc[cell_idx]['id']
            parent_id = df_current.iloc[cell_idx]['parent_id']
            center_current = (objects_center_coord_x[cell_idx], objects_center_coord_y[cell_idx])

            # Skip if the object doesn't match the specified IDs or Parent IDs
            if specified_ids and specified_parent_ids and specified_object_numbers:
                if ((cell_id not in specified_ids) and (parent_id not in specified_parent_ids) and
                        (cell_obj_num not in specified_object_numbers)):
                    continue
            elif specified_ids and specified_parent_ids:
                if (cell_id not in specified_ids) and (parent_id not in specified_parent_ids):
                    continue
            elif specified_ids and specified_object_numbers:
                if (cell_id not in specified_ids) and (cell_obj_num not in specified_object_numbers):
                    continue
            elif specified_object_numbers and specified_parent_ids:
                if (cell_obj_num not in specified_object_numbers) and (parent_id not in specified_parent_ids):
                    continue
            elif specified_ids:
                if cell_id not in specified_ids:
                    continue
            elif specified_parent_ids:
                if parent_id not in specified_parent_ids:
                    continue
            elif specified_object_numbers:
                if cell_obj_num not in specified_object_numbers:
                    continue

            # print("Object Number: " + str(cell_obj_num) + " Id: " + str(cell_id) + " parent id: " + str(parent_id))

            # Plot bacteria and labels
            major_current = (objects_major_current.iloc[cell_idx]) / 2
            angle_current = objects_orientation_current.iloc[cell_idx]
            ends = calculate_bac_endpoints([center_current[0], center_current[1]], major_current, angle_current)

            # Endpoints of the major axis of the ellipse
            node_x1_x_current = ends[0][0]
            node_x1_y_current = ends[0][1]
            node_x2_x_current = ends[1][0]
            node_x2_y_current = ends[1][1]

            line, = ax.plot([node_x1_x_current, node_x2_x_current], [node_x1_y_current, node_x2_y_current], lw=1,
                            solid_capstyle="round", color=self.object_color, picker=1)

            # Store the line object with its corresponding metadata (Object Number, ID, Parent ID)
            line_objects[line] = (cell_obj_num, cell_id, parent_id)

            # Add cell ID and parent ID as text
            pos1x = np.abs(node_x1_x_current + center_current[0]) / 2
            pos2x = np.abs(node_x2_x_current + center_current[0]) / 2

            pos1y = np.abs(node_x1_y_current + center_current[1]) / 2
            pos2y = np.abs(node_x2_y_current + center_current[1]) / 2

            final_pos1x = np.abs(pos1x + center_current[0]) / 2
            final_pos2x = np.abs(pos2x + center_current[0]) / 2

            final_pos1y = np.abs(pos1y + center_current[1]) / 2
            final_pos2y = np.abs(pos2y + center_current[1]) / 2

            # Add ID and Parent ID as text on the object
            ax.text(final_pos1x, final_pos1y, cell_id, fontsize=self.font_size, color="#ff0000")
            ax.text(final_pos2x, final_pos2y, parent_id, fontsize=self.font_size, color="#0000ff")

        # Add legend inside the window
        parent_patch = mpatches.Patch(color='#0000ff', label='Parent id')
        id_patch = mpatches.Patch(color='#ff0000', label='Id')
        plt.legend(handles=[parent_patch, id_patch], loc='upper center', ncol=6,
                   bbox_to_anchor=(.5, 1.1), prop={'size': 7})

        def on_click(event):
            """Callback function for mouse clicks on plot."""
            for line, metadata in line_objects.items():
                if line.contains(event.mouseevent)[0]:
                    obj_num, obj_id, parent_id = metadata
                    messagebox.showinfo("Object Clicked",
                                        f"Object Number: {obj_num}\nID: {obj_id}\nParent ID: {parent_id}")
                    break  # Stop after finding the first matching line

        # Add event listener for clicks (closure method, avoids global variables)
        fig.canvas.mpl_connect('pick_event', on_click)

        # Embed the figure into the Tkinter window with zoom functionality
        canvas = FigureCanvasTkAgg(fig, master=frame)
        canvas.draw()
        canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

        # Add toolbar for zooming
        toolbar = NavigationToolbar2Tk(canvas, frame)
        toolbar.update()
        toolbar.pack()

        # Resize image dynamically with the window
        fig.tight_layout()

    def show_sliding_image(self, df_t1, df_t2, image_path_t1, image_path_t2, specified_object_numbers=None,
                           specified_ids=None, specified_parent_ids=None, neighbor_distance=None, timestep1=None,
                           timestep2=None, df_neighbor=None):
        # Create a new window for the sliding plot
        slide_window = tk.Toplevel(self)
        slide_window.title("Trackrefiner Slider Control")

        # Keep track of this window
        self.open_windows.append(slide_window)

        # Configure the window to resize
        slide_window.grid_rowconfigure(0, weight=1)
        slide_window.grid_columnconfigure(0, weight=1)

        # Create a frame to hold the canvas and another frame for the toolbar
        canvas_frame = tk.Frame(slide_window)
        toolbar_frame = tk.Frame(slide_window)

        canvas_frame.grid(row=0, column=0, sticky='nsew')
        toolbar_frame.grid(row=1, column=0, sticky='ew')

        # Make sure the canvas_frame and canvas expand when the window is resized
        canvas_frame.grid_rowconfigure(0, weight=1)
        canvas_frame.grid_columnconfigure(0, weight=1)

        # Read the images for both time steps
        img_t1 = cv2.imread(image_path_t1)
        img_t1_rgb = cv2.cvtColor(img_t1, cv2.COLOR_BGR2RGB)

        img_t2 = cv2.imread(image_path_t2)
        img_t2_rgb = cv2.cvtColor(img_t2, cv2.COLOR_BGR2RGB)

        # Create a figure and adjust it to fit the canvas
        fig, ax = plt.subplots()

        # Label for showing time step name dynamically
        time_step_label = tk.Label(slide_window, text=f"Time Step {timestep1}")
        time_step_label.grid(row=3, column=0, sticky='ew')

        # Variables to store zoom level and position
        current_x_lim = None
        current_y_lim = None

        # Function to update the image based on slider value
        def update_image(val):

            nonlocal current_x_lim, current_y_lim

            time_step = int(val)

            # Store the current zoom level if it's the first time switching
            if current_x_lim is None and current_y_lim is None:
                current_x_lim = ax.get_xlim()
                current_y_lim = ax.get_ylim()

            ax.clear()

            # Choose image and dataframe based on slider value
            if time_step == 1:
                current_img = img_t1_rgb
                current_df = df_t1
                title = "Time Step " + str(timestep1)
                time_step_label.config(text=f"Time Step {timestep1}")  # Update label to show timestep1
            else:
                current_img = img_t2_rgb
                current_df = df_t2
                title = "Time Step " + str(timestep2)
                time_step_label.config(text=f"Time Step {timestep2}")  # Update label to show timestep2

            ax.imshow(current_img)
            # ax.set_title(title)

            # Plot bacteria objects (if specified) for the current time step
            if specified_ids:
                self.plot_bacteria_on_image(ax, current_df, specified_object_numbers, specified_ids.copy(),
                                            specified_parent_ids, neighbor_distance, df_neighbor, line_objects)
            else:
                self.plot_bacteria_on_image(ax, current_df, specified_object_numbers, specified_ids.copy(),
                                            specified_parent_ids, neighbor_distance, df_neighbor, line_objects)

            # Reapply the stored zoom level
            ax.set_xlim(current_x_lim)
            ax.set_ylim(current_y_lim)

            canvas.draw()

        # Display the initial image (time step 1)
        ax.imshow(img_t1_rgb)
        line_objects = {}
        # ax.set_title("Time Step 1")

        # Plot bacteria objects for time step 1
        if specified_ids:
            self.plot_bacteria_on_image(ax, df_t1, specified_object_numbers, specified_ids.copy(), specified_parent_ids,
                                        neighbor_distance, df_neighbor, line_objects)
        else:
            self.plot_bacteria_on_image(ax, df_t1, specified_object_numbers, specified_ids, specified_parent_ids,
                                        neighbor_distance, df_neighbor, line_objects)

        # Embed the figure into the Tkinter window with zoom functionality
        canvas = FigureCanvasTkAgg(fig, master=canvas_frame)
        canvas.draw()
        canvas.get_tk_widget().grid(row=0, column=0, sticky='nsew')

        def on_click(event):
            """Callback function for mouse clicks on plot."""
            for line, metadata in line_objects.items():
                if line.contains(event.mouseevent)[0]:  # Ensure mouseevent is used
                    obj_num, obj_id, parent_id = metadata
                    messagebox.showinfo("Object Clicked",
                                        f"Object Number: {obj_num}\nID: {obj_id}\nParent ID: {parent_id}")
                    break

        fig.canvas.mpl_connect('pick_event', on_click)

        # Configure resizing behavior for the figure
        def on_resize(event):
            # Adjust figure layout on resize
            fig.tight_layout()
            canvas.draw()

        # Bind the resize event to dynamically update the layout
        canvas.get_tk_widget().bind("<Configure>", on_resize)

        # Add toolbar for zooming, in the separate toolbar frame
        toolbar = NavigationToolbar2Tk(canvas, toolbar_frame)
        toolbar.update()
        toolbar.pack(side=tk.BOTTOM, fill=tk.X)

        # Create a slider widget to switch between the two time steps, using grid instead of pack
        slider = tk.Scale(slide_window, from_=1, to=2, orient=tk.HORIZONTAL, label="Slide", command=update_image)
        slider.grid(row=2, column=0, sticky='ew')

        # Allow the slider to resize along the X-axis
        slide_window.grid_rowconfigure(2, weight=0)
        slide_window.grid_columnconfigure(0, weight=1)

    def plot_bacteria_on_image(self, ax, df_current, specified_object_numbers, specified_ids, specified_parent_ids,
                               neighbor_distance, df_neighbor, line_objects):

        # Extract information about the bacteria: center coordinates, major axis length, and orientation
        objects_center_coord_x, objects_center_coord_y, objects_major_current, objects_orientation_current = \
            extract_bacteria_info(df_current, self.um_per_pixel, self.center_coord_cols,
                                  "AreaShape_MajorAxisLength", "AreaShape_Orientation")

        if specified_object_numbers or specified_ids or specified_parent_ids:
            selected_bac_by_user = df_current.loc[(df_current['ObjectNumber'].isin(specified_object_numbers)) |
                                                  (df_current['id'].isin(specified_ids)) |
                                                  (df_current['parent_id'].isin(specified_parent_ids))]

            if df_neighbor.shape[0] > 0:
                selected_bac_with_neighbors = \
                    selected_bac_by_user.merge(df_neighbor, left_on=['ImageNumber', 'ObjectNumber'],
                                               right_on=['First Image Number', 'First Object Number'], how='left')
                selected_bac_with_neighbors_with_info = \
                    selected_bac_with_neighbors.merge(df_current,
                                                      left_on=['Second Image Number', 'Second Object Number'],
                                                      right_on=['ImageNumber', 'ObjectNumber'], how='left',
                                                      suffixes=('_1', '_2'))

                neighbors_bac_id = [v for v in selected_bac_with_neighbors_with_info['id_2'].values.tolist() if
                                    str(v) != 'nan']
                specified_ids.extend(neighbors_bac_id)

            # Check if the object is within the neighbor distance from the specified objects
            if neighbor_distance:
                other_bac = df_current.loc[~ df_current['id'].isin(selected_bac_by_user['id'].values)]

                merged_df = selected_bac_by_user.merge(other_bac, how='cross', suffixes=('_1', '_2'))

                dist_cond = (
                        np.sqrt(np.power((merged_df[self.center_coord_cols['x'] + '_1'] -
                                          merged_df[self.center_coord_cols['x'] + '_2']), 2) +
                                np.power((merged_df[self.center_coord_cols['y'] + '_1'] -
                                          merged_df[self.center_coord_cols['y'] + '_2']), 2)) <= neighbor_distance)

                neighbors_bac_id = merged_df[dist_cond]['id_2'].values.tolist()
                specified_ids.extend(neighbors_bac_id)

        # Loop through each bacteria object and plot it

        for cell_idx in range(df_current.shape[0]):
            cell_obj_num = df_current.iloc[cell_idx]['ObjectNumber']
            cell_id = df_current.iloc[cell_idx]['id']
            parent_id = df_current.iloc[cell_idx]['parent_id']
            center_current = (objects_center_coord_x[cell_idx], objects_center_coord_y[cell_idx])

            # Skip objects that don't match the specified IDs or Parent IDs
            if specified_ids and specified_parent_ids and specified_object_numbers:
                if ((cell_id not in specified_ids) and (parent_id not in specified_parent_ids) and
                        (cell_obj_num not in specified_object_numbers)):
                    continue
            elif specified_ids and specified_parent_ids:
                if (cell_id not in specified_ids) and (parent_id not in specified_parent_ids):
                    continue
            elif specified_ids and specified_object_numbers:
                if (cell_id not in specified_ids) and (cell_obj_num not in specified_object_numbers):
                    continue
            elif specified_object_numbers and specified_parent_ids:
                if (cell_obj_num not in specified_object_numbers) and (parent_id not in specified_parent_ids):
                    continue
            elif specified_ids:
                if cell_id not in specified_ids:
                    continue
            elif specified_parent_ids:
                if parent_id not in specified_parent_ids:
                    continue
            elif specified_object_numbers:
                if cell_obj_num not in specified_object_numbers:
                    continue

            # Plot the bacteria's major axis as a line
            major_current = (objects_major_current.iloc[cell_idx]) / 2
            angle_current = objects_orientation_current.iloc[cell_idx]
            ends = calculate_bac_endpoints([center_current[0], center_current[1]], major_current, angle_current)

            # Endpoints of the major axis of the ellipse
            node_x1_x_current = ends[0][0]
            node_x1_y_current = ends[0][1]
            node_x2_x_current = ends[1][0]
            node_x2_y_current = ends[1][1]

            line, = ax.plot([node_x1_x_current, node_x2_x_current], [node_x1_y_current, node_x2_y_current], lw=1,
                            solid_capstyle="round", color=self.object_color, picker=1)

            # Store the line object with its corresponding metadata (Object Number, ID, Parent ID)
            line_objects[line] = (cell_obj_num, cell_id, parent_id)

            # Add the bacteria's ID and parent ID as text labels
            pos1x = np.abs(node_x1_x_current + center_current[0]) / 2
            pos2x = np.abs(node_x2_x_current + center_current[0]) / 2

            pos1y = np.abs(node_x1_y_current + center_current[1]) / 2
            pos2y = np.abs(node_x2_y_current + center_current[1]) / 2

            final_pos1x = np.abs(pos1x + center_current[0]) / 2
            final_pos2x = np.abs(pos2x + center_current[0]) / 2

            final_pos1y = np.abs(pos1y + center_current[1]) / 2
            final_pos2y = np.abs(pos2y + center_current[1]) / 2

            # Add ID and Parent ID labels to the image
            ax.text(final_pos1x, final_pos1y, cell_id, fontsize=self.font_size, color="#ff0000")
            ax.text(final_pos2x, final_pos2y, parent_id, fontsize=self.font_size, color="#0000ff")

        # Add a legend for the ID and Parent ID colors
        parent_patch = mpatches.Patch(color='#0000ff', label='Parent id')
        id_patch = mpatches.Patch(color='#ff0000', label='Id')
        ax.legend(handles=[parent_patch, id_patch], loc='upper center', ncol=6,
                  bbox_to_anchor=(.5, 1.1), prop={'size': 7})

    def show_all_time_steps(self):

        self.update_font_size()  # Update font size before performing actions
        self.update_um_per_pixel()

        # Close any previously opened windows
        self.close_previous_windows()

        # Get all available time steps from the dataset
        available_time_steps = self.bacteria_df['ImageNumber'].unique()

        if not available_time_steps.size:
            messagebox.showerror("No Time Steps", "No time steps found in the dataset.")
            return

        # Create a new window for the sliding plot
        slide_window = tk.Toplevel(self)
        slide_window.title("Trackrefiner Show All Time Steps")

        # Keep track of this window
        self.open_windows.append(slide_window)

        # Create a frame to hold the canvas and toolbar
        canvas_frame = tk.Frame(slide_window)
        toolbar_frame = tk.Frame(slide_window)

        canvas_frame.grid(row=0, column=0, sticky='nsew')
        toolbar_frame.grid(row=1, column=0, sticky='ew')

        canvas_frame.grid_rowconfigure(0, weight=1)
        canvas_frame.grid_columnconfigure(0, weight=1)

        # Create a figure and adjust it to fit the canvas
        fig, ax = plt.subplots()

        # Label for showing time step name dynamically
        time_step_label = tk.Label(slide_window, text=f"Time Step {available_time_steps[0]}")
        time_step_label.grid(row=3, column=0, sticky='ew')

        # Variables to store zoom level and position
        current_x_lim = None
        current_y_lim = None

        # Function to update the image based on the slider value
        def update_image(val):

            nonlocal current_x_lim, current_y_lim

            line_objects = {}

            time_step_index = int(val)
            timestep = available_time_steps[time_step_index]

            # Store the current zoom level if it's the first time switching
            if current_x_lim is None and current_y_lim is None:
                current_x_lim = ax.get_xlim()
                current_y_lim = ax.get_ylim()

            # Clear the axes
            ax.clear()

            # Update the time step label dynamically
            time_step_label.config(text=f"Time Step {timestep}")

            # Filter data for the current time step
            df_current = self.bacteria_df.loc[self.bacteria_df['ImageNumber'] == timestep].reset_index(drop=True)

            # Load the corresponding image
            img_path = self.raw_images[timestep - 1]
            img = cv2.imread(img_path)
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            ax.imshow(img_rgb)
            self.plot_bacteria_on_image(ax, df_current, [], [], [],
                                        None, None, line_objects)

            # Reapply the stored zoom level
            ax.set_xlim(current_x_lim)
            ax.set_ylim(current_y_lim)

            # Draw the updated canvas
            canvas.draw()

            def on_click(event):
                """Callback function for mouse clicks on plot."""
                for line, metadata in line_objects.items():
                    if line.contains(event.mouseevent)[0]:  # Ensure mouseevent is used
                        obj_num, obj_id, parent_id = metadata
                        messagebox.showinfo("Object Clicked",
                                            f"Object Number: {obj_num}\nID: {obj_id}\nParent ID: {parent_id}")
                        break

            fig.canvas.mpl_connect('pick_event', on_click)  # Rebind the event each time the image updates

        # Display the first image (corresponding to the first time step)
        first_time_step = available_time_steps[0]
        img = cv2.imread(self.raw_images[first_time_step - 1])
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        ax.imshow(img_rgb)

        line_objects = {}

        # Plot bacteria objects for the first time step
        df_t1 = self.bacteria_df[self.bacteria_df['ImageNumber'] == first_time_step].reset_index(drop=True)
        self.plot_bacteria_on_image(ax, df_t1, [], [], [],
                                    None, None, line_objects)

        # Embed the figure into the Tkinter window with zoom functionality
        canvas = FigureCanvasTkAgg(fig, master=canvas_frame)
        canvas.draw()
        canvas.get_tk_widget().grid(row=0, column=0, sticky='nsew')

        def on_click(event):
            """Callback function for mouse clicks on plot."""
            for line, metadata in line_objects.items():
                if line.contains(event.mouseevent)[0]:
                    obj_num, obj_id, parent_id = metadata
                    messagebox.showinfo("Object Clicked",
                                        f"Object Number: {obj_num}\nID: {obj_id}\nParent ID: {parent_id}")
                    break  # Stop after finding the first matching line

        # Add event listener for clicks (closure method, avoids global variables)
        fig.canvas.mpl_connect('pick_event', on_click)

        # Add toolbar for zooming
        toolbar = NavigationToolbar2Tk(canvas, toolbar_frame)
        toolbar.update()
        toolbar.pack()

        # Create a slider widget to switch between time steps
        slider = tk.Scale(slide_window, from_=0, to=len(available_time_steps) - 1, orient=tk.HORIZONTAL, label="Slide",
                          command=update_image)
        slider.grid(row=2, column=0, sticky='ew')

    def remove_specific_window(self, target_window_name_list):
        for window in self.open_windows[:]:  # Iterate over a copy to safely modify the list
            if window.winfo_exists():
                window_title = window.winfo_toplevel().title()
                if any(window_title.startswith(name) for name in target_window_name_list):
                    window.destroy()
                    self.open_windows.remove(window)  # Remove from tracking list

    def close_previous_windows(self):
        """Close all previously opened windows."""
        for window in self.open_windows:
            if window.winfo_exists():  # Check if the window still exists before destroying it
                window.destroy()
        # Clear the list after closing all windows
        self.open_windows.clear()

    def check_if_all_closed(self):
        """Check if all windows are closed and close the main app if true."""
        open_windows_exist = any(window.winfo_exists() for window in self.open_windows)
        if not open_windows_exist:
            self.quit()

    def on_main_window_close(self):
        """Close all open windows and then close the main app."""
        # Close all the opened windows first
        self.close_previous_windows()
        # Quit the main app
        self.quit()


# Run the GUI application
def main():
    app = TrackingGUI()
    app.mainloop()


if __name__ == "__main__":
    main()
