import agxSDK
import os

from agxPythonModules.utils.environment import is_running_tests

pyqtgraph_enabled = True
try:
    os.environ["QT_LOGGING_RULES"] = "qt.qpa.*=false;qt.core.qobject.connect=false"  # Silence annoying warnings

    from pyqtgraph.Qt import QtWidgets, QtCore  # PySide6/PyQt5 via pyqtgraph>=0.13
    import pyqtgraph as pg

    if is_running_tests():
        pyqtgraph_enabled = False
        print("Running tests. Plotting disabled")
    else:
        print("Plotting enabled")
except Exception as e:
    print("Plotting disabled: {}".format(e))
    pyqtgraph_enabled = False


class PlotGraph:
    def __init__(self, plotItem):
        self.plotItem = plotItem
        self.plots = []
        self.plotItem.addLegend(offset=(0, 0))
        self.timeStamps = []

    def add_plot(self, name, f=lambda t: 0.0, pen=(255, 255, 255)):
        self.plots.append([self.plotItem.plot([0], [0], name=name, pen=pen), f, []])

    def update(self, t):
        self.timeStamps.append(t)
        for [plotData, f, fData] in self.plots:
            fData.append(f(t))
            plotData.setData(self.timeStamps, fData)


class PlotWindow(agxSDK.GuiEventListener):
    def __init__(self, agx_app):
        super().__init__(agxSDK.GuiEventListener.UPDATE)
        self.agx_app = agx_app

        # Reuse existing QApplication if present; otherwise create one
        self.app = QtWidgets.QApplication.instance() or QtWidgets.QApplication([])

        self.win = pg.GraphicsLayoutWidget()
        self.win.setWindowTitle(self.agx_app.getArguments().getArgumentName(1))
        self.win.resize(1024, 768)
        self.win.show()

        self.graphs = []

    def update(self, x, y):
        for graphs in self.graphs:
            graphs.update(self.getSimulation().getTimeStamp())
        QtWidgets.QApplication.processEvents()

    def add_graph(self, title, new_row=False, units=""):
        if new_row:
            self.win.nextRow()
        p = self.win.addPlot(title=title)
        p.showGrid(x=True, y=True)
        p.setLabel("left", "Value", units=units)
        p.setLabel("bottom", "Time", units="s")
        plotGraph = PlotGraph(p)
        self.graphs.append(plotGraph)
        return plotGraph

