plot_window.py

"""PlotWindow is a window in which simulation results are plotted."""

import re
import sys
from typing import Dict, List, Mapping, Optional, Sequence  # , Union

# from numpy import (array, real, imag, real_if_close, absolute, angle)
import numpy as np
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure
from matplotlib.ticker import MaxNLocator
from qtpy.QtCore import Qt
from qtpy.QtWidgets import (  # QFrame,; QScrollArea,; QLineEdit,; QSizePolicy,; QLabel,; QFileDialog,; QShortcut,
    QApplication,
    QCheckBox,
    QHBoxLayout,
    QListWidget,
    QListWidgetItem,
    QPushButton,
    QSizePolicy,
    QVBoxLayout,
    QWidget,
)

from b_asic import Simulation
from b_asic.gui_utils.icons import get_icon
from b_asic.operation import ResultKey
from b_asic.types import Num


class PlotWindow(QWidget):
    """
    Dialog for plotting the result of simulations.

    Parameters
    ----------
    sim_result : dict
        Simulation results of the form obtained from :attr:`~b_asic.simulation.Simulation.results`.
        Alternative, sim_result can be a list, ['name1', sim_result1, 'name2', sim_result2, ...]
    sfg_name : str, optional
        The name of the SFG.
    parent : optional
        The parent window.
    """

    def __init__(
        self,
        sim_result: Mapping[ResultKey, Sequence[Num]],
        figure_name: Optional[str] = None,
    ):
        super().__init__()
        self.setWindowFlags(
            Qt.WindowTitleHint
            | Qt.WindowCloseButtonHint
            | Qt.WindowMinimizeButtonHint
            | Qt.WindowMaximizeButtonHint
        )
        title = (
            f"Simulation results: {figure_name}"
            if figure_name is not None
            else "Simulation results"
        )
        self.setWindowTitle(title)
        self._auto_redraw = False

        ########### Flattening sim_result, if it is a list of results #######
        # take: sim_result (possibly on form ['name1', simres1, 'name2', simres2, ...]
        # generate: sim_result (dict)
        if isinstance(sim_result, Simulation):
            sim_result = sim_result._results
            assert isinstance(sim_result, dict), TypeError(
                "Parsing sim_result as a Simulation, but the _result seems broken."
            )
        elif isinstance(sim_result, list):
            new_result = dict()
            nr = 0  # value number. Used for automatic names.
            name = None
            for element in sim_result:
                if isinstance(element, str):
                    assert not name, Exception(
                        "Parsing sim_result as a list. Did you provide two names after"
                        " each other?"
                    )
                    name = element
                else:
                    if not name:
                        nr = nr + 1
                        name = "(res" + str(nr) + ")"
                    if isinstance(element, dict):
                        res = element
                    elif isinstance(element, Simulation):
                        res = element._results
                        assert isinstance(res, dict), TypeError(
                            f"Parsing sim_result as a list. Result '{name}' is a"
                            " Simulation, and its _result seems broken."
                        )
                    else:
                        raise TypeError(
                            "Parsing sim_result as a list. Supported results are dict"
                            f" and Simulation, but result '{name}' is {type(element)}"
                        )
                    for key, result in res.items():
                        if re.fullmatch("[0-9]+", key):  # it's an output
                            key = "out" + key
                        new_result[name + "." + key] = result
                    name = None
            assert not name, Exception(
                "Parsing sim_result as a list. Did you provide a name as the last item"
                " in the list?"
            )
            sim_result = new_result
        elif not isinstance(sim_result, dict):
            raise TypeError(
                "sim_result must be a dict, Simulation or list."
                f" Found {type(sim_result)}"
            )

        ########### Categorizing/sorting/renaming sim_results: ##############
        #  take: sim_result
        #  generate: key_order, initially_checked
        #  generate: updated_result
        initially_checked = []
        dict_to_sort = {}  # use key=m+1/n, where m:3=input, 4=output, 2=T, 1=others
        updated_result = {}
        n = 0
        for key, result in sim_result.items():
            key2 = key  # in most cases
            if re.fullmatch(r"(.*\.)?in[0-9]+", key):
                m = 4
            elif re.fullmatch(r"[0-9]+", key):  # It's an output as just a number.
                m = 3
                key2 = 'out' + key
            elif re.fullmatch(
                r"(.*\.)?out[0-9]+", key
            ):  # It's an output already formulated.
                m = 3
            elif re.fullmatch(r"(.*\.)?t[0-9]+", key):
                m = 2
            else:
                m = 1

            if all(np.imag(np.real_if_close(result)) == 0):
                n = n + 1
                dict_to_sort[m + 1 / n] = key2
                updated_result[key2] = np.real(result)
                if m == 3:  # output
                    initially_checked.append(key2)
            else:
                # The same again, but split into several lines
                dict_to_sort[m + 1 / (n + 1)] = key2 + '_re'
                dict_to_sort[m + 1 / (n + 2)] = key2 + '_im'
                dict_to_sort[m + 1 / (n + 3)] = key2 + '_mag'
                dict_to_sort[m + 1 / (n + 4)] = key2 + '_ang'
                updated_result[key2 + '_re'] = np.real(result)
                updated_result[key2 + '_im'] = np.imag(result)
                updated_result[key2 + '_mag'] = np.absolute(result)
                updated_result[key2 + '_ang'] = np.angle(result)
                n = n + 4
                if m == 3:  # output
                    initially_checked.append(key2 + '_re')
                    initially_checked.append(key2 + '_im')

        key_order = list(dict(sorted(dict_to_sort.items(), reverse=True)).values())

        # Layout: ############################################
        # | list | icons |
        # | ...  | plot  |
        # | misc |       |

        self._dialog_layout = QHBoxLayout()
        self.setLayout(self._dialog_layout)

        listlayout = QVBoxLayout()
        plotlayout = QVBoxLayout()

        self._dialog_layout.addLayout(listlayout)
        self._dialog_layout.addLayout(plotlayout)

        ########### Plot: ##############
        # Do this before the list layout, as the list layout will re/set visibility
        # Note: The order is of importens. Interesting lines last, to be on top.
        # self.plotcanvas = PlotCanvas(
        #    logger=logger, parent=self, width=5, height=4, dpi=100
        # )
        self._plot_fig = Figure(figsize=(5, 4), layout="compressed")
        self._plot_axes = self._plot_fig.add_subplot(111)
        self._plot_axes.xaxis.set_major_locator(MaxNLocator(integer=True))

        self._lines = {}
        markers = ".ov<^>s+*xd|_"
        fmt = '-'
        ix = 0  # index bytes in markers
        for key in key_order:
            if len(updated_result[key]) <= 100:
                fmt = markers[ix] + '-'
                ix = (ix + 1) % len(markers)
            else:
                fmt = '-'
            line = self._plot_axes.plot(updated_result[key], fmt, label=key)
            self._lines[key] = line[0]

        self._plot_canvas = FigureCanvas(self._plot_fig)
        self._toolbar = NavigationToolbar(self._plot_canvas, self)

        plotlayout.addWidget(self._toolbar)
        plotlayout.addWidget(self._plot_canvas)

        ########### List layout: ##############

        # Add two buttons for selecting all/none:
        hlayout = QHBoxLayout()
        self._button_all = QPushButton(get_icon('all'), "&All")
        self._button_all.clicked.connect(self._button_all_click)
        hlayout.addWidget(self._button_all)
        self._button_none = QPushButton(get_icon('none'), "&None")
        self._button_none.clicked.connect(self._button_none_click)
        hlayout.addWidget(self._button_none)
        listlayout.addLayout(hlayout)

        # Add the entire list
        self._checklist = QListWidget()
        self._checklist.setSizePolicy(QSizePolicy.Preferred, QSizePolicy.Preferred)
        self._checklist.itemChanged.connect(self._item_change)
        listitems = {}
        for key in key_order:
            list_item = QListWidgetItem(key)
            listitems[key] = list_item
            self._checklist.addItem(list_item)
            list_item.setCheckState(
                Qt.CheckState.Unchecked  # CheckState: Qt.CheckState.{Unchecked, PartiallyChecked, Checked}
            )
        for key in initially_checked:
            listitems[key].setCheckState(Qt.CheckState.Checked)
        # self._checklist.setFixedWidth(150)
        listlayout.addWidget(self._checklist)

        # Add additional checkboxes
        self._legend = self._plot_axes.legend()
        self._legend_checkbox = QCheckBox("&Legend")
        self._legend_checkbox.setCheckState(Qt.CheckState.Checked)
        self._legend_checkbox.stateChanged.connect(self._legend_checkbox_change)
        self._legend_checkbox.setIcon(get_icon('legend'))
        listlayout.addWidget(self._legend_checkbox)
        # self.ontop_checkbox = QCheckBox("&On top")
        # self.ontop_checkbox.stateChanged.connect(self._ontop_checkbox_change)
        # self.ontop_checkbox.setCheckState(Qt.CheckState.Unchecked)
        # listlayout.addWidget(self.ontop_checkbox)

        relim_button = QPushButton("&Recompute limits")
        relim_button.clicked.connect(self._relim)
        listlayout.addWidget(relim_button)

        # Add "Close" buttons
        button_close = QPushButton(get_icon('close'), "&Close", self)
        button_close.clicked.connect(self.close)
        listlayout.addWidget(button_close)
        self._relim()

        # Done. Tell the functions below to redraw the canvas when needed.
        # self.plotcanvas.draw()
        self._auto_redraw = True

    def add_result(self, name, result) -> None:
        """
        Add another result to the plot.

        The new signals are added in natural order without sorting them.
        ---
        name: str
          The name of the result
        result: dict
          The result from one simulation.
        """
        markers = ".ov<^>s+*xd|_"
        ix = 0
        self._auto_redraw = False
        for key, res in result.items():
            key2 = name + "." + key
            ischecked = Qt.CheckState.Unchecked
            if re.fullmatch(r"[0-9]+", key):
                key2 = name + '.out' + key
                ischecked = Qt.CheckState.Checked

            if len(res) <= 100:
                fmt = markers[ix] + '-'
                ix = (ix + 1) % len(markers)
            else:
                fmt = '-'

            def addline(key, vector, checked):
                line = self._plot_axes.plot(np.real(res), fmt, label=key)
                self._lines[key] = line[0]
                list_item = QListWidgetItem(key)
                list_item.setCheckState(
                    Qt.CheckState.Unchecked
                )  # will add it if checked
                self._checklist.addItem(list_item)
                self._lines[
                    key
                ].remove()  # remove the line from plot. Keep it in _lines.
                list_item.setCheckState(checked)  # will add it if checked

            if all(np.imag(np.real_if_close(res)) == 0):
                # real: add one line with corresponding checkbox
                addline(key2, np.real(res), ischecked)
            else:
                # complex: add '_re', '_im', '_mag', '_ang'
                addline(key2 + "_re", np.real(res), ischecked)
                addline(key2 + "_im", np.imag(res), ischecked)
                addline(key2 + "_mag", np.absolute(res), Qt.CheckState.Unchecked)
                addline(key2 + "_ang", np.angle(res), Qt.CheckState.Unchecked)
        self._auto_redraw = True
        self._update_legend()
        self._plot_canvas.draw()

    def _legend_checkbox_change(self, check_state):
        self._legend.set(visible=(check_state == Qt.CheckState.Checked))
        self._update_legend()
        self._plot_canvas.draw()

    # def _ontop_checkbox_change(self, checkState):
    # Bugg: It seems the window closes if you change the WindowStaysOnTopHint.
    # (Nothing happens if "changing" from False to False or True to True)
    # self.setWindowFlag(Qt.WindowStaysOnTopHint, on = (checkState == Qt.CheckState.Checked))
    # self.setWindowFlag(Qt.WindowStaysOnTopHint, on = True)
    # print("_ontop_checkbox_change")

    def _button_all_click(self, event):
        self._auto_redraw = False
        for x in range(self._checklist.count()):
            self._checklist.item(x).setCheckState(Qt.CheckState.Checked)
        self._auto_redraw = True
        self._update_legend()
        self._plot_canvas.draw()

    def _update_legend(self):
        # if self._legend_checkbox.checkState == Qt.CheckState.Checked:
        if self._legend_checkbox.isChecked():
            self._legend = self._plot_axes.legend()

    def _button_none_click(self, event):
        self._auto_redraw = False
        for x in range(self._checklist.count()):
            self._checklist.item(x).setCheckState(Qt.CheckState.Unchecked)
        self._auto_redraw = True
        self._update_legend()
        self._plot_canvas.draw()

    def _item_change(self, listitem):
        key = listitem.text()
        if listitem.checkState() == Qt.CheckState.Checked:
            self._plot_axes.add_line(self._lines[key])
        else:
            self._lines[key].remove()
        if self._auto_redraw:
            self._update_legend()
            self._plot_canvas.draw()

    def _relim(self, event=None):
        self._plot_axes.relim(True)
        self._plot_axes.autoscale(True)
        self._plot_axes.autoscale(axis='x', tight=True)
        self._plot_axes.autoscale(axis='y')
        self._plot_canvas.draw()


def start_simulation_dialog(
    sim_results: Dict[str, List[complex]], sfg_name: Optional[str] = None
):
    """Deprecated. Use `show_simulation_result` instead."""
    show_simulation_result(sim_results, sfg_name)


def show_simulation_result(
    sim_results: Dict[str, List[complex]], figure_name: Optional[str] = None
):
    """
    Display the simulation results window.

    Parameters
    ----------
    sim_results : dict
        Simulation results of the form obtained from :attr:`~b_asic.simulation.Simulation.results`.
        Alternative, sim_result can be a list, ['name1', sim_result1, 'name2', sim_result2, ...]
    sfg_name : str, optional
        The name of the SFG.
    """
    if not QApplication.instance():
        app = QApplication(sys.argv)
    else:
        app = QApplication.instance()
    win = PlotWindow(sim_result=sim_results, figure_name=figure_name)
    win.show()
    app.exec_()


# Simple test of the dialog
if __name__ == "__main__":
    res1 = {
        '0': [1.5, 1.6, 1.5, 1],
        '1': [1.0, 2.0, 1.5, 1.1],
        'add1': [1.5, 1.5, 1, 1],
        'cmul1': [1, 1.5, 1, 1],
        'cmul2': [1.5, 1, 1, 1],
        'in1': [2, 1, 1, 1],
        'in2': [1.1, 3, 1, 1],
        't1': [1, 2, 1, 1],
        't2': [1, 1, 2, 1],
        't3': [1, 1, 1, 2],
    }
    res2 = {
        '0': [0.5, 0.6, 0.5, 0],
        '1': [0.0, 1.0 + 0.3j, 0.5, 0.1j],
        'add1': [0.5, 0.5, 0, 0],
        'cmul1': [0, 0.5, 0, 0],
        'cmul2': [0.5, 0, 0, 0],
        'in1': [1, 0, 0, 0],
        'in2': [0.1, 2, 0, 0],
        't1': [0, 1, 0, 0],
        't2': [0, 0, 1, 0],
        't3': [0, 0, 0, 1],
    }
    res3 = {
        #'0': np.random.rand(50).tolist(),
        '0': np.random.rand(50),
        '1': np.random.rand(200).tolist(),
    }
    res4 = {
        '0': np.random.rand(60).tolist(),
        '1': np.random.rand(220).tolist(),
        't4': np.random.rand(50).tolist(),
    }
    # start_simulation_dialog(res3)
    app = QApplication(sys.argv)
    win2 = PlotWindow(['Real', res1, 'Cpx', res2, res3], "Test data")
    win2.add_result('res4', res4)
    win2.show()
    app.exec_()