from __future__ import annotations

from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg, NavigationToolbar2QT
from matplotlib.figure import Figure
from PySide6.QtCore import Signal
from PySide6.QtWidgets import QVBoxLayout, QWidget

from bedit_gui.simulation_models import SimulationPlotSettings
from bedit_simulation import SimulationResult


class SimulationPlotWidget(QWidget):
    settings_requested = Signal()

    def __init__(self, parent: QWidget | None = None) -> None:
        super().__init__(parent)
        self.figure = Figure(layout="constrained")
        self.canvas = FigureCanvasQTAgg(self.figure)
        self.toolbar = NavigationToolbar2QT(self.canvas, self)
        for action in self.toolbar.actions():
            if action.text() in ("Customize", "Subplots"):
                self.toolbar.removeAction(action)
        self.toolbar.addSeparator()
        self.settings_action = self.toolbar.addAction("Plot settings…")
        self.settings_action.setToolTip("Edit plot title, axes, traces, grid, and legend")
        self.settings_action.triggered.connect(self.settings_requested)
        self.axes = self.figure.add_subplot()
        layout = QVBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.addWidget(self.toolbar)
        layout.addWidget(self.canvas)

    def set_plot(self, results: list[SimulationResult], signals: list[str], x_axis: str | None = None, settings: SimulationPlotSettings | None = None) -> None:
        settings = settings or SimulationPlotSettings()
        self.axes.clear()
        for signal in signals:
            trace = settings.traces.get(signal)
            if trace is not None and not trace.visible:
                continue
            x_values: list[float] = []
            values: list[float] = []
            sample_offset = 0
            for result in results:
                result_values = result.data.get(signal)
                if result_values is None:
                    continue
                result_x = result.data.get(x_axis or "time")
                if result_x is not None and len(result_x) == len(result_values):
                    x_values.extend(result_x)
                elif x_axis is not None:
                    continue
                else:
                    x_values.extend(range(sample_offset, sample_offset + len(result_values)))
                values.extend(result_values)
                sample_offset += len(result_values)
            if values:
                options = {"label": trace.label or signal, "linestyle": trace.line_style, "linewidth": trace.line_width, "marker": trace.marker or None, "markersize": trace.marker_size} if trace is not None else {"label": signal}
                if trace is not None and trace.color:
                    options["color"] = trace.color
                self.axes.plot(x_values, values, **options)
        self.axes.set_title(settings.title)
        self.axes.set_xlabel(settings.x_label or x_axis or "Time")
        self.axes.set_ylabel(settings.y_label)
        self.axes.set_xscale(settings.x_scale)
        self.axes.set_yscale(settings.y_scale)
        if not settings.x_auto:
            self.axes.set_xlim(settings.x_min, settings.x_max)
        if not settings.y_auto:
            self.axes.set_ylim(settings.y_min, settings.y_max)
        self.axes.grid(settings.grid_visible, axis=settings.grid_axis, linestyle=settings.grid_style, alpha=settings.grid_alpha)
        if self.axes.lines and settings.legend_visible:
            self.axes.legend(loc=settings.legend_location)
        self.canvas.draw_idle()