import sys
import os
import numpy as np
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QListWidget, QSpinBox,QDoubleSpinBox, QLabel, QCheckBox, QFileDialog, QTabWidget, QSizePolicy, QGridLayout, QSpinBox, QMessageBox, QLineEdit, QFormLayout, QHBoxLayout, QStyleFactory,)
from PySide6.QtCore import Qt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from functools import partial
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter
from PySide6.QtGui import QPalette, QColor, QFont
import math

class MplCanvas(FigureCanvas):
    def __init__(self, width=5, height=4, dpi=300):
        fig = Figure(figsize=(width, height), dpi=dpi)
        self.axes = fig.add_subplot(111)
        super().__init__(fig)

class PlottingApplication(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("SD Frequency and multi-plots")
        self.histogram_bins = 100  # Default number of bins
        self.scatter_marker_size = 3  # Default scatter marker size
        self.global_limits = {}
        self.initUI()

    def initUI(self):
        # Main widget and layout
        mainWidget = QWidget()
        self.setCentralWidget(mainWidget)
        mainLayout = QHBoxLayout(mainWidget)

        # Left panel for file selection and plot options
        self.leftPanel = QTabWidget()  # Assign to self to make it accessible class-wide
        fileSelectionTab = QWidget()
        self.plotOptionsTab = QWidget()
        self.leftPanel.addTab(fileSelectionTab, "Files")
        self.leftPanel.addTab(self.plotOptionsTab, "Options")
        
        # Right panel for displaying plots
        self.rightPanel = QTabWidget()

        mainLayout.addWidget(self.leftPanel, 25)
        mainLayout.addWidget(self.rightPanel, 75)

        self.setupFileSelectionTab(fileSelectionTab)
        self.setupPlotOptionsTab(self.plotOptionsTab)
    


    def calculateGridSize(self, numberOfPlots):
        # Calculate the square root of the number of plots to get an approximation
        # of the grid size
        gridSize = math.ceil(math.sqrt(numberOfPlots))
        
        # For a perfectly square layout, both the number of rows and columns
        # would be equal to gridSize. However, depending on your exact needs,
        # you might want to adjust this to have more rows than columns or vice versa.
        nrows = gridSize
        ncols = gridSize
        
        # Adjust ncols if there are too many empty spots
        if gridSize * (gridSize - 1) >= numberOfPlots:
            ncols -= 1
        
        return nrows, ncols



    def setupFileSelectionTab(self, tab):
        layout = QVBoxLayout(tab)  # Create a layout for this tab

        headerLabel = QLabel("SD Frequency and multi-plots\nshankar.dutt@anu.edu.au")
        headerLabel.setWordWrap(True)
        headerLabel.setAlignment(Qt.AlignmentFlag.AlignCenter)  # Update alignment flag
        headerLabel.setFont(QFont('Arial', 18, QFont.Weight.Bold))  # Update font weight
        headerLabel.setStyleSheet("QLabel { color : white; margin-top: 20px; margin-bottom: 20px; }")

        # Update the stylesheet for QPushButton for PyQt6
        app.setStyleSheet("""
        QPushButton {
            background-color: #2b5b84;
            color: white;
            border-radius: 5px;
            padding: 10px;
        }
        QPushButton:hover {
            background-color: #1d4560;
        }
        """)

        layout.addWidget(headerLabel)

        selectFolderBtn = QPushButton("Select Folder")
        selectFolderBtn.clicked.connect(self.selectFolder)
        layout.addWidget(selectFolderBtn)

        self.fileListWidget = QListWidget()
        self.fileListWidget.setSelectionMode(QListWidget.SelectionMode.MultiSelection)
        layout.addWidget(self.fileListWidget)

        self.includeSubfoldersCheckbox = QCheckBox("Include Subfolders")
        layout.addWidget(self.includeSubfoldersCheckbox)

        plotButton = QPushButton("Plot")
        plotButton.clicked.connect(self.plotSelectedFiles)
        layout.addWidget(plotButton)

        # Initialize plot options and add them to the layout
        self.plotOptions = {
            'Histogram of dI': QCheckBox('Histogram of dI'),
            'Histogram of dt': QCheckBox('Histogram of dt'),
            'Histogram of Area': QCheckBox('Histogram of Area'),
            'Scatter plot dI vs dt': QCheckBox('Scatter plot dI vs dt'),
            'Scatter plot area vs dt': QCheckBox('Scatter plot area vs dt'),
            'Histogram of dt (log)': QCheckBox('Histogram of dt (log)'),
            'Scatter plot dI vs dt (log)': QCheckBox('Scatter plot dI vs dt (log)'),
            'Frequency plot': QCheckBox('Frequency plot')
        }
        for option in self.plotOptions.values():
            layout.addWidget(option)

        syncZoomButton = QPushButton("Sync Zoom")
        syncZoomButton.clicked.connect(self.syncZoomAcrossPlots)
        layout.addWidget(syncZoomButton)
        self.selectedFilesOrder = []
        
        # Connect the selection changed signal
        self.fileListWidget.itemSelectionChanged.connect(self.updateSelectionOrder)


    def setupPlotOptionsTab(self, tab):
        layout = QVBoxLayout(tab)
        
        # Use QFormLayout for a neat alignment of labels and fields including buttons
        formLayout = QFormLayout()

        # Histogram Bins
        self.histogramBinsSpinBox = QSpinBox()
        self.histogramBinsSpinBox.setMinimum(1)
        self.histogramBinsSpinBox.setMaximum(1000)
        self.histogramBinsSpinBox.setValue(100)
        formLayout.addRow("Histogram Bins:", self.histogramBinsSpinBox)

        # Marker Size
        self.markerSizeSpinBox = QSpinBox()
        self.markerSizeSpinBox.setMinimum(1)
        self.markerSizeSpinBox.setMaximum(100)
        self.markerSizeSpinBox.setValue(3)
        formLayout.addRow("Marker Size:", self.markerSizeSpinBox)

        # Linear Fit
        self.linearFitCheckBox = QCheckBox("Enable Linear Fit in the Frequency Plots")
        formLayout.addRow(self.linearFitCheckBox)

        # Fit between two points
        self.fitBetweenPointsCheckBox = QCheckBox("Fit Between Two Points")
        formLayout.addRow(self.fitBetweenPointsCheckBox)

        # Start point for fit
        self.fitStartPointSpinBox = QDoubleSpinBox()
        self.fitStartPointSpinBox.setMinimum(-np.inf)
        self.fitStartPointSpinBox.setMaximum(np.inf)
        self.fitStartPointSpinBox.setValue(0)

        # End point for fit
        self.fitEndPointSpinBox = QDoubleSpinBox()
        self.fitEndPointSpinBox.setMinimum(-np.inf)
        self.fitEndPointSpinBox.setMaximum(np.inf)
        self.fitEndPointSpinBox.setValue(1)

        self.fitRangeLabel = QLabel("Fit Range:")
        self.fitRangeLayout = QHBoxLayout()
        self.fitRangeLayout.addWidget(self.fitRangeLabel)
        self.fitRangeLayout.addWidget(self.fitStartPointSpinBox)
        self.fitRangeLayout.addWidget(self.fitEndPointSpinBox)
        formLayout.addRow(self.fitRangeLayout)

        self.fitBetweenPointsCheckBox.stateChanged.connect(self.toggleFitRangeVisibility)
        self.toggleFitRangeVisibility(False)

        self.overlayPlotsCheckBox = QCheckBox("Overlay Plots")
        formLayout.addRow(self.overlayPlotsCheckBox)

        self.alphaSpin = QDoubleSpinBox()
        self.alphaSpin.setMinimum(0.0)
        self.alphaSpin.setMaximum(1.0)
        self.alphaSpin.setSingleStep(0.1)
        self.alphaSpin.setValue(0.5)  # Default alpha value
        formLayout.addRow("Alpha:", self.alphaSpin)

        # Save directory input and button
        self.saveDirectoryLineEdit = QLineEdit()
        self.saveDirectoryButton = QPushButton("Select Save Folder")
        self.saveDirectoryButton.clicked.connect(self.selectSaveFolder)
        saveFolderLayout = QHBoxLayout()
        saveFolderLayout.addWidget(self.saveDirectoryLineEdit)
        saveFolderLayout.addWidget(self.saveDirectoryButton)
        formLayout.addRow("Save Figures To:", saveFolderLayout)

        # Update plots button
        updatePlotsButton = QPushButton("Update Plots")
        updatePlotsButton.clicked.connect(self.updatePlots)
        formLayout.addRow("", updatePlotsButton)  # Add an empty label for alignment

        # Save figures button
        saveFiguresButton = QPushButton("Save Figures")
        saveFiguresButton.clicked.connect(self.saveAllFigures)
        formLayout.addRow("", saveFiguresButton)  # Add an empty label for alignment

        layout.addLayout(formLayout)

        # Adjust overall layout settings
        layout.setSpacing(10)  # Adjust the spacing as needed
        tab.setLayout(layout)

    def toggleFitRangeVisibility(self, visible):
        self.fitRangeLabel.setVisible(visible)
        self.fitStartPointSpinBox.setVisible(visible)
        self.fitEndPointSpinBox.setVisible(visible)


    def selectSaveFolder(self):
        folder = QFileDialog.getExistingDirectory(self, "Select Folder")
        if folder:
            self.saveDirectoryLineEdit.setText(folder)

    def saveAllFigures(self):
        saveFolder = self.saveDirectoryLineEdit.text().strip()
        if not saveFolder:
            QMessageBox.information(self, "Info", "Please select a folder to save the figures.")
            return

        for i in range(self.rightPanel.count()):
            tab = self.rightPanel.widget(i)
            if hasattr(tab, "property"):
                tabCanvases = tab.property("canvases")
                if tabCanvases:
                    tabFolder = os.path.join(saveFolder, self.rightPanel.tabText(i))
                    if not os.path.exists(tabFolder):
                        os.makedirs(tabFolder)
                    for j, canvas in enumerate(tabCanvases):
                        figure = canvas.figure
                        figure.savefig(os.path.join(tabFolder, f"figure_{j+1}.png"), dpi=300, bbox_inches='tight')
        QMessageBox.information(self, "Info", "Figures saved successfully.")


           
    def selectFolder(self):
        folderPath = QFileDialog.getExistingDirectory(self, "Select Folder")
        if folderPath:
            self.populateFileList(folderPath,self.includeSubfoldersCheckbox.isChecked() )

    def populateFileList(self, folderPath, includeSubfolders):
        self.fileListWidget.clear()
        if includeSubfolders:
            for root, dirs, files in os.walk(folderPath):
                for file in files:
                    if file.endswith('.dataset.npz'):
                        fullPath = os.path.join(root, file)
                        self.fileListWidget.addItem(fullPath)
        else:
            for file in os.listdir(folderPath):
                if file.endswith('.dataset.npz'):
                    fullPath = os.path.join(folderPath, file)
                    self.fileListWidget.addItem(fullPath)
    
    def updateSelectionOrder(self):
        # Get the current selection
        currentSelection = {self.fileListWidget.item(i).text() for i in range(self.fileListWidget.count()) if self.fileListWidget.item(i).isSelected()}
        
        # Filter out any deselected items
        self.selectedFilesOrder = [f for f in self.selectedFilesOrder if f in currentSelection]
        
        # Add new selections in order
        for i in range(self.fileListWidget.count()):
            itemText = self.fileListWidget.item(i).text()
            if self.fileListWidget.item(i).isSelected() and itemText not in self.selectedFilesOrder:
                self.selectedFilesOrder.append(itemText)
    

    def updatePlots(self):
        # Update the class attributes with the new values from the UI
        self.histogram_bins = self.histogramBinsSpinBox.value()
        self.scatter_marker_size = self.markerSizeSpinBox.value()

        # Re-plot all selected files with the new settings
        self.plotSelectedFiles()


    @staticmethod
    def create_joint_plot(self, fig, data_x, data_y, min_x, max_x, min_y, max_y, plotType, alpha):
        # Create grid layout for the plots
        gs = fig.add_gridspec(4, 4)
        
        # Main scatter plot with specified axis limits
        ax_scatter = fig.add_subplot(gs[1:4, 0:3])
        if data_x is not None and data_y is not None:
            ax_scatter.scatter(data_x, data_y, s=int(self.markerSizeSpinBox.value()), alpha=alpha)
        ax_scatter.set_xlim(min_x, max_x)
        ax_scatter.set_ylim(min_y, max_y)
        
        # Histogram for X distribution
        ax_histx = fig.add_subplot(gs[0, 0:3], sharex=ax_scatter)
        if data_x is not None:
            ax_histx.hist(data_x, bins=self.histogramBinsSpinBox.value())
        ax_histx.set_ylabel(' ')
        
        # Histogram for Y distribution
        ax_histy = fig.add_subplot(gs[1:4, 3], sharey=ax_scatter)
        if data_y is not None:
            ax_histy.hist(data_y, bins=self.histogramBinsSpinBox.value(), orientation='horizontal')
        ax_histy.set_xlabel(' ')
        
        # Hide tick labels for histograms to avoid overlap
        plt.setp(ax_histx.get_xticklabels(), visible=False)
        plt.setp(ax_histy.get_yticklabels(), visible=False)

        ax_scatter.xaxis.set_major_formatter(ScalarFormatter(useMathText=True))
        ax_scatter.yaxis.set_major_formatter(ScalarFormatter(useMathText=True))
        if 'dI vs dt' in plotType and 'log' not in plotType:
            ax_scatter.set_xlabel('Δt (ms)')
            ax_scatter.set_ylabel('ΔI (nA)')
        elif 'area vs dt' in plotType:
            ax_scatter.set_xlabel('Δt (ms)')
            ax_scatter.set_ylabel('Area (C)')
        elif 'dI vs dt (log)' in plotType:
            ax_scatter.set_xlabel('Δt (ms) [log]')
            ax_scatter.set_ylabel('ΔI (nA)')

        ax_scatter.ticklabel_format(style='sci', axis='both', scilimits=(0,0))
        ax_histx.ticklabel_format(style='sci', axis='both', scilimits=(0,0))
        ax_histy.ticklabel_format(style='sci', axis='both', scilimits=(0,0))
        fig.tight_layout()

        # Return the axes dictionary
        return {'ax_scatter': ax_scatter, 'ax_histx': ax_histx, 'ax_histy': ax_histy}
    
    def plotSelectedFiles(self):
        '''selectedFiles = [self.fileListWidget.item(i).text() for i in range(self.fileListWidget.count()) if self.fileListWidget.item(i).isSelected()]
        plotOptions = {key: option.isChecked() for key, option in self.plotOptions.items() if option.isChecked()}'''

        plotOptions = {key: option.isChecked() for key, option in self.plotOptions.items() if option.isChecked()}
        self.rightPanel.clear()

        self.rightPanel.clear()

        if not self.selectedFilesOrder or not plotOptions:
            return

        data = self.prepareData(self.selectedFilesOrder, plotOptions)

        for plotType, plotData in data.items():
            tab = QWidget()
            layout = QVBoxLayout(tab)
            plotLayout = QGridLayout()
            alpha = self.alphaSpin.value()
            tabCanvases = []

            if self.overlayPlotsCheckBox.isChecked():
                # Create a single figure and canvas for overlaid plots
                fig = Figure(figsize=(10, 5))
                canvas = FigureCanvas(fig)
                tabCanvases.append(canvas)

                min_x, max_x, min_y, max_y = self.calculateGlobalAxisLimits(plotData, plotType)

                if 'Scatter plot dI vs dt' in plotType and 'log' not in plotType:
                    axes_dict = self.create_joint_plot(self, fig, None, None, min_x, max_x, min_y, max_y, 'dI vs dt', alpha)
                    ax_scatter = axes_dict['ax_scatter']
                    ax_histx = axes_dict['ax_histx']
                    ax_histy = axes_dict['ax_histy']

                    for filePath, values in plotData.items():
                        alpha = self.alphaSpin.value()
                        label = os.path.basename(filePath)

                        ax_scatter.scatter(values[0], values[1], s=int(self.markerSizeSpinBox.value()), alpha=alpha, label=label)
                        ax_histx.hist(values[0], bins=self.histogramBinsSpinBox.value(), alpha=alpha)
                        ax_histy.hist(values[1], bins=self.histogramBinsSpinBox.value(), orientation='horizontal', alpha=alpha)
                    fig.tight_layout()
                    ax_scatter.legend()

                elif 'Scatter plot area vs dt' in plotType:
                    axes_dict = self.create_joint_plot(self, fig, None, None, min_x, max_x, min_y, max_y, 'area vs dt', alpha)
                    ax_scatter = axes_dict['ax_scatter']
                    ax_histx = axes_dict['ax_histx']
                    ax_histy = axes_dict['ax_histy']

                    for filePath, values in plotData.items():
                        alpha = self.alphaSpin.value()
                        label = os.path.basename(filePath)

                        ax_scatter.scatter(values[0], values[1], s=int(self.markerSizeSpinBox.value()), alpha=alpha, label=label)
                        ax_histx.hist(values[0], bins=self.histogramBinsSpinBox.value(), alpha=alpha)
                        ax_histy.hist(values[1], bins=self.histogramBinsSpinBox.value(), orientation='horizontal', alpha=alpha)
                    fig.tight_layout()
                    ax_scatter.legend()

                elif 'Scatter plot dI vs dt (log)' in plotType:
                    axes_dict = self.create_joint_plot(self, fig, None, None, np.log(min_x), np.log(max_x), min_y, max_y, 'dI vs dt (log)', alpha)
                    ax_scatter = axes_dict['ax_scatter']
                    ax_histx = axes_dict['ax_histx']
                    ax_histy = axes_dict['ax_histy']

                    for filePath, values in plotData.items():
                        alpha = self.alphaSpin.value()
                        label = os.path.basename(filePath)

                        ax_scatter.scatter(np.log(values[0]), values[1], s=int(self.markerSizeSpinBox.value()), alpha=alpha, label=label)
                        ax_histx.hist(np.log(values[0]), bins=self.histogramBinsSpinBox.value(), alpha=alpha)
                        ax_histy.hist(values[1], bins=self.histogramBinsSpinBox.value(), orientation='horizontal', alpha=alpha)
                    fig.tight_layout()
                    ax_scatter.legend()
                    

                elif 'Histogram of dt (log)' in plotType:
                    ax = fig.add_subplot(111)

                    for filePath, values in plotData.items():
                        alpha = self.alphaSpin.value()
                        label = os.path.basename(filePath)

                        log_values = np.log(values[values > 0])
                        ax.hist(log_values, bins=self.histogramBinsSpinBox.value(), range=(np.log(min_x), np.log(max_x)), alpha=alpha, label=label)
                    fig.tight_layout()
                    ax.legend()

                elif 'Frequency plot' in plotType:
                    ax = fig.add_subplot(111)

                    for filePath, values in plotData.items():
                        alpha = self.alphaSpin.value()
                        label = os.path.basename(filePath)

                        ax.plot(values[0], values[1], 'o', markersize=self.scatter_marker_size, alpha=alpha, label=label)

                        if self.linearFitCheckBox.isChecked():
                            if self.fitBetweenPointsCheckBox.isChecked():
                                start_point = self.fitStartPointSpinBox.value()
                                end_point = self.fitEndPointSpinBox.value()
                                mask = (values[0] >= start_point) & (values[0] <= end_point)
                                x_fit = values[0][mask]
                                y_fit = values[1][mask]
                                slope, intercept = np.polyfit(x_fit, y_fit, 1)
                            else:
                                slope, intercept = np.polyfit(values[0], values[1], 1)
                                x_fit = np.array([min_x, max_x])
                                y_fit = slope * x_fit + intercept

                            ax.plot(x_fit, slope * x_fit + intercept, '--', alpha=alpha)

                    ax.legend()

                else:  # Regular histogram
                    ax = fig.add_subplot(111)

                    for filePath, values in plotData.items():
                        alpha = self.alphaSpin.value()
                        label = os.path.basename(filePath)

                        ax.hist(values, bins=self.histogramBinsSpinBox.value(), range=(min_x, max_x), alpha=alpha, label=label)

                    ax.legend()

                fig.suptitle("Overlaid Plots", fontsize=12)
                fig.tight_layout()
                fig.tight_layout()
                fig.tight_layout()
                fig.tight_layout()
                fig.tight_layout()
                plotLayout.addWidget(canvas, 0, 0)
                

                
            else:
                row = col = 0
                max_col = 1
                
                min_x, max_x, min_y, max_y = self.calculateGlobalAxisLimits(plotData, plotType)

                for filePath, values in plotData.items():
                    fig = Figure(figsize=(5, 5))
                    canvas = FigureCanvas(fig)

                    if 'Scatter plot dI vs dt' in plotType and 'log' not in plotType:
                        axes_dict = self.create_joint_plot(self, fig, values[0], values[1], min_x, max_x, min_y, max_y, 'dI vs dt', alpha=1.0)
                        canvas.axes_dict = axes_dict
                        canvas.plotType = 'Scatter plot dI vs dt'
                        title = f"{os.path.basename(filePath)}\nEvents: {len(values[0])}"

                    elif 'Scatter plot area vs dt' in plotType:
                        axes_dict = self.create_joint_plot(self, fig, values[0], values[1], min_x, max_x, min_y, max_y, 'area vs dt', alpha=1.0)
                        canvas.axes_dict = axes_dict
                        canvas.plotType = 'Scatter plot area vs dt'
                        title = f"{os.path.basename(filePath)}\nEvents: {len(values[0])}"

                    elif 'Scatter plot dI vs dt (log)' in plotType:
                        axes_dict = self.create_joint_plot(self, fig, np.log(values[0]), values[1], np.log(min_x), np.log(max_x), min_y, max_y, 'dI vs dt (log)', alpha=1.0)
                        canvas.axes_dict = axes_dict
                        canvas.plotType = 'Scatter plot dI vs dt (log)'
                        title = f"{os.path.basename(filePath)}\nEvents: {len(values[0])}"

                    elif 'Histogram of dt (log)' in plotType:
                        log_values = np.log(values[values > 0])
                        axes_dict = self.create_histogram(fig, log_values, np.log(min_x), np.log(max_x), 'dt (log)', alpha)
                        canvas.axes_dict = axes_dict
                        canvas.plotType = 'Histogram of dt (log)'
                        title = f"{os.path.basename(filePath)}\nEvents: {len(log_values)}"

                    elif 'Frequency plot' in plotType:
                        self.create_frequency_plot(fig, values[0], values[1], min_x, max_x, min_y, max_y, alpha=1.0)
                        canvas.plotType = 'Frequency plot'
                        title = f"{os.path.basename(filePath)}\nEvents: {len(values[0])}"
                        if self.linearFitCheckBox.isChecked():
                            if self.fitBetweenPointsCheckBox.isChecked():
                                start_point = self.fitStartPointSpinBox.value()
                                end_point = self.fitEndPointSpinBox.value()
                                mask = (values[0] >= start_point) & (values[0] <= end_point)
                                x_fit = values[0][mask]
                                y_fit = values[1][mask]
                                slope, intercept = np.polyfit(x_fit, y_fit, 1)
                            else:
                                slope, intercept = np.polyfit(values[0], values[1], 1)
                                x_fit = np.array([min_x, max_x])
                                y_fit = slope * x_fit + intercept

                            ax = fig.axes[0]
                            ax.plot(x_fit, slope * x_fit + intercept, 'r--', label=f'Fit: {slope:.2f} Events/s, Intercept: {intercept:.2f}')
                            ax.legend()
                            title += f"\nFit: {slope:.2f} Events/s, Intercept: {intercept:.2f}"

                    else:
                        # Regular histogram
                        axes_dict = self.create_histogram(fig, values, min_x, max_x, plotType, alpha=1.0)
                        canvas.axes_dict = axes_dict
                        canvas.plotType = plotType
                        title = f"{os.path.basename(filePath)}\nEvents: {len(values)}"

                    fig.suptitle(title, fontsize=8)
                    
                    plotLayout.addWidget(canvas, row, col)
                    tabCanvases.append(canvas)

                    col += 1
                    if col > max_col:
                        col = 0
                        row += 1
                    
                    # Adjust layout
                    fig.tight_layout()

            layout.addLayout(plotLayout)

            if tabCanvases:
                toolbar = NavigationToolbar(tabCanvases[0], self)
                layout.addWidget(toolbar)

            fig.tight_layout()

            resetButton = QPushButton("Reset Zoom")
            resetButton.clicked.connect(self.resetZoom)
            layout.addWidget(resetButton)

            self.rightPanel.addTab(tab, plotType)
            tab.setProperty("canvases", tabCanvases)


    def create_frequency_plot(self, fig, times, event_numbers, min_x, max_x, min_y, max_y, alpha):
        ax = fig.add_subplot(111)
        ax.plot(times, event_numbers, 'o', markersize=self.scatter_marker_size, alpha = alpha)
        ax.set_xlim(min_x, max_x)
        #ax.set_ylim(min_y, max_y)
        ax.set_xlabel('Time (s)')
        ax.set_ylabel('Event Number')
        ax.ticklabel_format(style='sci', axis='both', scilimits=(0,0))
        fig.tight_layout()
    

    def syncZoomAcrossPlots(self):
        currentTab = self.rightPanel.currentWidget()
        if hasattr(currentTab, "property"):
            tabCanvases = currentTab.property("canvases")
            if tabCanvases and len(tabCanvases) > 0:
                # Assuming the first canvas' figure's first axes as the reference
                referenceAxes = tabCanvases[0].figure.axes[0]
                xlim = referenceAxes.get_xlim()
                ylim = referenceAxes.get_ylim()
                
                for canvas in tabCanvases:
                    plotType = getattr(canvas, 'plotType', '')
                    if 'Scatter plot' in plotType:
                        if hasattr(canvas, 'axes_dict'):
                            canvas.axes_dict['ax_scatter'].set_xlim(xlim)
                            canvas.axes_dict['ax_scatter'].set_ylim(ylim)
                            canvas.axes_dict['ax_histx'].set_xlim(xlim)
                            canvas.axes_dict['ax_histy'].set_ylim(ylim)  # Reverse y-limits for horizontal histogram
                    elif 'Histogram' in plotType or 'Frequency plot' in plotType:
                        # For histograms and frequency plots, synchronize only the X-axis
                        ax = canvas.figure.axes[0]
                        ax.set_xlim(xlim)
                    else:
                        # For other plots, synchronize both axes
                        for ax in canvas.figure.axes:
                            ax.set_xlim(xlim)
                            ax.set_ylim(ylim)

                    canvas.draw_idle()

  

    def create_histogram(self, fig, values, min_x, max_x, plotType, alpha):
        ax = fig.add_subplot(111)
        ax.hist(values, bins=self.histogramBinsSpinBox.value(), range=(min_x, max_x),  alpha = alpha)
        ax.xaxis.set_major_formatter(ScalarFormatter(useMathText=True))
        if 'dI' in plotType:
            ax.set_xlabel('ΔI (nA)')
        elif 'dt' in plotType and 'log' not in plotType:
            ax.set_xlabel('Δt (s)')
        elif 'dt (log)' in plotType:
            ax.set_xlabel('log(Δt (ms))')
        ax.set_ylabel('Count')
        ax.ticklabel_format(style='sci', axis='both', scilimits=(0,0))
        fig.tight_layout()

        # Return the axis dictionary
        return {'ax_hist': ax}


    def resetZoom(self):
        currentTab = self.rightPanel.currentWidget()
        if currentTab is None:
            return

        plotType = self.rightPanel.tabText(self.rightPanel.indexOf(currentTab))

        if plotType in self.global_limits:
            min_x, max_x, min_y, max_y = self.global_limits[plotType]

            tabCanvases = currentTab.property("canvases")
            if tabCanvases:
                for canvas in tabCanvases:
                    if plotType == 'Scatter plot dI vs dt' or plotType == 'Scatter plot area vs dt' or plotType == 'Scatter plot dI vs dt (log)':
                        if hasattr(canvas, 'axes_dict'):
                            canvas.axes_dict['ax_scatter'].set_xlim(min_x, max_x)
                            canvas.axes_dict['ax_scatter'].set_ylim(min_y, max_y)
                            canvas.axes_dict['ax_histx'].set_xlim(min_x, max_x)
                            canvas.axes_dict['ax_histy'].set_ylim(min_y, max_y)
                    elif plotType == 'Frequency plot':
                        ax = canvas.figure.axes[0]
                        ax.set_xlim(min_x, max_x)
                        #ax.set_ylim(min_y, max_y)
                    elif plotType in ['Histogram of dI', 'Histogram of dt', 'Histogram of Area', 'Histogram of dt (log)']:
                        ax = canvas.figure.axes[0]
                        ax.set_xlim(min_x, max_x)
                    else:
                        print(f"Unsupported plot type: {plotType}")

                    canvas.draw_idle()

    
    def prepareData(self, selectedFiles, plotOptions):
        data = {}
        for filePath in selectedFiles:
            fileData = np.load(filePath)
            X = fileData['X']
            
            for option in plotOptions:
                if option not in data:
                    data[option] = {}
                
                if 'Histogram' in option:
                    index = 0 if 'dI' in option else 3 if 'Area' in option else 4  # Assuming the order dI, dt, Area
                    values = X[:, index].astype(float)
                    # Remove NaN values
                    values = values[~np.isnan(values)]
                    data[option][filePath] = values
                    #print(f"Data type for {option}: {type(values)}")
                elif 'Histogram of dt (log)' in option:
                    # Assuming dt values are in the second column (index 1)
                    dt_values = fileData['X'][:, 4]
                    # Filter out non-positive values to avoid log(0) and log(negative)
                    positive_dt_values = dt_values[dt_values > 0]
                    # Apply logarithmic transformation
                    log_dt_values = np.log(positive_dt_values*1e3)
                    # Remove NaN values
                    log_dt_values = log_dt_values[~np.isnan(log_dt_values)]
                    data['Histogram of dt (log)'][filePath] = log_dt_values.astype(float)
                    #print(f"Data type for Histogram of dt (log): {type(log_dt_values)}")
                elif 'Frequency plot' in option:
                    event_numbers = np.arange(1, len(fileData['X']) + 1).astype(float)
                    times = fileData['X'][:, -1].astype(float)
                    # Remove NaN values
                    valid_indices = ~np.isnan(times)
                    event_numbers = event_numbers[valid_indices]
                    times = times[valid_indices]
                    data[option][filePath] = (times, event_numbers)
                    #print(f"Data type for Frequency plot: {type(times)}, {type(event_numbers)}")
                else:  # Scatter plots
                    x_index, y_index = (4, 0) if 'dI vs dt' in option else (4, 3)  # Assuming dI vs dt, Area vs dt
                    x_values = X[:, x_index]*1e3
                    y_values = X[:, y_index]
                    # Remove NaN values
                    valid_indices = ~np.isnan(x_values) & ~np.isnan(y_values)
                    x_values = x_values[valid_indices]
                    y_values = y_values[valid_indices]
                    data[option][filePath] = (x_values, y_values)
                    #print(f"Data type for {option}: {type(x_values)}, {type(y_values)}")
                    
        return data
    
    def calculateGlobalAxisLimits(self, plotData, plotType):
        min_x = min_y = np.inf
        max_x = max_y = -np.inf
        for values in plotData.values():
            if isinstance(values, dict):
                # If values is a dictionary, iterate over its values
                for subvalues in values.values():
                    if 'Histogram' in plotType:
                        curr_min_x = np.min(subvalues)
                        curr_max_x = np.max(subvalues)
                        #print(f"curr_min_x: {curr_min_x}, curr_max_x: {curr_max_x}")
                        if not np.isfinite(curr_min_x) or not np.isfinite(curr_max_x):
                            continue  # Skip this file if min_x or max_x is inf or -inf
                        if curr_min_x < curr_max_x:
                            min_x = min(min_x, curr_min_x)
                            max_x = max(max_x, curr_max_x)
                        else:
                            min_x = min(min_x, curr_min_x)
                            max_x = max(max_x, curr_min_x + 1e-6)  # Adjust max_x to be slightly larger than min_x
                    elif 'Frequency plot' in plotType:
                        x, y = subvalues
                        min_x = min(min_x, np.min(x))
                        max_x = max(max_x, np.max(x))
                        curr_min_y = np.min(y)
                        curr_max_y = np.max(y)
                        if not np.isfinite(curr_min_y) or not np.isfinite(curr_max_y):
                            continue  # Skip this file if min_y or max_y is inf or -inf
                        min_y = min(min_y, curr_min_y)
                        max_y = max(max_y, curr_max_y)
                    else:  # Scatter plots
                        x, y = subvalues
                        min_x = min(min_x, np.min(x))
                        max_x = max(max_x, np.max(x))
                        curr_min_y = np.min(y)
                        curr_max_y = np.max(y)
                        if not np.isfinite(curr_min_y) or not np.isfinite(curr_max_y):
                            continue  # Skip this file if min_y or max_y is inf or -inf
                        min_y = min(min_y, curr_min_y)
                        max_y = max(max_y, curr_max_y)
            else:
                if 'Histogram' in plotType:
                    curr_min_x = np.min(values)
                    curr_max_x = np.max(values)
                    #print(f"curr_min_x: {curr_min_x}, curr_max_x: {curr_max_x}")
                    if not np.isfinite(curr_min_x) or not np.isfinite(curr_max_x):
                        continue  # Skip this file if min_x or max_x is inf or -inf
                    if curr_min_x < curr_max_x:
                        min_x = min(min_x, curr_min_x)
                        max_x = max(max_x, curr_max_x)
                    else:
                        min_x = min(min_x, curr_min_x)
                        max_x = max(max_x, curr_min_x + 1e-6)  # Adjust max_x to be slightly larger than min_x
                elif 'Frequency plot' in plotType:
                    x, y = values
                    min_x = min(min_x, np.min(x))
                    max_x = max(max_x, np.max(x))
                    curr_min_y = np.min(y)
                    curr_max_y = np.max(y)
                    if not np.isfinite(curr_min_y) or not np.isfinite(curr_max_y):
                        continue  # Skip this file if min_y or max_y is inf or -inf
                    min_y = min(min_y, curr_min_y)
                    max_y = max(max_y, curr_max_y)
                else:  # Scatter plots
                    x, y = values
                    min_x = min(min_x, np.min(x))
                    max_x = max(max_x, np.max(x))
                    curr_min_y = np.min(y)
                    curr_max_y = np.max(y)
                    if not np.isfinite(curr_min_y) or not np.isfinite(curr_max_y):
                        continue  # Skip this file if min_y or max_y is inf or -inf
                    min_y = min(min_y, curr_min_y)
                    max_y = max(max_y, curr_max_y)

        if not np.isfinite(min_x) or not np.isfinite(max_x):
            min_x, max_x = 0, 1  # Set default values if min_x or max_x is inf or -inf
        if not np.isfinite(min_y) or not np.isfinite(max_y):
            min_y, max_y = 0, 1  # Set default values if min_y or max_y is inf or -inf

        # Store the global limits in the dictionary
        self.global_limits[plotType] = (min_x, max_x, min_y, max_y)

        return min_x, max_x, min_y, max_y
    

if __name__ == "__main__":

    app = QApplication(sys.argv)
    app.setStyle(QStyleFactory.create('Fusion'))  # Use Fusion or other available styles

    # Customize the palette for a darker, more modern look
    palette = QPalette()
    palette.setColor(QPalette.ColorRole.Window, QColor(53, 53, 53))
    palette.setColor(QPalette.ColorRole.WindowText, Qt.GlobalColor.white)
    palette.setColor(QPalette.ColorRole.Base, QColor(25, 25, 25))
    palette.setColor(QPalette.ColorRole.AlternateBase, QColor(53, 53, 53))
    palette.setColor(QPalette.ColorRole.ToolTipBase, Qt.GlobalColor.white)
    palette.setColor(QPalette.ColorRole.ToolTipText, Qt.GlobalColor.white)
    palette.setColor(QPalette.ColorRole.Text, Qt.GlobalColor.white)
    palette.setColor(QPalette.ColorRole.Button, QColor(53, 53, 53))
    palette.setColor(QPalette.ColorRole.ButtonText, Qt.GlobalColor.white)
    palette.setColor(QPalette.ColorRole.BrightText, Qt.GlobalColor.red)
    palette.setColor(QPalette.ColorRole.Link, QColor(42, 130, 218))
    palette.setColor(QPalette.ColorRole.Highlight, QColor(42, 130, 218))
    palette.setColor(QPalette.ColorRole.HighlightedText, Qt.GlobalColor.black)
    app.setPalette(palette)
    
    mainWindow = PlottingApplication()
    mainWindow.showMaximized()
    sys.exit(app.exec())
