test.py

import pyabf
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('TkAgg')
from tkinter import *
from tkinter import Tk, filedialog, StringVar, Button, Frame, Label
from matplotlib.widgets import SpanSelector
from matplotlib.backends.backend_tkagg import NavigationToolbar2Tk
from tkinter.ttk import OptionMenu
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.backend_bases import PickEvent
from matplotlib.lines import Line2D
import threading
import csv
import numpy as np
from scipy import signal
import pickle 

   
def open_trace(file_name):
    #abf = pyabf.ABF(file_name)
    trace = Trace(file_name)  # Create a Trace object linked to the file
    #trace.initialize_abf()
    #trace.time, trace.current = abf.sweepX, abf.sweepY
    return trace

class Trace():
    def __init__(self, file_name):
        self.file_name = file_name
        self.abf = None
        self.time, self.current = None, None  # Initialize time and current as None
        self.segments = []
        self.baseline = None
        self.peaks = []

    def initialize_abf(self):
        if self.abf is None:
            self.abf = pyabf.ABF(self.file_name)
            self.time, self.current = self.abf.sweepX, self.abf.sweepY

    def serialize(self):
        # Before pickling, initialize the abf object
        # Exclude the abf object from pickling by setting it to None
        abf = self.abf
        self.abf = None
        # Serialize the object without the abf attribute
        data = pickle.dumps(self)
        # Restore the abf attribute after pickling
        self.abf = abf
        return data

    @staticmethod
    def deserialize(data):
        trace = pickle.loads(data)
        # After deserialization, reinitialize the abf object
        trace.initialize_abf()
        return trace


class Segment:
    def __init__(self, file_name, start_time, end_time, category):
        self.file_name = file_name
        self.start_time = start_time
        self.end_time = end_time
        self.duration = end_time - start_time
        self.category = category
    
    def serialize(self):
       return pickle.dumps(self)

    @staticmethod
    def deserialize(data):
        return pickle.loads(data) 
    
class Peak():
    def __init__(self, trace):
        self.file_name = trace.file_name
        self.trace = trace
        self.peak_x = None
        self.peak_y = None
    
    def set_peak_position (self, x, y):
        self.peak_x = x
        self.peak_y =y

    def highest_point(self, x_clicked):
        index_clicked = np.argmin(np.abs(self.trace.time - x_clicked))
        search_window = 100 #search window for setting peaks
        
        # Create a 1-second window surrounding the clicked point
        points_per_second = int(1000 * self.trace.abf.dataPointsPerMs) #20 points per Ms given a sampling freq of 20 Hz
        start_index = max(0, index_clicked - search_window)
        end_index = min(len(self.trace.time) - search_window, index_clicked + search_window)
        
        range_y = self.trace.current[start_index:end_index] #change so that it uses plottedtrace
        max_index = np.argmax(abs(range_y))
        index_highest = start_index + max_index
        self.x_highest = self.trace.time[index_highest]
        self.y_highest = self.trace.current[index_highest]
        self.set_peak_position(self.x_highest, self.y_highest)

        #plot
        ax.plot(self.x_highest, self.y_highest, 'ro')
        fig.canvas.draw()

    def on_line_click(self, event):
        if isinstance(event, PickEvent):
            index = event.ind[0]
            x_clicked = self.trace.time[index]
            y_clicked = self.trace.current[index]
            self.highest_point(x_clicked)
        else:
            print('not instance')

    def serialize(self):
       return pickle.dumps(self)

    @staticmethod
    def deserialize(data):
        return pickle.loads(data) 

    

class FileManager:
    def __init__(self, file_names, segments, peaks, traces):
        self.file_names = file_names
        self.segments = segments
        self.current_file_index = 0
        self.current_file_name = ""
        self.current_segment = None
        self.peaks = peaks
        self.window_size = 10
        self.traces = traces
    
    def update_graph(self):
        self.current_file_name = self.file_names[self.current_file_index]
        ax.clear()
        #self.time, self.current = open_trace(self.current_file_name)
        self.trace = open_trace(self.current_file_name)
        self.traces.append(self.trace)
        self.trace.initialize_abf()
        self.time, self.current = self.trace.time, self.trace.current
        #self.traces.append(self.trace)
        
        #self.plottedtrace, = ax.plot(self.time, self.current, linewidth=0.5)
        self.plottedtrace, = ax.plot(self.time, self.current, linewidth=0.5)
        self.plottedtrace.set_picker(1)
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Current (pA)")
        
        # Load and plot segments for the current file, if available
        current_segments = self.get_segments_for_file(self.current_file_name)
        if current_segments:
            for segment in current_segments:
                color = 'green' if segment.category == 'burst' else 'red'
                span = ax.axvspan(segment.start_time, segment.end_time, color=color, alpha=0.3)

        current_peaks = self.get_peaks_for_file(self.current_file_name)
        if current_peaks:
            for peak in current_peaks:
                ax.plot(peak.peak_x, peak.peak_y, 'ro')
                fig.canvas.draw()
        
        canvas.draw()
      
    
    def filter_trace(self, event=None):
        def ButterBandStop(current):
            fs = 20000  # Sampling frequency (Hz)
            fstop = [600, 5000]

            # Design Butterworth filter
            b, a = signal.butter(2, fstop, btype='bandstop', analog=False, fs=fs)
            current_filtered = signal.filtfilt(b, a, current)

            return current, current_filtered

        # open current trace
        #self.time, self.current = open_trace(self.current_file_name)
        xlim = ax.get_xlim()
        ylim = ax.get_ylim()
        
        # Apply the Butterworth bandstop filter if it's not already applied
        if not filter_button.is_filtered:
            current, current_filtered = ButterBandStop(self.current)

            # Plot the filtered trace
            ax.clear()
            self.plottedtrace, = ax.plot(self.time, current_filtered, linewidth=0.5, color='black')
            self.plottedtrace.set_picker(1)
            ax.set_xlabel("Time (s)")
            ax.set_ylabel("Current (pA)")
            filter_button.is_filtered = True
            filter_button.config(text='Raw Trace')

        else:
            ax.clear()
            self.plottedtrace, = ax.plot(self.time, self.current, linewidth=0.5)
            self.plottedtrace.set_picker(1)
            ax.set_xlabel("Time (s)")
            ax.set_ylabel("Current (pA)")
            filter_button.is_filtered = False
            filter_button.config(text='Filter')

        # Restore the previous zoom settings
        ax.set_xlim(xlim)
        ax.set_ylim(ylim)
        
        # Load and plot segments for the current file, if available
        current_segments = self.get_segments_for_file(self.current_file_name)
        if current_segments:
            for segment in current_segments:
                color = 'green' if segment.category == 'burst' else 'red'
                span = ax.axvspan(segment.start_time, segment.end_time, color=color, alpha=0.3)

        canvas.draw()
        #self.plottedtrace.draw()

        current_peaks = self.get_peaks_for_file(self.current_file_name)
        if current_peaks:
            for peak in current_peaks:
                ax.plot(peak.peak_x, peak.peak_y, 'ro')
                
        canvas.draw()

    def add_files(self, event=None):
        #root.withdraw()
        new_files = filedialog.askopenfilenames(filetypes=[('ABF Files', '*.abf')])
        self.file_names.extend(new_files)
        print(f"Added files: {', '.join(new_files)}")
        if len(self.file_names) > 0:
            #self.update_graph()
            self.update_dropdown_menu()
    
    def next_file(self, event=None):
        self.current_file_index = (self.current_file_index + 1) % len(self.file_names)
        filter_button.is_filtered = False  # Reset the filtering state
        filter_button.config(text='Filter')
        #self.update_graph()
        self.update_dropdown_menu()
    
    def previous_file(self, event=None):
        self.current_file_index = (self.current_file_index - 1) % len(self.file_names)
        filter_button.is_filtered = False # Reset the filtering state
        filter_button.config(text='Filter')
        #self.update_graph()
        self.update_dropdown_menu()

    def save_traces(self):
        if len(traces) == 0:
            print("No data to save.")
            return
        
        filename = filedialog.asksaveasfilename(defaultextension='.pickle', filetypes=[('Pickle Files', '*.pickle')])
        if filename:
            with open(filename, 'wb') as file:
                pickle.dump(traces, file)
            print(f"All traces saved as {filename}")

    def load_traces(self):
        filename = filedialog.askopenfilename(filetypes=[('Pickle Files', '*.pickle')])
        if filename:
            with open(filename, 'rb') as file:
                loaded_traces = pickle.load(file)

            # Update the list of traces with the loaded data
            self.traces = loaded_traces
            
            # Update the GUI display and plots accordingly
            self.update_graph()
    
    def save_table(self, event=None):
        if len(segments) == 0:
            print("No segments to save.")
            return

        filename = filedialog.asksaveasfilename(defaultextension='.csv', filetypes=[('CSV Files', '*.csv')])
        if filename:
            with open(filename, 'w', newline='') as file:
                writer = csv.writer(file)
                writer.writerow(["File Name", "Start Time", "End Time", "Duration", "Category", "Events"])
                for segment in segments:
                    writer.writerow([segment.file_name, segment.start_time, segment.end_time, segment.duration, segment.category, segment.events])
            print(f"Table saved as {filename}")
    
    def save_segments(self, event=None):
        if len(segments) == 0:
            print("No segments to save.")
            return

        filename = filedialog.asksaveasfilename(defaultextension='.pickle', filetypes=[('Pickle Files', '*.pickle')])
        if filename:
            with open(filename, 'wb') as file:
                pickle.dump(segments, file)
                #pickle.dump(self.peaks, file)
            print(f"Segments saved as {filename}")
   
    def get_segments_for_file(self, file_name):
        self.file_segments = []
        for segment in self.segments:
            if segment.file_name == file_name:
                self.file_segments.append(segment)
        return self.file_segments
    
    def get_peaks_for_file(self, file_name):
        self.file_peaks = []
        for peak in self.peaks:
            if peak.file_name == file_name:
                self.file_peaks.append(peak)
        return self.file_peaks

    def load_segments(self, event=None):
        filename = filedialog.askopenfilename(filetypes=[('Pickle Files', '*.pickle')])
        if filename:
            with open(filename, 'rb') as file:
                loaded_segments = pickle.load(file)
            self.segments = loaded_segments
            print(f"Segments loaded from {filename}")

            self.update_graph()
    
    def update_dropdown_menu(self):
        file_menu["menu"].delete(0, "end")  # Clear the current menu options

        if len(self.file_names) > 0:
            for file_name in self.file_names:
                name= file_name.split("/")[-1].split(".")[0]
                #file_menu["menu"].add_command(label=file_name, command=lambda value=file_name: set_current_file(value))
                file_menu["menu"].add_command(label=name, command=lambda value=file_name: self.set_current_file(value))
            self.current_file_name = self.file_names[self.current_file_index]
            display_name = str(self.current_file_index) + ": " + self.current_file_name.split("/")[-1].split(".")[0]
            current_file_menu_var.set(display_name)
            self.update_graph()
        else:
            display_name = "No files added"
            current_file_menu_var.set(display_name)

    def set_current_file(self, value):
        self.current_file_name = str(value)
        selected_index = self.file_names.index(value)
        if selected_index != self.current_file_index:
            self.current_file_index = selected_index
            self.current_file_name= value
            display_name = str(self.current_file_index) + ": " + self.current_file_name.split("/")[-1].split(".")[0]
            current_file_menu_var.set(display_name)
            self.update_graph()

    def enable_zoom_and_scroll(self, canvas):
        def on_scroll(event):
            if event.button == 'down':
                direction = 1  # Scroll Up (Move Right)
            elif event.button == 'up':
                direction = -1  # Scroll Down (Move Left)
            else:
                direction = 0  # No scroll

            x_length = ax.get_xlim()[1] - ax.get_xlim()[0]
            x_shift = x_length * 0.1 * direction

            ax.set_xlim(ax.get_xlim()[0] + x_shift, ax.get_xlim()[1] + x_shift)
            canvas.draw_idle()

        def on_key_press(event):
            if event.key == '+':
                ax.set_xlim(ax.get_xlim()[0] * 1.005, ax.get_xlim()[1] * 0.995)
                canvas.draw_idle()
            elif event.key == '-':
                ax.set_xlim(ax.get_xlim()[0] * 0.995, ax.get_xlim()[1] * 1.005)
                canvas.draw_idle()
            elif event.key == '{':
                ax.set_ylim(ax.get_ylim()[0] * 1.05, ax.get_ylim()[1] * 0.95)
                canvas.draw_idle()
            elif event.key == '}':
                ax.set_ylim(ax.get_ylim()[0] * 0.95, ax.get_ylim()[1] * 1.05)
                canvas.draw_idle()
            elif event.key == 'z':
                ax.set_xlim(ax.get_xlim()[0], ax.get_xlim()[0] + 2)
                canvas.draw_idle()
            elif event.key == 'a':
                ax.set_xlim(ax.get_xlim()[0], ax.get_xlim()[0]+10)
                ax.set_ylim(-500, 50)
                canvas.draw_idle()
            elif event.key == 'f':
                ax.set_xlim(self.time[0], self.time[-1] - self.time[0])
                ax.set_ylim(-800, 50)
                canvas.draw_idle()

        def on_motion(event):
            if event.button == '2' and event.xdata and event.ydata:
                x_shift = event.xdata * 0.001
                y_shift = event.ydata * 0.001
                x_limits = ax.get_xlim()
                y_limits = ax.get_ylim()
                x_center = (x_limits[0] + x_limits[1]) / 2
                y_center = (y_limits[0] + y_limits[1]) / 2
                x_scale = (x_limits[1] - x_limits[0]) / 2
                y_scale = (y_limits[1] - y_limits[0]) / 2
                new_x_limits = [x_center - x_scale * 1.05, x_center + x_scale * 1.05]
                new_y_limits = [y_center - y_scale * 1.05, y_center + y_scale * 1.05]
                ax.set_xlim(new_x_limits[0] + x_shift, new_x_limits[1] + x_shift)
                ax.set_ylim(new_y_limits[0] + y_shift, new_y_limits[1] + y_shift)
                canvas.draw_idle()


        def toggle_zoom_and_scroll():
            if zoom_scroll_button.config('text')[-1] == 'Disable Zoom/Scroll':
                zoom_scroll_button.config(text='Enable Zoom/Scroll')
                canvas.mpl_disconnect(scroll_cid)
                canvas.mpl_disconnect(motion_cid)
                canvas.mpl_disconnect(key_press_cid)
            else:
                zoom_scroll_button.config(text='Disable Zoom/Scroll')
                scroll_cid = canvas.mpl_connect('scroll_event', on_scroll)
                motion_cid = canvas.mpl_connect('motion_notify_event', on_motion)
                key_press_cid = canvas.mpl_connect('key_press_event', on_key_press)

        zoom_scroll_button = Button(master=root, text='Disable Zoom/Scroll', command=toggle_zoom_and_scroll)
        zoom_scroll_button.pack()

        scroll_cid = canvas.mpl_connect('scroll_event', on_scroll)
        motion_cid = canvas.mpl_connect('motion_notify_event', on_motion)
        key_press_cid = canvas.mpl_connect('key_press_event', on_key_press)

    def onselect(self, xmin, xmax):
        if not segmentation_enabled:
            return

        self.current_segment = Segment(self.current_file_name, xmin, xmax, "not selected")
        span = ax.axvspan(xmin, xmax, color='grey', alpha=0.3)
        plt.draw()
    
    def mark_as_burst(self, event=None):
        self.current_segment.category = "burst"
        self.current_segment.color = 'green'
        span = ax.axvspan(self.current_segment.start_time, self.current_segment.end_time, color='green', alpha=0.3)
        plt.draw()
        self.segments.append(self.current_segment)  # Add current_segment to segments list
        print(f"Segment marked as 'burst': {self.current_segment.file_name}, {self.current_segment.start_time}s to {self.current_segment.end_time}s")
    
    def mark_as_discard(self, event=None):
        self.current_segment.category = "discard"
        self.current_segment.color = 'red'
        span = ax.axvspan(self.current_segment.start_time, self.current_segment.end_time, color='red', alpha=0.3)
        plt.draw()
        self.segments.append(self.current_segment)  # Add current_segment to segments list
        print(f"Segment marked as 'discard': {self.current_segment.file_name}, {self.current_segment.start_time}s to {self.current_segment.end_time}s")

    def edit_segment(self, segment):
        # Open a dialog box to edit the start and end time of the segment
        edit_window = Tk()
        edit_window.title("Edit Segment")
        
        start_time_label = Label(edit_window, text="Start Time (s):")
        start_time_label.pack()
        start_time_entry = Entry(edit_window)
        start_time_entry.insert(END, str(segment.start_time))
        start_time_entry.pack()

        end_time_label = Label(edit_window, text="End Time (s):")
        end_time_label.pack()
        end_time_entry = Entry(edit_window)
        end_time_entry.insert(END, str(segment.end_time))
        end_time_entry.pack()

        def save_changes():
            new_start_time = float(start_time_entry.get())
            new_end_time = float(end_time_entry.get())
            segment.start_time = new_start_time
            segment.end_time = new_end_time
            segment.duration = new_end_time - new_start_time
            self.filter_trace()
            plt.draw()
            edit_window.destroy()
            edit_window.quit()
        
        save_button = Button(edit_window, text="Save", command=save_changes)
        save_button.pack()
        
        edit_window.mainloop()
    
    def delete_segment(self, segment):
        self.segments.remove(segment)
        self.filter_trace()
        plt.draw()

    def toggle_segmentation(self):
        global segmentation_enabled
        segmentation_enabled = not segmentation_enabled

        if segmentation_enabled:
            toggle_segmentation_button.config(text='Disable Segmentation')
        else:
            toggle_segmentation_button.config(text='Enable Segmentation')
    
    def on_double_click(self, event):
        if event.button == 1 and event.dblclick:  # Left button (double-click)
            for segment in self.segments:
                if segment.start_time <= event.xdata <= segment.end_time:
                    self.edit_segment(segment)
                    break
    
    def on_right_click(self, event):
        if event.button == 3:  # Right button (click)
            for segment in self.segments:
                if segment.start_time <= event.xdata <= segment.end_time:
                    self.delete_segment(segment)
                    break

    def toggle_peaks(self):
        global peaks_enabled
        peaks_enabled = not peaks_enabled

        if peaks_enabled:
            toggle_peaks_button.config(text='Disable Peaks Selection')
        else:
            toggle_peaks_button.config(text='Enable Peaks Selection')

    def setPeak(self, event):
        if not peaks_enabled:
            return

        self.current_peak = Peak(self.trace)
        self.current_peak.on_line_click(event)
        self.peaks.append(self.current_peak)
        self.trace.peaks.append(self.current_peak)

        for peak in self.trace.peaks:
            ax.plot(peak.peak_x, peak.peak_y, 'ro')
            fig.canvas.draw()
    


        

def EphysApp():
    global canvas, ax, fig,  filter_button, segmentation_enabled, segments, peaks, traces, toggle_segmentation_button, file_menu, current_file_menu_var, root, peaks_enabled, toggle_peaks_button
    root= Tk()
    root.title("Ephys Analysis Tool")
    file_names = []
    segments= []
    peaks = []
    traces= []
    #current_file_index = 0
    #current_file_name = ""
    FileMan= FileManager(file_names, segments, peaks, traces)
    
    def on_closing():
        root.quit()
        root.destroy()

    root.protocol("WM_DELETE_WINDOW", on_closing)

    # Add figure canvas
    fig, ax = plt.subplots(facecolor="gainsboro")
    fig.tight_layout()

    canvas = FigureCanvasTkAgg(fig, master=root)
    canvas.draw()
    canvas.get_tk_widget().pack(side='bottom', fill='both', expand=True)

    # Add toolbar
    toolbar_frame = Frame(root)
    toolbar_frame.pack(side='top', fill='x')
    toolbar = NavigationToolbar2Tk(canvas, root)
    toolbar.update()
    canvas.get_tk_widget().pack(side='bottom', fill='both', expand=True)

    # Add dropdown menu of file list
    current_file_menu_var = StringVar(root)
    
    file_menu = OptionMenu(root, current_file_menu_var, *file_names)
    file_menu.pack(side="top", anchor='nw')

    # Call the update_dropdown_menu function
    FileMan.update_dropdown_menu()

    # Load files
    add_files_button = Button(master=root, command= FileMan.add_files, text='Add Files')
    add_files_button.pack(side="left")

    # Add functionality to switch files
    previous_button = Button(master=root, command= FileMan.previous_file, text='Previous')
    previous_button.pack(side="left")

    next_button = Button(master=root, command= FileMan.next_file, text='Next')
    next_button.pack(side="left")

    span_selector = SpanSelector(ax, FileMan.onselect, 'horizontal', useblit=True)

    # Add the Filter button to the toolbar
    filter_button = Button(master=root, command= FileMan.filter_trace, text='Filter')
    filter_button.pack(side="left")
    filter_button.is_filtered = False

    # Enable and disable segmentation

    segmentation_enabled = False
    toggle_segmentation_button = Button(master=root, command=FileMan.toggle_segmentation, text='Enable Segmentation')
    toggle_segmentation_button.pack(side="left")

    # Mark events as bursts or discard
    burst_button = Button(master=root, command=FileMan.mark_as_burst, text='Mark Burst')
    burst_button.pack(side="left")

    discard_button = Button(master=root, command=FileMan.mark_as_discard, text='Mark Discard')
    discard_button.pack(side="left")

    #Enable and disable peak selection
    peaks_enabled = False
    toggle_peaks_button = Button(master=root, command=FileMan.toggle_peaks, text='Enable Peak Selection')
    toggle_peaks_button.pack(side="left")

    # Export table
    save_button = Button(master=root, command=FileMan.save_table, text='Save Table')
    save_button.pack(side="left")

    # save and load segments
    Button(root, text="Save Segments", command=FileMan.save_segments).pack(side="left", padx=5, pady=5)
    Button(root, text="Load Segments", command=FileMan.load_segments).pack(side="left", padx=5, pady=5)

    button_frame = Frame(root)
    button_frame.pack(side='top', fill='x')
    
    FileMan.enable_zoom_and_scroll(canvas)

    # Add save and load data buttons
    save_data_button = Button(master=root, command=FileMan.save_traces, text='Save Trace Data')
    save_data_button.pack(side="left")

    load_data_button = Button(master=root, command=FileMan.load_traces, text='Load Trace Data')
    load_data_button.pack(side="left")
    
    #calculatebase = Button(master=root, command= FileMan.calculate_baseline, text='Calculate Baseline')
    #calculatebase.pack(side="left")

    canvas.mpl_connect('button_press_event', FileMan.on_double_click)
    canvas.mpl_connect('button_press_event', FileMan.on_right_click)
    canvas.mpl_connect('pick_event', lambda event: FileMan.setPeak(event))

    root.mainloop()

EphysApp()