#!/usr/bin/env python3
"""
Graph Widget (PyQt6)
====================
GUI for visualizing telemetry data from SQLite database.

Features:
- Session/run tree selector
- Data series checkboxes
- Plot type selector (overlay, subplots, drift)
- Embedded matplotlib canvas
- Zoom/pan controls
- Export PNG/CSV
- Independent from UART/I2C widgets

Author: Kynsight
Version: 1.0.0
"""

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QPushButton, QComboBox,
    QGroupBox, QTreeWidget, QTreeWidgetItem,
    QCheckBox, QFileDialog, QMessageBox, QSplitter
)
from PyQt6.QtCore import Qt, pyqtSignal
from PyQt6.QtGui import QFont
from typing import Optional

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure

# Import from table_query (data access layer)
from graph_table_query import (
    DataAdapter,
    SQLiteAdapter,
    CSVAdapter,
    TelemetryData,
    get_column_label,
    get_available_columns,
    get_default_columns
)

# Import from graph_core (pure plotting)
from .graph_core import (
    PlotConfig,
    plot_overlay,
    plot_subplots,
    plot_comparison,
    plot_multi_series,
    plot_xy_scatter,
    export_png,
    export_csv
)


class GraphWidget(QWidget):
    """
    Telemetry data visualization widget.
    
    Loads data from SQLite and generates matplotlib plots.
    """
    
    # Signals
    plot_updated = pyqtSignal()
    
    def __init__(self, db_path: str = "./database/ehinge.db", parent=None):
        super().__init__(parent)
        
        self.db_path = db_path
        self.adapter: Optional[DataAdapter] = None  # Use adapter pattern
        
        # Data
        self.sessions = []
        self.selected_runs = []  # List of (session_id, run_no) tuples
        self.loaded_data = []     # List of TelemetryData objects
        
        # Build UI
        self._init_ui()
        
        # Connect to database
        self._connect_database()
    
    def _init_ui(self):
        """Create user interface."""
        layout = QHBoxLayout()
        self.setLayout(layout)
        
        # Left panel: Session/run selection
        left_panel = self._create_selection_panel()
        
        # Right panel: Plot area
        right_panel = self._create_plot_panel()
        
        # Splitter
        splitter = QSplitter(Qt.Orientation.Horizontal)
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        splitter.setStretchFactor(0, 1)  # Left: 25%
        splitter.setStretchFactor(1, 3)  # Right: 75%
        
        layout.addWidget(splitter)
        
        self.setMinimumWidth(1200)
        self.setMinimumHeight(700)
    
    # =========================================================================
    # UI Creation
    # =========================================================================
    
    def _create_selection_panel(self):
        """Create left panel with session/run selection."""
        panel = QWidget()
        layout = QVBoxLayout()
        panel.setLayout(layout)
        
        # 1. Database selector (top)
        db_group = QGroupBox("Database")
        db_layout = QVBoxLayout()
        db_group.setLayout(db_layout)
        
        db_row = QHBoxLayout()
        self.combo_database = QComboBox()
        self.combo_database.setEditable(True)
        self.combo_database.addItem("./database/ehinge.db")
        self.combo_database.setCurrentText(self.db_path)
        db_row.addWidget(self.combo_database)
        
        btn_browse = QPushButton("Browse...")
        btn_browse.clicked.connect(self._browse_database)
        btn_browse.setMaximumWidth(80)
        db_row.addWidget(btn_browse)
        
        db_layout.addLayout(db_row)
        
        btn_connect = QPushButton("🔗 Connect")
        btn_connect.clicked.connect(self._connect_database)
        db_layout.addWidget(btn_connect)
        
        layout.addWidget(db_group)
        
        # 2. Plot controls (right after database!)
        plot_controls = self._create_plot_controls()
        layout.addWidget(plot_controls)
        
        # 3. Refresh button
        btn_refresh = QPushButton("🔄 Refresh Sessions")
        btn_refresh.clicked.connect(self._load_sessions)
        layout.addWidget(btn_refresh)
        
        # 4. Data series selection (MOVED UP - before Sessions!)
        self.series_group = self._create_series_selector()
        layout.addWidget(self.series_group)
        
        # 5. Session/run tree (MOVED DOWN - after Data Series!)
        tree_group = QGroupBox("Sessions & Runs")
        tree_layout = QVBoxLayout()
        tree_group.setLayout(tree_layout)
        
        self.tree_sessions = QTreeWidget()
        self.tree_sessions.setHeaderLabels(["Session / Run", "Samples", "Duration"])
        self.tree_sessions.itemChanged.connect(self._on_selection_changed)
        tree_layout.addWidget(self.tree_sessions)
        
        # Buttons under tree
        btn_row = QHBoxLayout()
        btn_generate = QPushButton("📊 Generate Plot")
        btn_generate.clicked.connect(self._generate_plot)
        btn_row.addWidget(btn_generate)
        
        btn_export_png = QPushButton("💾 PNG")
        btn_export_png.clicked.connect(self._export_png)
        btn_row.addWidget(btn_export_png)
        
        btn_export_csv = QPushButton("📄 CSV")
        btn_export_csv.clicked.connect(self._export_csv)
        btn_row.addWidget(btn_export_csv)
        
        tree_layout.addLayout(btn_row)
        
        layout.addWidget(tree_group)
        
        layout.addStretch()
        
        return panel
    
    def _create_series_selector(self):
        """Create data series checkboxes."""
        group = QGroupBox("Data Series (Y-Axis)")
        layout = QVBoxLayout()
        group.setLayout(layout)
        
        # Simple vertical list - no grouping
        self.check_t_ns = QCheckBox("t_ns")
        layout.addWidget(self.check_t_ns)
        
        self.check_time_ms = QCheckBox("time_ms")
        layout.addWidget(self.check_time_ms)
        
        self.check_motor_current = QCheckBox("motor_current")
        self.check_motor_current.setChecked(True)  # Default
        layout.addWidget(self.check_motor_current)
        
        self.check_encoder_value = QCheckBox("encoder_value")
        self.check_encoder_value.setChecked(True)  # Default
        layout.addWidget(self.check_encoder_value)
        
        self.check_relative_encoder_value = QCheckBox("relative_encoder_value")
        self.check_relative_encoder_value.setChecked(True)  # Default
        layout.addWidget(self.check_relative_encoder_value)
        
        self.check_v24_pec_diff = QCheckBox("v24_pec_diff")
        layout.addWidget(self.check_v24_pec_diff)
        
        self.check_pwm = QCheckBox("pwm")
        self.check_pwm.setChecked(True)  # Default
        layout.addWidget(self.check_pwm)
        
        self.check_i2c_raw14 = QCheckBox("i2c_raw14")
        layout.addWidget(self.check_i2c_raw14)
        
        self.check_i2c_zero_raw14 = QCheckBox("i2c_zero_raw14")
        layout.addWidget(self.check_i2c_zero_raw14)
        
        self.check_i2c_delta_raw14 = QCheckBox("i2c_delta_raw14")
        layout.addWidget(self.check_i2c_delta_raw14)
        
        self.check_i2c_angle_deg = QCheckBox("i2c_angle_deg")
        self.check_i2c_angle_deg.setChecked(True)  # Default
        layout.addWidget(self.check_i2c_angle_deg)
        
        self.check_i2c_zero_angle_deg = QCheckBox("i2c_zero_angle_deg")
        layout.addWidget(self.check_i2c_zero_angle_deg)
        
        self.check_angular_velocity = QCheckBox("angular_velocity")
        self.check_angular_velocity.setChecked(True)  # Default
        layout.addWidget(self.check_angular_velocity)
        
        return group
    
    def _create_plot_panel(self):
        """Create right panel with plot area - FULL HEIGHT!"""
        panel = QWidget()
        layout = QVBoxLayout()
        panel.setLayout(layout)
        
        # Remove all margins for maximum space
        layout.setContentsMargins(0, 0, 0, 0)
        
        # Store layout reference for canvas recreation
        self.plot_layout = layout
        
        # Matplotlib canvas (FULL SPACE - top to bottom!)
        self.figure = Figure(figsize=(10, 6), dpi=100)
        self.canvas = FigureCanvas(self.figure)
        layout.addWidget(self.canvas)
        
        # Matplotlib toolbar (zoom, pan, save)
        self.toolbar = NavigationToolbar(self.canvas, self)
        layout.addWidget(self.toolbar)
        
        return panel
    
    def _create_plot_controls(self):
        """Create plot control buttons."""
        group = QGroupBox("Plot Configuration")
        layout = QVBoxLayout()
        group.setLayout(layout)
        
        # Plot mode selector
        mode_row = QHBoxLayout()
        mode_row.addWidget(QLabel("Mode:"))
        self.combo_plot_mode = QComboBox()
        self.combo_plot_mode.addItems(["Time Series", "XY Plot"])
        self.combo_plot_mode.currentTextChanged.connect(self._on_plot_mode_changed)
        mode_row.addWidget(self.combo_plot_mode)
        layout.addLayout(mode_row)
        
        # Line type selector
        line_type_row = QHBoxLayout()
        line_type_row.addWidget(QLabel("Line Type:"))
        self.combo_line_type = QComboBox()
        self.combo_line_type.addItems([
            "Line",
            "Line + Markers",
            "Markers Only",
            "Steps"
        ])
        line_type_row.addWidget(self.combo_line_type)
        layout.addLayout(line_type_row)
        
        # Separate window checkbox
        self.check_separate_window = QCheckBox("Open in separate window")
        layout.addWidget(self.check_separate_window)
        
        # X-axis selector (for time series mode)
        self.xaxis_row = QHBoxLayout()
        self.xaxis_label = QLabel("X-Axis:")
        self.xaxis_row.addWidget(self.xaxis_label)
        self.combo_xaxis = QComboBox()
        self.combo_xaxis.addItems(["time_ms", "t_ns"])
        self.xaxis_row.addWidget(self.combo_xaxis)
        layout.addLayout(self.xaxis_row)
        
        # Plot type selector (for time series mode)
        self.type_row = QHBoxLayout()
        self.type_label = QLabel("Type:")
        self.type_row.addWidget(self.type_label)
        self.combo_plot_type = QComboBox()
        self.combo_plot_type.addItems([
            "Overlay",
            "Subplots",
            "Drift Comparison",
            "Multi-Series"
        ])
        self.type_row.addWidget(self.combo_plot_type)
        layout.addLayout(self.type_row)
        
        # XY plot selectors (hidden by default)
        self.xy_row = QHBoxLayout()
        self.xy_x_label = QLabel("X:")
        self.xy_row.addWidget(self.xy_x_label)
        self.combo_xy_x = QComboBox()
        self.combo_xy_x.addItems([
            "time_ms", "t_ns",
            "motor_current", "encoder_value", "relative_encoder_value", "v24_pec_diff", "pwm",
            "i2c_raw14", "i2c_zero_raw14", "i2c_delta_raw14", "i2c_angle_deg", "i2c_zero_angle_deg",
            "angular_velocity"
        ])
        self.combo_xy_x.setCurrentText("motor_current")
        self.xy_row.addWidget(self.combo_xy_x)
        
        self.xy_y_label = QLabel("Y:")
        self.xy_row.addWidget(self.xy_y_label)
        self.combo_xy_y = QComboBox()
        self.combo_xy_y.addItems([
            "time_ms", "t_ns",
            "motor_current", "encoder_value", "relative_encoder_value", "v24_pec_diff", "pwm",
            "i2c_raw14", "i2c_zero_raw14", "i2c_delta_raw14", "i2c_angle_deg", "i2c_zero_angle_deg",
            "angular_velocity"
        ])
        self.combo_xy_y.setCurrentText("pwm")
        self.xy_row.addWidget(self.combo_xy_y)
        layout.addLayout(self.xy_row)
        
        # Hide XY controls initially
        self.xy_x_label.setVisible(False)
        self.xy_y_label.setVisible(False)
        self.combo_xy_x.setVisible(False)
        self.combo_xy_y.setVisible(False)
        
        return group
    
    def _on_plot_mode_changed(self, mode: str):
        """Handle plot mode change."""
        if mode == "XY Plot":
            # Hide time series controls
            self.xaxis_label.setVisible(False)
            self.combo_xaxis.setVisible(False)
            self.type_label.setVisible(False)
            self.combo_plot_type.setVisible(False)
            
            # Show XY controls
            self.xy_x_label.setVisible(True)
            self.xy_y_label.setVisible(True)
            self.combo_xy_x.setVisible(True)
            self.combo_xy_y.setVisible(True)
            
            # Hide data series panel (not needed for XY)
            self.series_group.setVisible(False)
        else:
            # Show time series controls
            self.xaxis_label.setVisible(True)
            self.combo_xaxis.setVisible(True)
            self.type_label.setVisible(True)
            self.combo_plot_type.setVisible(True)
            
            # Hide XY controls
            self.xy_x_label.setVisible(False)
            self.xy_y_label.setVisible(False)
            self.combo_xy_x.setVisible(False)
            self.combo_xy_y.setVisible(False)
            
            # Show data series panel
            self.series_group.setVisible(True)
    
    # =========================================================================
    # Database Connection
    # =========================================================================
    
    def _browse_database(self):
        """Browse for database file."""
        filepath, _ = QFileDialog.getOpenFileName(
            self,
            "Select Database",
            "./database",
            "Database Files (*.db *.sqlite);;All Files (*)"
        )
        
        if filepath:
            self.combo_database.setCurrentText(filepath)
            self._connect_database()
    
    def _connect_database(self):
        """Connect to selected database."""
        db_path = self.combo_database.currentText()
        
        # Close previous connection
        if self.adapter:
            self.adapter.close()
        
        # Determine adapter type based on file extension
        if db_path.endswith('.csv'):
            self.adapter = CSVAdapter(db_path)
            adapter_type = "CSV"
        else:
            self.adapter = SQLiteAdapter(db_path)
            adapter_type = "SQLite"
        
        # Open new connection
        self.db_path = db_path
        
        if self.adapter.connect():
            self._load_sessions()
        else:
            QMessageBox.critical(self, "Connection Error", f"Failed to connect to database:\n{db_path}")
    
    # =========================================================================
    # Session/Run Loading
    # =========================================================================
    
    def _load_sessions(self):
        """Load sessions from database and populate tree."""
        self.sessions = self.adapter.get_sessions()
        self.tree_sessions.clear()
        
        for session in self.sessions:
            # Create session item
            session_item = QTreeWidgetItem(self.tree_sessions)
            session_item.setText(0, session.session_id)
            session_item.setText(1, f"{session.run_count} runs")
            session_item.setText(2, session.created_at)
            session_item.setFlags(session_item.flags() | Qt.ItemFlag.ItemIsUserCheckable)
            session_item.setCheckState(0, Qt.CheckState.Unchecked)
            session_item.setData(0, Qt.ItemDataRole.UserRole, session.session_id)
            
            # Load runs for this session
            runs = self.adapter.get_runs(session.session_id)
            
            for run in runs:
                run_item = QTreeWidgetItem(session_item)
                run_item.setText(0, f"Run {run.run_number} ({run.session_name})")
                run_item.setText(1, f"{run.sample_count} samples")
                run_item.setText(2, f"{run.duration_ms:.1f}ms")
                run_item.setFlags(run_item.flags() | Qt.ItemFlag.ItemIsUserCheckable)
                run_item.setCheckState(0, Qt.CheckState.Unchecked)
                run_item.setData(0, Qt.ItemDataRole.UserRole, (session.session_id, run.run_number))
        
        self.tree_sessions.expandAll()
    
    def _on_selection_changed(self, item: QTreeWidgetItem, column: int):
        """Handle checkbox changes in tree."""
        # Update parent/child checkboxes
        if item.childCount() > 0:  # Session item
            # Update all children
            state = item.checkState(0)
            for i in range(item.childCount()):
                item.child(i).setCheckState(0, state)
        else:  # Run item
            # Check if all siblings checked -> check parent
            parent = item.parent()
            if parent:
                all_checked = all(
                    parent.child(i).checkState(0) == Qt.CheckState.Checked
                    for i in range(parent.childCount())
                )
                parent.setCheckState(0, Qt.CheckState.Checked if all_checked else Qt.CheckState.Unchecked)
    
    def _get_selected_runs(self):
        """Get list of selected (session_id, run_number) tuples."""
        selected = []
        
        # Iterate through tree
        for i in range(self.tree_sessions.topLevelItemCount()):
            session_item = self.tree_sessions.topLevelItem(i)
            
            for j in range(session_item.childCount()):
                run_item = session_item.child(j)
                
                if run_item.checkState(0) == Qt.CheckState.Checked:
                    data = run_item.data(0, Qt.ItemDataRole.UserRole)
                    if data:
                        selected.append(data)
        
        return selected
    
    def _load_selected_data(self):
        """Load telemetry data for selected runs."""
        self.selected_runs = self._get_selected_runs()
        
        if not self.selected_runs:
            QMessageBox.warning(self, "No Selection", "Please select at least one run")
            return
        
        self.loaded_data = []
        
        for session_id, run_no in self.selected_runs:
            data = self.adapter.load_run_data(session_id, run_no)
            if data:
                self.loaded_data.append(data)
        
        if self.loaded_data:
            QMessageBox.information(
                self,
                "Data Loaded",
                f"Successfully loaded {len(self.loaded_data)} run(s)"
            )
        else:
            QMessageBox.warning(self, "Load Error", "Failed to load data from database")
    
    # =========================================================================
    # Plotting
    # =========================================================================
    
    def _get_selected_series(self):
        """Get list of selected data series."""
        series = []
        labels = []
        
        # UART data
        if self.check_motor_current.isChecked():
            series.append('motor_current')
            labels.append('Motor Current')
        
        if self.check_encoder_value.isChecked():
            series.append('encoder_value')
            labels.append('Encoder Value')
        
        if self.check_relative_encoder_value.isChecked():
            series.append('relative_encoder_value')
            labels.append('Relative Encoder')
        
        if self.check_v24_pec_diff.isChecked():
            series.append('v24_pec_diff')
            labels.append('V24 PEC Diff')
        
        if self.check_pwm.isChecked():
            series.append('pwm')
            labels.append('PWM')
        
        # I2C data
        if self.check_i2c_raw14.isChecked():
            series.append('i2c_raw14')
            labels.append('I2C Raw (14-bit)')
        
        if self.check_i2c_zero_raw14.isChecked():
            series.append('i2c_zero_raw14')
            labels.append('I2C Zero Raw')
        
        if self.check_i2c_delta_raw14.isChecked():
            series.append('i2c_delta_raw14')
            labels.append('I2C Delta Raw')
        
        if self.check_i2c_angle_deg.isChecked():
            series.append('i2c_angle_deg')
            labels.append('Angle (degrees)')
        
        if self.check_i2c_zero_angle_deg.isChecked():
            series.append('i2c_zero_angle_deg')
            labels.append('Zero Angle (degrees)')
        
        # Derived data
        if self.check_angular_velocity.isChecked():
            series.append('angular_velocity')
            labels.append('Angular Velocity')
        
        return series, labels
    
    def _generate_plot(self):
        """Generate plot based on selected data and settings."""
        # Auto-load selected runs first
        self.selected_runs = self._get_selected_runs()
        
        if not self.selected_runs:
            QMessageBox.warning(self, "No Selection", "Please select at least one run")
            return
        
        # Load data
        self.loaded_data = []
        
        for session_id, run_no in self.selected_runs:
            data = self.adapter.load_run_data(session_id, run_no)
            if data:
                self.loaded_data.append(data)
        
        if not self.loaded_data:
            QMessageBox.warning(self, "Load Error", "Failed to load data from database")
            return
        
        # Get plot mode
        plot_mode = self.combo_plot_mode.currentText()
        
        # Get line style settings
        line_type = self.combo_line_type.currentText()
        linestyle, marker, markersize = self._get_line_style(line_type)
        
        config = PlotConfig(
            title="Telemetry Data Visualization",
            figsize=(10, 6),
            linestyle=linestyle,
            marker=marker,
            markersize=markersize
        )
        
        try:
            if plot_mode == "XY Plot":
                # XY scatter plot mode
                x_col = self.combo_xy_x.currentText()
                y_col = self.combo_xy_y.currentText()
                
                xlabel = self._get_axis_label(x_col)
                ylabel = self._get_axis_label(y_col)
                
                self.figure = plot_xy_scatter(
                    self.loaded_data,
                    x_col,
                    y_col,
                    xlabel,
                    ylabel,
                    config
                )
            
            else:
                # Time series mode
                series, labels = self._get_selected_series()
                
                if not series:
                    QMessageBox.warning(self, "No Series", "Please select at least one data series")
                    return
                
                # Get X-axis
                x_col = self.combo_xaxis.currentText()
                xlabel = self._get_axis_label(x_col)
                
                plot_type = self.combo_plot_type.currentText()
                
                if plot_type == "Overlay":
                    # FIXED: All series on SAME plot (not stacked)
                    self.figure = self._plot_overlay_all(
                        self.loaded_data,
                        x_col,
                        series,
                        xlabel,
                        labels,
                        config
                    )
                
                elif plot_type == "Subplots":
                    # One subplot per run (first series only)
                    self.figure = plot_subplots(
                        self.loaded_data,
                        x_col,
                        series[0],
                        xlabel,
                        labels[0],
                        config
                    )
                
                elif plot_type == "Drift Comparison":
                    # Drift comparison (first series only)
                    self.figure = plot_comparison(
                        self.loaded_data,
                        x_col,
                        series[0],
                        xlabel,
                        labels[0],
                        reference_index=0,
                        config=config
                    )
                
                elif plot_type == "Multi-Series":
                    # Multi-series stacked subplots
                    self.figure = plot_multi_series(
                        self.loaded_data,
                        x_col,
                        series,
                        xlabel,
                        labels,
                        config
                    )
            
            # Display plot
            if self.check_separate_window.isChecked():
                # Open in separate window
                self._show_plot_window(self.figure)
            else:
                # Embed in main window - recreate canvas to fix matplotlib toolbar
                self._update_canvas(self.figure)
            
            self.plot_updated.emit()
        
        except Exception as e:
            import traceback
            traceback.print_exc()
            QMessageBox.critical(self, "Plot Error", f"Failed to generate plot:\n{e}")
    
    def _plot_overlay_all(self, data_list, x_col, y_columns, xlabel, ylabels, config):
        """
        Plot all series from all runs on SAME axes (true overlay).
        
        Fixes the stacking issue - everything goes on one plot.
        """
        import matplotlib.pyplot as plt
        
        fig, ax = plt.subplots(figsize=config.figsize, dpi=config.dpi)
        
        # Set color cycle for distinguishable colors
        ax.set_prop_cycle(color=plt.cm.tab10.colors)  # 20 colors
        
        # Get line type settings
        line_type = self.combo_line_type.currentText()
        linestyle, marker, markersize = self._get_line_style(line_type)
        
        # Plot each combination of run × series
        for data in data_list:
            x_data = getattr(data, x_col, None)
            
            if x_data is None:
                continue
            
            for y_col, ylabel in zip(y_columns, ylabels):
                y_data = getattr(data, y_col, None)
                
                if y_data is None:
                    continue
                
                # Label format: "session_name run_no (name) - series"
                label = f"{data.session_id} {data.run_no} ({data.session_name}) - {ylabel}"
                ax.plot(x_data, y_data, label=label, alpha=0.8, 
                       linestyle=linestyle, marker=marker, markersize=markersize)
        
        # Formatting
        ax.set_xlabel(xlabel)
        ax.set_ylabel("Value")
        ax.set_title(config.title)
        
        if config.grid:
            ax.grid(True, alpha=0.3)
        
        if config.legend:
            # Legend at bottom, outside plot, 2 columns
            ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.12), ncol=2)
        
        fig.tight_layout(rect=[0, 0.08, 1, 1])  # Leave space for legend
        
        return fig
    
    def _get_line_style(self, line_type):
        """
        Get matplotlib line style parameters.
        
        Returns:
            (linestyle, marker, markersize)
        """
        if line_type == "Line":
            return '-', None, 0
        elif line_type == "Line + Markers":
            return '-', 'o', 3
        elif line_type == "Markers Only":
            return '', 'o', 4
        elif line_type == "Steps":
            return 'steps-post', None, 0
        else:
            return '-', None, 0
    
    def _update_canvas(self, figure):
        """Update canvas with new figure - recreates canvas to fix matplotlib toolbar."""
        from PyQt6.QtWidgets import QApplication
        import matplotlib.pyplot as plt

        # Close old matplotlib figure to free memory
        if hasattr(self, 'figure') and self.figure is not None:
            try:
                plt.close(self.figure)
            except Exception:
                pass

        # Remove old canvas and toolbar safely
        if hasattr(self, 'canvas') and self.canvas is not None:
            try:
                self.plot_layout.removeWidget(self.canvas)
                self.plot_layout.removeWidget(self.toolbar)
                self.canvas.close()
                self.toolbar.close()
                self.canvas.deleteLater()
                self.toolbar.deleteLater()
                self.canvas = None
                self.toolbar = None
                # Process events to ensure cleanup
                QApplication.processEvents()
            except RuntimeError:
                # Widget already deleted
                pass

        # Store figure reference
        self.figure = figure

        # Create fresh canvas and toolbar
        self.canvas = FigureCanvas(figure)
        self.toolbar = NavigationToolbar(self.canvas, self)

        # Add to layout
        self.plot_layout.insertWidget(1, self.canvas)  # After controls
        self.plot_layout.insertWidget(2, self.toolbar)  # After canvas

        # Draw the canvas
        try:
            self.canvas.draw()
        except Exception as e:
            print(f"Warning: Canvas draw failed: {e}")
    
    def _show_plot_window(self, figure):
        """Show plot in separate window."""
        from PyQt6.QtWidgets import QDialog, QVBoxLayout
        
        dialog = QDialog(self)
        dialog.setWindowTitle("Telemetry Plot")
        dialog.resize(1000, 700)
        
        layout = QVBoxLayout()
        dialog.setLayout(layout)
        
        # Canvas and toolbar
        canvas = FigureCanvas(figure)
        toolbar = NavigationToolbar(canvas, dialog)
        
        layout.addWidget(canvas)
        layout.addWidget(toolbar)
        
        dialog.show()
    
    def _get_axis_label(self, column: str) -> str:
        """Get human-readable label for column."""
        labels = {
            't_ns': 'Time (ns)',
            'time_ms': 'Time (ms)',
            'motor_current': 'Motor Current',
            'encoder_value': 'Encoder Value',
            'relative_encoder_value': 'Relative Encoder',
            'v24_pec_diff': 'V24 PEC Diff',
            'pwm': 'PWM',
            'i2c_raw14': 'I2C Raw (14-bit)',
            'i2c_zero_raw14': 'I2C Zero Raw',
            'i2c_delta_raw14': 'I2C Delta Raw',
            'i2c_angle_deg': 'Angle (degrees)',
            'i2c_zero_angle_deg': 'Zero Angle (degrees)',
            'angular_velocity': 'Angular Velocity'
        }
        return labels.get(column, column)
    
    # =========================================================================
    # Export
    # =========================================================================
    
    def _export_png(self):
        """Export current plot to PNG."""
        if not self.loaded_data:
            QMessageBox.warning(self, "No Data", "Please generate a plot first")
            return
        
        filepath, _ = QFileDialog.getSaveFileName(
            self,
            "Export PNG",
            "",
            "PNG Files (*.png)"
        )
        
        if filepath:
            if export_png(self.figure, filepath, dpi=300):
                QMessageBox.information(self, "Success", f"Plot saved to:\n{filepath}")
            else:
                QMessageBox.critical(self, "Error", "Failed to export PNG")
    
    def _export_csv(self):
        """Export loaded data to CSV (only selected columns)."""
        if not self.loaded_data:
            QMessageBox.warning(self, "No Data", "Please load data first")
            return
        
        # Get selected columns based on plot mode
        plot_mode = self.combo_plot_mode.currentText()
        
        if plot_mode == "XY Plot":
            # XY mode: X and Y from dropdowns
            x_col = self.combo_xy_x.currentText()
            y_cols = [self.combo_xy_y.currentText()]
        else:
            # Time series mode: X from dropdown, Y from checkboxes
            x_col = self.combo_xaxis.currentText()
            series, labels = self._get_selected_series()
            
            if not series:
                QMessageBox.warning(self, "No Series", "Please select at least one data series")
                return
            
            y_cols = series
        
        # Get save path
        filepath, _ = QFileDialog.getSaveFileName(
            self,
            "Export CSV",
            "",
            "CSV Files (*.csv)"
        )
        
        if filepath:
            if export_csv(self.loaded_data, filepath, x_col, y_cols):
                QMessageBox.information(self, "Success", f"Data saved to:\n{filepath}")
            else:
                QMessageBox.critical(self, "Error", "Failed to export CSV")
    
    # =========================================================================
    # Cleanup
    # =========================================================================
    
    def closeEvent(self, event):
        """Close database connection on widget close."""
        if self.adapter:
            self.adapter.close()
        event.accept()


# =============================================================================
# Demo
# =============================================================================

if __name__ == "__main__":
    import sys
    from PyQt6.QtWidgets import QApplication
    
    app = QApplication(sys.argv)
    
    # You need to provide a database path
    widget = GraphWidget(db_path="telemetry.db")
    widget.setWindowTitle("Telemetry Graph Visualizer")
    widget.show()
    
    sys.exit(app.exec())
