FigTimeDischarge.py

"""
QRame
Copyright (C) 2023  INRAE

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published
by the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

import numpy as np
from PyQt5 import QtCore
from datetime import datetime
import matplotlib.dates as mdates

class FigTimeDischarge(object):
    """Class to plot chronological discharge graph.

        Attributes
        ----------
        canvas: MplCanvas
            Object of MplCanvas a FigureCanvas
        fig: Object
            Figure object of the canvas
        units: dict
            Dictionary of units from units_conversion
        _translate: QCoreApplication.translate object
            Save words which need to be translated
        hover_connection: int
            Index to data cursor connection
        annot: Annotation
            Annotation object for data cursor
        """

    def __init__(self, canvas, units):
        """Initialize object using the specified canvas.

        Parameters
        ----------
        canvas: MplCanvas
            Object of MplCanvas
        units: dict
            Dictionary of units from units_conversion
        """

        # Initialize attributes
        self.canvas = canvas
        self.fig = canvas.fig
        self.units = units
        self.hover_connection = None
        self.annot = None
        self._translate = QtCore.QCoreApplication.translate

    def create(self, mean_selected_meas, selected_meas=None):
        """Create the axes and lines for the figure.

        Parameters
        ----------
        mean_selected_meas: pandas DataFrame
            Measurement results dataframe
        selected_meas: str
            Name of the selected measurement
        """

        # Clear the plot
        self.fig.clear()

        # Configure axis
        self.fig.ax = self.fig.add_subplot(1, 1, 1)

        # Set margins and padding for figure
        # self.fig.subplots_adjust(left=0.03, bottom=0.25, right=0.99, top=0.99, wspace=0.1, hspace=0)

        df = mean_selected_meas
        df['color'] = 'k'
        # Change color for selected
        if selected_meas and selected_meas in df.meas_name.values:
            df.loc[df['meas_name'] == selected_meas, 'color'] = 'darkorange'

        # Plot horizontal line if there is less than 3 days
        if df[['start_time', 'end_time']].max().max() - df[['start_time', 'end_time']].min().min() < 1210000:
            is_time = ~np.logical_or(np.isnan(df.start_time), np.isnan(df.end_time))
            df_time = df.loc[is_time, ['meas_name', 'tr_q_total', 'start_time', 'end_time', 'color']]

            start_datetime = [datetime.utcfromtimestamp(i) for i in df_time.start_time]
            end_datetime = [datetime.utcfromtimestamp(i) for i in df_time.end_time]
            self.fig.ax.hlines(df_time.tr_q_total * self.units['Q'], start_datetime, end_datetime,
                               color=df_time.color, linewidth=2, zorder=0)

            # Highlight selected measurement in orange
            if selected_meas:
                df_selected = df.loc[df['meas_name'] == selected_meas]
                start_selected = [datetime.utcfromtimestamp(i) for i in df_selected.start_time if not np.isnan(i)]
                end_selected = [datetime.utcfromtimestamp(i) for i in df_selected.end_time if not np.isnan(i)]
                if len(start_selected) > 0 and len(start_selected) == len(end_selected):
                    self.fig.ax.hlines(df_selected.tr_q_total * self.units['Q'], start_selected, end_selected,
                                       color=df_selected.color, linewidth=3, zorder=1)

            # Plot data with only one date (start or end time) as scatter
            df_missing_time = df.loc[~is_time, ['meas_name', 'tr_q_total', 'start_time', 'end_time', 'color']]
            any_time = np.logical_or(~np.isnan(df_missing_time.start_time), ~np.isnan(df_missing_time.end_time))
            df_scatter = df_missing_time[any_time]
            df_scatter['time'] = df_scatter['start_time'].fillna(0) + df_scatter['end_time'].fillna(0)
            scatter_datetime = [datetime.utcfromtimestamp(i) for i in df_scatter.time]
            self.fig.ax.scatter(scatter_datetime, df_scatter['tr_q_total'] * self.units['Q'],
                                color=df_scatter.color, zorder=0)

            self.fig.ax.xaxis.set_major_formatter(mdates.DateFormatter('%d %b %H:%M'))
        else:
            df = df.groupby('meas_name').agg({'start_time': 'min', 'end_time': 'max', 'meas_mean_q': 'mean',
                                              'color': 'first'})

            mid_time = df[['start_time', 'end_time']].mean(axis=1)
            is_time = ~np.isnan(mid_time)
            mid_datetime = [datetime.utcfromtimestamp(i) for i in mid_time[is_time]]
            self.fig.ax.scatter(mid_datetime, df.loc[is_time, 'meas_mean_q'] * self.units['Q'],
                                color=df.loc[is_time, 'color'], zorder=0)
            self.fig.ax.xaxis.set_major_formatter(mdates.DateFormatter('%d %b %y'))

        self.fig.ax.tick_params(which="major", axis="x", pad=14, size=2)
        self.fig.autofmt_xdate()

        self.fig.ax.set_xlabel(self._translate('Main', 'Time'))
        self.fig.ax.set_ylabel(self._translate('Main', 'Discharge') + ' (' + self.units['label_Q'] + ')')
        self.fig.ax.xaxis.label.set_fontsize(12)
        self.fig.ax.yaxis.label.set_fontsize(12)
        self.fig.ax.tick_params(axis='both', direction='in', bottom=True, top=True, left=True, right=True)
        self.fig.ax.grid()


    def hover(self, event):
        """Determines if the user has selected a location with temperature data and makes
        annotation visible and calls method to update the text of the annotation. If the
        location is not valid the existing annotation is hidden.

        Parameters
        ----------
        event: MouseEvent
            Triggered when mouse button is pressed.
        """

        # Set annotation to visible
        vis = self.annot.get_visible()

        # Determine if mouse location references a data point in the plot and update the annotation.
        if event.inaxes == self.fig.ax and event.button != 3:
            cont = False
            ind = None
            plotted_line = None

            # Find the transect(line) that contains the mouse click
            for plotted_line in self.fig.ax.lines:
                cont, ind = plotted_line.contains(event)
                if cont:
                    break
            if cont:
                self.update_annot(ind, plotted_line)
                self.annot.set_visible(True)
                self.canvas.draw_idle()
            else:
                # If the cursor location is not associated with the plotted data hide the annotation.
                if vis:
                    self.annot.set_visible(False)
                    self.canvas.draw_idle()

    def update_annot(self, ind, plt_ref):
        """Updates the location and text and makes visible the previously initialized and hidden annotation.

        Parameters
        ----------
        ind: dict
            Contains data selected.
        plt_ref: Line2D
            Reference containing plotted data
        vector_ref: Quiver
            Refernece containing plotted data
        ref_label: str
            Label used to ID data type in annotation
        """

        pos = plt_ref._xy[ind["ind"][0]]

        # Shift annotation box left or right depending on which half of the axis the pos x is located and the
        # direction of x increasing.
        if plt_ref.axes.viewLim.intervalx[0] < plt_ref.axes.viewLim.intervalx[1]:
            if pos[0] < (plt_ref.axes.viewLim.intervalx[0] + plt_ref.axes.viewLim.intervalx[1]) / 2:
                self.annot._x = -20
            else:
                self.annot._x = -80
        else:
            if pos[0] < (plt_ref.axes.viewLim.intervalx[0] + plt_ref.axes.viewLim.intervalx[1]) / 2:
                self.annot._x = -80
            else:
                self.annot._x = -20

        # Shift annotation box up or down depending on which half of the axis the pos y is located and the
        # direction of y increasing.
        if plt_ref.axes.viewLim.intervaly[0] < plt_ref.axes.viewLim.intervaly[1]:
            if pos[1] > (plt_ref.axes.viewLim.intervaly[0] + plt_ref.axes.viewLim.intervaly[1]) / 2:
                self.annot._y = -40
            else:
                self.annot._y = 20
        else:
            if pos[1] > (plt_ref.axes.viewLim.intervaly[0] + plt_ref.axes.viewLim.intervaly[1]) / 2:
                self.annot._y = 20
            else:
                self.annot._y = -40
        self.annot.xy = pos

        text = 'x: {:.2f}, y: {:.2f}'.format(pos[0], pos[1])
        self.annot.set_text(text)

    def set_hover_connection(self, setting):
        """Turns the connection to the mouse event on or off.

        Parameters
        ----------
        setting: bool
            Boolean to specify whether the connection for the mouse event is active or not.
        """

        if setting and self.hover_connection is None:
            self.hover_connection = self.canvas.mpl_connect('button_press_event', self.hover)
        elif not setting:
            self.canvas.mpl_disconnect(self.hover_connection)
            self.hover_connection = None
            self.annot.set_visible(False)
            self.canvas.draw_idle()