import sys
import numpy as np
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                               QHBoxLayout, QPushButton, QFileDialog, QListWidget, 
                               QMessageBox, QLabel, QSplitter)
from PySide6.QtCore import Qt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure
from util import utiltool
class MatplotlibCanvas(FigureCanvas):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)
        self.setParent(parent)
        self.axes = self.fig.add_subplot(111)
        
    def plot_data(self, data_list, normalized=False):
        self.axes.clear()
        
        for file_path, data in data_list:
            x = data[:, 0]
            y = data[:, 1]
            
            if normalized:
                # Normalize y-axis data using numpy
                y_min = np.min(y)
                y_max = np.max(y)
                if y_max != y_min:  # Avoid division by zero
                    y = (y - y_min) / (y_max - y_min)
                else:
                    y = np.zeros_like(y)
            
            # Extract filename from path for the label
            label = file_path.split('/')[-1].split('\\')[-1]
            self.axes.plot(x, y, label=label)
        
        self.axes.set_xlabel('X Axis')
        if normalized:
            self.axes.set_ylabel('Normalized Y Axis')
        else:
            self.axes.set_ylabel('Y Axis')
        self.axes.set_title('Data Plot' + (' (Normalized)' if normalized else ''))
        self.axes.legend()
        self.axes.grid(True)
        self.fig.tight_layout()
        self.draw()

@utiltool.window_shortcut
class PlotDataWindow(QMainWindow):
    def __init__(self, files=None):
        super().__init__()
        self.setWindowTitle("Data Plotting Tool")
        self.setGeometry(100, 100, 1200, 800)
        
        # Initialize data storage
        self.file_data = {}  # file path -> data array
        
        # Create central widget and main layout
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)
        
        # Create left panel
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        left_layout.setContentsMargins(0, 0, 0, 0)
        
        # File list label
        file_label = QLabel("Loaded Files:")
        left_layout.addWidget(file_label)
        
        # File list
        self.file_list_widget = QListWidget()
        left_layout.addWidget(self.file_list_widget)
        
        # Button layout
        button_layout = QHBoxLayout()
        
        # Load file button
        self.load_button = QPushButton("Load Files")
        self.load_button.clicked.connect(self.load_files)
        button_layout.addWidget(self.load_button)
        
        # Remove file button
        self.remove_button = QPushButton("Remove Selected")
        self.remove_button.clicked.connect(self.remove_selected_file)
        button_layout.addWidget(self.remove_button)
        
        left_layout.addLayout(button_layout)
        
        # Normalization option
        self.normalize_button = QPushButton("Toggle Normalization")
        self.normalize_button.setCheckable(True)
        self.normalize_button.clicked.connect(self.toggle_normalization)
        left_layout.addWidget(self.normalize_button)
        
        # Create right plotting area
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)
        right_layout.setContentsMargins(0, 0, 0, 0)
        
        # Create Matplotlib canvas
        self.canvas = MatplotlibCanvas(self, width=8, height=6, dpi=100)
        
        # Create navigation toolbar
        self.toolbar = NavigationToolbar(self.canvas, self)
        right_layout.addWidget(self.toolbar)
        right_layout.addWidget(self.canvas)
        
        # Use splitter for resizable layout
        splitter = QSplitter(Qt.Horizontal)
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        splitter.setSizes([300, 900])  # Initial size ratio
        
        main_layout.addWidget(splitter)
        
        # Initialize normalization state
        self.normalized = False
        
        # Load files if provided in constructor
        if files:
            self.load_files_from_list(files)
    
    def load_files_from_list(self, file_paths):
        """Load data files from a provided list"""
        if not file_paths:
            return
        for file_path in file_paths:
            self.file_data[file_path] = utiltool.load_data(file_path=file_path)
            # Update file list
            if self.file_list_widget.findItems(file_path, Qt.MatchExactly):
                self.file_list_widget.takeItem(
                    self.file_list_widget.row(
                        self.file_list_widget.findItems(file_path, Qt.MatchExactly)[0]
                    )
                )
            self.file_list_widget.addItem(file_path)

        
        # Update plot
        self.update_plot()
    
    def load_files(self):
        """Load data files through file dialog"""
        file_paths, _ = QFileDialog.getOpenFileNames(
            self, "Select Data Files", "", "Data Files (*.txt *.csv *.dat);;All Files (*)"
        )
        self.load_files_from_list(file_paths)
    
    def remove_selected_file(self):
        """Remove selected file"""
        current_row = self.file_list_widget.currentRow()
        if current_row >= 0:
            item = self.file_list_widget.takeItem(current_row)
            file_path = item.text()
            del self.file_data[file_path]
            self.update_plot()
    
    def toggle_normalization(self):
        """Toggle normalization state"""
        self.normalized = self.normalize_button.isChecked()
        self.update_plot()
    
    def update_plot(self):
        """Update plot display"""
        if not self.file_data:
            self.canvas.axes.clear()
            self.canvas.axes.text(0.5, 0.5, 'Please load data files', 
                                 horizontalalignment='center',
                                 verticalalignment='center',
                                 transform=self.canvas.axes.transAxes)
            self.canvas.axes.set_xlabel('X Axis')
            self.canvas.axes.set_ylabel('Y Axis')
            self.canvas.axes.set_title('Data Plot')
            self.canvas.draw()
            return
        
        # Prepare data list
        data_list = [(path, data) for path, data in self.file_data.items()]
        
        # Plot data
        self.canvas.plot_data(data_list, self.normalized)

def main():
    app = QApplication(sys.argv)
    window = PlotDataWindow()
    window.show()
    sys.exit(app.exec())

if __name__ == "__main__":
    main()