###############################################################################
# Example of a Python program for plotting graphs using matplotlib in
# combination with Qt6 (PySide6).  Run the program from the command line:
# 
#   python3 qt-matplotlib.py
#
# The easiest way to install Matplotlib and Qt6 (Pyside6) is with pip:
#
#   pip3 install matplotlib
#   pip3 install pyside6
#
# (c) Copyright 2018-2021 Anders Andersen, UiT The Arctic University of Norway
###############################################################################


# Standard modules
import sys
from numpy import *

# Qt and matplotlib integration
import matplotlib
matplotlib.use("QtAgg")
from matplotlib.backends.backend_qtagg import FigureCanvas

# Matplotlib
from matplotlib.figure import Figure

# Use Qt with the PySide6 mapping (Qt's official Python mapping)
from PySide6.QtWidgets import QApplication, QWidget, QHBoxLayout, QVBoxLayout
from PySide6.QtWidgets import QPushButton, QLabel, QLineEdit, QCheckBox

# Enhanced mode?
try:
    enhanced_mode = (sys.argv[1] == "-x")
except:
    enhanced_mode = False
    

class MyPlotWidget(QWidget):
    """A widget that plots graphs

    An example Qt widget that plots graphs with matplotlib.

    """

    # The explanation text (when no errors are detected)
    labeltext = "Write an expression to plot (Python syntax, " + \
        "including numpy mathematical functions with one variable): "

    def __init__(self, app):
        """Build widget"""

        # Initialize widet
        QWidget.__init__(self)

        # Create plots (figsize in inches, width and height)
        self.fig = Figure(figsize=(10,8), dpi=100,
                          facecolor=(0.8,0.8,0.9), edgecolor=(0,0,0))
        self.ax = self.fig.add_subplot(1, 1, 1)

        # Add plot to Qt GUI usig a canvas:
        # A canvas must be manually attached to the figure (pyplot would
        # automatically do it). This is done by instantiating the canvas
        # with the figure as argument.
        self.canvas = FigureCanvas(self.fig)

        # Title and text
        self.setWindowTitle("Plot a graph")
        self.plotlabel = QLabel(self.labeltext)

        # Edit fields
        self.exprlabel = QLabel("Expression (with variable x):")
        self.expr = QLineEdit("sin(1/x)")
        self.minlabel = QLabel("Min x:")
        self.min = QLineEdit("-pi/4")
        self.min.setFixedWidth(80)
        self.maxlabel = QLabel("Max x:")
        self.max = QLineEdit("pi/4")
        self.max.setFixedWidth(80)

        # Create buttons
        self.button1 = QPushButton("Plot it")
        self.button2 = QPushButton("Exit")

        # Define layout
        self.minmaxline = QHBoxLayout()
        self.layout = QVBoxLayout()
        self.layout.addWidget(self.canvas)
        self.layout.addWidget(self.plotlabel)
        self.minmaxline.addWidget(self.exprlabel)
        self.minmaxline.addWidget(self.expr)
        self.minmaxline.addSpacing(20)
        self.minmaxline.addWidget(self.minlabel)
        self.minmaxline.addWidget(self.min)
        self.minmaxline.addSpacing(10)
        self.minmaxline.addWidget(self.maxlabel)
        self.minmaxline.addWidget(self.max)
        self.minmaxline.addSpacing(30)
        self.minmaxline.addWidget(self.button1)
        self.minmaxline.addWidget(self.button2)
        self.layout.addLayout(self.minmaxline)

        # Add extra line in enhanced mode (&-shortcut does not work on macos)
        self.emmulti =  QCheckBox("Plot &multiple graphs on the same canvas")
        if enhanced_mode:
            self.emline = QHBoxLayout()
            self.emline.addWidget(self.emmulti)
            self.layout.addLayout(self.emline)

        # Set layout of widget
        self.setLayout(self.layout)

        # Connect button clicks to actions
        self.button1.clicked.connect(self.plotit)
        self.button2.clicked.connect(app.exit)

        # Connect return to actions
        self.expr.returnPressed.connect(self.plotit)
        self.min.returnPressed.connect(self.plotit)
        self.max.returnPressed.connect(self.plotit)

    def plotit(self):
        """Update widget on user input"""
        
        try:
            # Evaluate min, max and expression (unsafe code!)
            x = linspace(eval(self.min.text()), eval(self.max.text()), 100001)
            y = eval(self.expr.text())
            
        except:
            # Not valid expression, min, and/or max
            self.ax.clear()
            self.ax.figure.canvas.draw()
            self.ax.figure.canvas.blit(self.ax.figure.bbox)
            self.plotlabel.setText(
                "Non valid expression or min and max value for x: " + \
                "Use only 'x' in expression and valid numbers for max and min:")
            
        else:
            # Update plot
            self.plotlabel.setText(self.labeltext)
            if not self.emmulti.isChecked():
                self.ax.clear()
            self.ax.plot(x, y)

        # Redraw canvas (blit: only redraw the graph part)
        self.ax.figure.canvas.draw()
        self.ax.figure.canvas.blit(self.ax.figure.bbox)


# Create Qt app and widget
app = QApplication(sys.argv)
widget = MyPlotWidget(app)
widget.show()
sys.exit(app.exec())