icbktransitionuim.py

import datetime
import warnings
import pyqtgraph as pg
import scipy.optimize.optimize
from PyQt5.QtWidgets import *
from PyQt5.QtGui import QColor
from PyQt5.QtCore import *
import pandas as pd
import pytz
from pandas.api.types import is_numeric_dtype
from pyqtgraph.GraphicsScene.mouseEvents import MouseClickEvent

import utils
from config import Config
from dataprovider.icbktransitionprovider import IcbktransitionProvider
from uim.explouim import ExploUim
from gui.ui_mainwindow import Ui_MainWindow
from gui.icbktransition_addon_dialog import IcbktransitionAddonDialog
from cfatools.logreader.dataset import DatasetReader


class IcbktransitionUim:

    # "Variables" tableWidget columns identifiers.
    ID_COL = 0
    NAME_COL = 1
    START_COL = 2
    END_COL = 3
    HEIGHT_COL = 4
    ACTIONS_COL = 5

    PROBES_COLORS = {"C1": QColor(228, 26, 28),  # Red
                     "C2": QColor(255, 255, 51),  # Yellow
                     "C3": QColor(77, 175, 74),  # Green
                     "C4": QColor(152, 78, 163),  # Violet
                     "C5": QColor(247, 129, 191)}  # Pink
                      # QColor(166, 86, 40),  # Brown
                      # }

    PROBES = ["C1", "C2", "C3", "C4", "C5"]

    PEN_TRANSITION_BASE = pg.mkPen(style=Qt.SolidLine, width=1, color=QColor(55, 126, 184))  # Blue
    PEN_TRANSITION_HOVERED = pg.mkPen(style=Qt.DotLine, width=3, color=QColor(55, 126, 184))  # Blue
    PEN_TRANSITION_SELECTED = pg.mkPen(style=Qt.DotLine, width=3, color=QColor(255, 127, 0))  # Orange

    def __init__(self, icbktransition_prvd: IcbktransitionProvider, explo_uim: ExploUim,
                 main_ui: Ui_MainWindow, config: Config, icbktransition_dialog: IcbktransitionAddonDialog):
        self.main_ui = main_ui
        self.icbktransition_prvd = icbktransition_prvd
        self.explo_uim = explo_uim
        self.config = config
        self.icbktransition_dialog = icbktransition_dialog
        self.icbktransition_ui = self.icbktransition_dialog.ui

        self.dataset_readers = dict()
        self.current_dataset = None

        self.__init_plot__()
        self.main_ui.explo_pushbutton_icbktransition.clicked.connect(self.__show_icbktransition_dialog__)
        self.icbktransition_ui.timeedit_reftime.dateTimeChanged.connect(self.__update_conduct_curves__)
        self.icbktransition_ui.checkbox_reftime_disable.clicked.connect(self.__update_conduct_curves__)

        self.__initialize_transition_table__()
        self.icbktransition_ui.toolbutton_reload_icbkctrl.clicked.connect(self.__reload_from_icbkctrl__)
        self.icbktransition_ui.toolbutton_reload_iceblocks.clicked.connect(self.__reload_from_iceblocks__)
        self.icbktransition_ui.toolbutton_save_iceblocks.clicked.connect(self.icbktransition_prvd.save_iceblock_df_as_csv)
        self.icbktransition_dialog.sig_escape_pressed.connect(self.__unselect_transition__)

        self.hovered_transition = None
        self.selected_transition = None

    def __show_icbktransition_dialog__(self):
        self.current_dataset = self.explo_uim.current_dataset
        self.icbktransition_prvd.load_dataset(self.current_dataset)
        self.conduct_timeseries = self.current_dataset.get_timeseries("CONDUCTI-R_periodic").copy()
        self.picarro_df = self.current_dataset.get_timeseries("PICARRO_periodic").copy()
        self.__init_reftime__()
        self.__update_conduct_curves__()
        self.icbktransition_dialog.show()
        self.__load_iceblocks_in_transition_table__()
        self.__reload_transitions_in_plot__()

    ####################################################################################################################
    # "Transition" table

    def __initialize_transition_table__(self):
        # Set column widths
        self.icbktransition_ui.tablewidget_icbk_list.setColumnWidth(self.ID_COL, 30)
        self.icbktransition_ui.tablewidget_icbk_list.setColumnWidth(self.NAME_COL, 130)
        self.icbktransition_ui.tablewidget_icbk_list.setColumnWidth(self.START_COL, 90)
        self.icbktransition_ui.tablewidget_icbk_list.setColumnWidth(self.END_COL, 90)
        self.icbktransition_ui.tablewidget_icbk_list.setColumnWidth(self.HEIGHT_COL, 100)
        self.icbktransition_ui.tablewidget_icbk_list.setColumnWidth(self.ACTIONS_COL, 120)

    def __reload_from_icbkctrl__(self):
        iceblock_df = self.icbktransition_prvd.build_iceblock_df(retrieve_iceblocks_from_processed_csv=False)
        self.__load_iceblocks_in_transition_table__(iceblock_df)
        self.__reload_transitions_in_plot__()

    def __reload_from_iceblocks__(self):
        iceblock_df = self.icbktransition_prvd.build_iceblock_df(retrieve_iceblocks_from_processed_csv=True)
        self.__load_iceblocks_in_transition_table__(iceblock_df)
        self.__reload_transitions_in_plot__()

    ####################################################################
    # Transitions

    def __select_transition__(self):
        if self.hovered_transition is None:
            return
        self.selected_transition = self.hovered_transition
        self.__paint_transitions_in_table__()
        self.__paint_transitions_in_plot__()

    def __unselect_transition__(self):
        self.selected_transition = None
        self.__paint_transitions_in_table__()
        self.__paint_transitions_in_plot__()

    def __move_transition__(self, timestamp_sec: int):
        self.icbktransition_prvd.move_transition(self.selected_transition, timestamp_sec)

        self.__load_iceblocks_in_transition_table__(self.icbktransition_prvd.iceblock_df)

        self.originaltransition_vlines[self.selected_transition].setPos(timestamp_sec)
        self.selected_transition = None
        self.__paint_transitions_in_plot__()

    def __delete_transition__(self, transition_id: int):
        self.icbktransition_prvd.delete_transition(transition_id)
        self.__load_iceblocks_in_transition_table__(self.icbktransition_prvd.iceblock_df)
        self.__reload_transitions_in_plot__()

    def __mark_hovered_transition__(self, timestamp_sec: int):
        """Determine which transition is hovered on the plot (if any) and repaint plot and table"""
        df = self.icbktransition_prvd.transitions_df.copy()
        df["interval"] = abs(df["timestamp"] - timestamp_sec)
        df = df[df["interval"] < (5 * 60)]

        if len(df.index) == 0:
            self.hovered_transition = None
        else:
            self.hovered_transition = df.loc[[df["interval"].idxmin()]].index[0]

        self.__paint_transitions_in_plot__()
        self.__paint_transitions_in_table__()

    ####################################################################################################################
    # Table modifications

    def __load_iceblocks_in_transition_table__(self, iceblock_df: pd.DataFrame = None):
        table = self.icbktransition_ui.tablewidget_icbk_list

        if iceblock_df is None:
            iceblock_df = self.icbktransition_prvd.build_iceblock_df()

        # Clear table
        table.setRowCount(0)

        # Load iceblock_df in the table
        for row_id, iceblock in iceblock_df.iterrows():
            table.insertRow(row_id)

            # ID
            id_item = QTableWidgetItem()
            id_item.setFlags(Qt.ItemIsEnabled)  # Read only
            id_item.setText(str(iceblock["id"]))
            table.setItem(row_id, self.ID_COL, id_item)

            # NAME
            name_item = QTableWidgetItem()
            name_item.setText(str(iceblock["name"]))
            table.setItem(row_id, self.NAME_COL, name_item)

            # START
            start_item = QTimeEdit()
            start_item.setDateTime(iceblock["datetime_start"])
            start_item.setButtonSymbols(QAbstractSpinBox.NoButtons)
            start_item.setReadOnly(True)
            table.setCellWidget(row_id, self.START_COL, start_item)

            # END
            end_item = QTimeEdit()
            end_item.setDateTime(iceblock["datetime_end"])
            end_item.setButtonSymbols(QAbstractSpinBox.NoButtons)
            end_item.setReadOnly(True)
            table.setCellWidget(row_id, self.END_COL, end_item)

            # HEIGHT
            height_item = QTableWidgetItem()
            height_item.setFlags(Qt.ItemIsEnabled)  # Read only
            height_item.setText(str(iceblock["initial_height"]))
            table.setItem(row_id, self.HEIGHT_COL, height_item)

            # ACTIONS
            actions_item = QPushButton('Delete')
            table.setCellWidget(row_id, self.ACTIONS_COL, actions_item)
            table.cellWidget(row_id, self.ACTIONS_COL).clicked.connect(
                lambda checked, row_id=row_id: self.__delete_transition__(transition_id=row_id))

    def __paint_transitions_in_table__(self):
        # First, reset the display
        table = self.icbktransition_ui.tablewidget_icbk_list
        for row_id in range(table.rowCount()):
            table.cellWidget(row_id, self.START_COL).setStyleSheet("color: 'black';")
            table.cellWidget(row_id, self.END_COL).setStyleSheet("color: 'black';")

        if self.hovered_transition is not None:
            self.__paint_single_transition_in_table__(self.hovered_transition, "blue")
        if self.selected_transition is not None:
            self.__paint_single_transition_in_table__(self.selected_transition, "orange")

    def __paint_single_transition_in_table__(self, transition_id: int, color: str):
        table = self.icbktransition_ui.tablewidget_icbk_list
        if transition_id < table.rowCount():
            table.cellWidget(transition_id, self.START_COL).setStyleSheet("color: '"+color+"';")
            table.scrollToItem(table.item(transition_id, self.ID_COL))
        if transition_id > 0:
            table.cellWidget(transition_id - 1, self.END_COL).setStyleSheet("color: '"+color+"';")

    ####################################################################################################################
    # Plot modifications

    def __init_plot__(self):
        self.plot_item = pg.PlotItem(axisItems={'bottom': utils.TimeAxisItem(orientation='bottom')})
        self.icbktransition_dialog.ui.graphicsview.setCentralItem(self.plot_item)
        self.plot_item.sigRangeChanged.connect(self.__autorescale_picarro_curve__)

        # Conduct curves
        self.conduct_curves = dict()
        for probe in self.PROBES:
            self.conduct_curves[probe] = pg.PlotCurveItem()
            self.plot_item.addItem(self.conduct_curves[probe])

        # Picarro curve
        self.picarro_curve = pg.PlotCurveItem()
        self.plot_item.addItem(self.picarro_curve)

        # Vertical lines for the original iceblocks transition
        self.originaltransition_vlines = []
        self.originaltransition_texts = []

        # Vertical lines for the original iceblocks transition
        self.newtransition_vlines = []
        self.newtransition_texts = []

        # Vertical line following the cursor
        self.cursor_vline = pg.InfiniteLine(angle=90, movable=False, pen=pg.mkPen(style=Qt.DotLine))
        self.plot_item.addItem(self.cursor_vline, ignoreBounds=True)
        self.conduct_curves["C1"].scene().sigMouseMoved.connect(lambda event: self.__mouse_moved__(event))
        self.conduct_curves["C1"].scene().sigMouseClicked.connect(lambda event: self.__mouse_clicked__(event))

    def __autorescale_picarro_curve__(self, viewbox: pg.ViewBox, range: list):
        range_x = range[0]
        range_y = range[1]

        rescaled_picarro_df = self.picarro_df.copy()
        x_min = pd.Timestamp(range_x[0]*1000000000, tz="UTC")
        x_max = pd.Timestamp(range_x[1]*1000000000, tz="UTC")
        top = range_y[1]
        bottom = range_y[0]
        # top = self.conduct_timeseries[x_min:x_max].max().max()
        # bottom = self.conduct_timeseries[x_min:x_max].min().min()
        print(str(top) + " - " + str(bottom))
        picarro_visible_df = self.picarro_df[x_min:x_max]
        picarro_min = picarro_visible_df["Delta_D_H"].min()
        picarro_max = picarro_visible_df["Delta_D_H"].max()

        rescaled_picarro_df["rescaled_Delta_D_H"] = (top - bottom) * (rescaled_picarro_df["Delta_D_H"] - picarro_min) / (picarro_max-picarro_min) + bottom

        x_values = utils.pd_time_to_epoch_ms(rescaled_picarro_df.index)
        x_values.append(x_values[-1] + 1)
        y_values = list(rescaled_picarro_df["rescaled_Delta_D_H"])
        self.plot_item.sigRangeChanged.disconnect(self.__autorescale_picarro_curve__)
        self.picarro_curve.setData(x=x_values,
                                   y=y_values,
                                   pen=QColor(166, 86, 40),  # Brown
                                   stepMode=True)
        self.plot_item.sigRangeChanged.connect(self.__autorescale_picarro_curve__)

    def __mouse_moved__(self, pos: QPointF) -> None:
        """Function triggered when the user's mouse cursor hovers over the plot. Display a vertical line where the
        cursor is.

        Parameters
        ----------
        pos: PyQt5.QtCore.QPointF
        """
        mouse_point = self.conduct_curves["C1"].getViewBox().mapSceneToView(pos)

        self.plot_item.sigRangeChanged.disconnect(self.__autorescale_picarro_curve__)
        self.cursor_vline.setPos(mouse_point.x())
        self.plot_item.sigRangeChanged.connect(self.__autorescale_picarro_curve__)
        self.__update_timeedit_cursor__(mouse_point.x())

        if self.selected_transition is None:
            self.__mark_hovered_transition__(mouse_point.x())

    def __mouse_clicked__(self, event: MouseClickEvent) -> None:
        """Function triggered when the user clicks on the plot. Display a vertical line under the mouse click and call
        function ``__update_xy_original__``

        Parameters
        ----------
        event: pyqtgraph.MouseClickEvent
        """
        pos = event.scenePos()
        mouse_point = self.conduct_curves["C1"].getViewBox().mapSceneToView(pos)
        if self.selected_transition is None:
            self.__select_transition__()
        else:
            self.__move_transition__(mouse_point.x())

    def __reload_transitions_in_plot__(self):
        """Delete any existing transition vline and re-create then, based on Provider's transitions_df."""
        for transition_vline in self.originaltransition_vlines:
            self.plot_item.removeItem(transition_vline)
        self.originaltransition_vlines = []

        transitions_df = self.icbktransition_prvd.get_transitions()
        transitions_df["timestamp"] = utils.pd_time_to_epoch_ms(transitions_df["datetime"])
        self.plot_item.sigRangeChanged.disconnect(self.__autorescale_picarro_curve__)
        for index, transition in transitions_df.iterrows():
            transition_line = pg.InfiniteLine(angle=90, movable=False, pen=self.PEN_TRANSITION_BASE)
            transition_line.setPos(transition["timestamp"])
            self.originaltransition_vlines.append(transition_line)
            self.plot_item.addItem(transition_line, ignoreBounds=True)
        self.plot_item.sigRangeChanged.connect(self.__autorescale_picarro_curve__)

    def __paint_transitions_in_plot__(self):
        # First, reset the display
        for transition_vline in self.originaltransition_vlines:
            transition_vline.setPen(self.PEN_TRANSITION_BASE)

        if self.hovered_transition is not None:
            self.originaltransition_vlines[self.hovered_transition].setPen(self.PEN_TRANSITION_HOVERED)

        if self.selected_transition is not None:
            self.originaltransition_vlines[self.selected_transition].setPen(self.PEN_TRANSITION_SELECTED)

    def __update_conduct_curves__(self):
        self.plot_item.sigRangeChanged.disconnect(self.__autorescale_picarro_curve__)
        # Set data to step curve
        x_values = utils.pd_time_to_epoch_ms(self.conduct_timeseries.index)
        x_values.append(x_values[-1] + 1)
        for probe in self.PROBES:
            if self.icbktransition_ui.checkbox_reftime_disable.isChecked():
                y_values = list(self.conduct_timeseries[probe])
            else:
                y_values = list(self.__shift_to_reftime__(self.conduct_timeseries[probe]))
            self.conduct_curves[probe].setData(x=x_values,
                                               y=y_values,
                                               pen=self.PROBES_COLORS[probe],
                                               stepMode=True)
        self.plot_item.sigRangeChanged.connect(self.__autorescale_picarro_curve__)

    ####################################################################################################################
    # Reftime: use a common reference point for all conducti probes

    def __init_reftime__(self):
        min_reftime = min(self.conduct_timeseries.index)
        self.icbktransition_ui.timeedit_reftime.blockSignals(True)
        self.icbktransition_ui.timeedit_reftime.setMinimumDateTime(min_reftime)
        self.icbktransition_ui.timeedit_reftime.setMaximumDateTime(max(self.conduct_timeseries.index))
        self.icbktransition_ui.timeedit_reftime.setDateTime(min_reftime + datetime.timedelta(minutes=20))
        self.icbktransition_ui.timeedit_reftime.blockSignals(False)

    def __shift_to_reftime__(self, series: pd.Series) -> pd.Series:
        timezone = pytz.timezone("UTC")
        reftime = self.icbktransition_ui.timeedit_reftime.dateTime().toPyDateTime()
        reftime = timezone.localize(reftime)
        value_at_reftime = series.iloc[series.index.get_loc(reftime,
                                                            method='nearest')]
        series = series - value_at_reftime
        return series

    ####################################################################################################################
    # Tiemedit cursor

    def __update_timeedit_cursor__(self, timestamp_sec: int):
        pos_x = QDateTime()
        pos_x.setMSecsSinceEpoch(timestamp_sec * 1000)
        pos_x = pos_x.toUTC()
        self.icbktransition_ui.timeedit_cursor.setDateTime(pos_x)