Pamhyr2/src/View/Results/PlotH.py

# PlotH.py -- Pamhyr
# Copyright (C) 2023-2024  INRAE
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

# -*- coding: utf-8 -*-

import logging

from functools import reduce
from datetime import datetime

from tools import timer, trace
from View.Tools.PamhyrPlot import PamhyrPlot

from PyQt5.QtCore import (
    QCoreApplication
)

_translate = QCoreApplication.translate

logger = logging.getLogger()


class PlotH(PamhyrPlot):
    def __init__(self, canvas=None, trad=None, toolbar=None,
                 results=None, reach_id=0, profile_id=[0], res_id=[0],
                 parent=None):
        super(PlotH, self).__init__(
            canvas=canvas,
            trad=trad,
            data=results,
            toolbar=toolbar,
            parent=parent
        )

        self._mode = "time"

        self._current_reach_id = reach_id
        self._current_profile_id = profile_id
        self._current_res_id = res_id
        self._timestamps = parent._timestamps
        self._current_timestamp = max(self._timestamps)

        self.label_x = self._trad["unit_time_s"]
        self.label_y = self._trad["unit_discharge"]

        self.label_discharge = _translate("Results", "Cross-section discharge")
        self.label_discharge_max = _translate("Results", "Max discharge")
        self.label_timestamp = _translate("Results", "Current timestamp")

        self._isometric_axis = False

        # self._auto_relim_update = False
        # self._autoscale_update = False

        self._line_max = []
        self._line = []

    @property
    def results(self):
        return self.data

    @results.setter
    def results(self, results):
        self.data = results
        self._current_timestamp = max(self._timestamps)

    @timer
    def draw(self, highlight=None):
        self.canvas.axes.cla()
        self.canvas.axes.grid(color='grey', linestyle='--', linewidth=0.5)

        if self.results is None:
            return

        if len(self._current_res_id) < 1:
            return

        if len(self._current_res_id) == 1:
            self.draw_max(self._current_res_id[0])

        for res_id in self._current_res_id:
            if self.results[res_id] is None:
                continue

            results = self.results[res_id]
            reach = results.river.reach(self._current_reach_id)

            if reach.geometry.number_profiles == 0:
                self._init = False
                return

            self.draw_data(res_id)

        self.canvas.axes.set_xlabel(
            self._trad["unit_time_s"],
            color='black', fontsize=10
        )

        self.canvas.axes.set_ylabel(
            self._trad["unit_discharge"],
            color='black', fontsize=10
        )

        self.set_ticks_time_formater()

        self.enable_legend()

        self.canvas.draw_idle()
        self.update_idle
        self.draw_current()
        self._init = True

    def draw_data(self, res_id):

        results = self.results[res_id]
        reach = results.river.reach(self._current_reach_id)
        for i, p in enumerate(self._current_profile_id):
            profile = reach.profile(p)

            x = self._timestamps
            y = profile.get_key("Q")

            if res_id == 2:
                label = "Δ " + self.label_discharge + f" {profile.name}"
            else:
                label = self.label_discharge + f" {profile.name}"

            if len(self._current_res_id) > 1:
                if res_id != 2:
                    label += f" ({results._solver.name})"

            line, = self.canvas.axes.plot(
                x, y,
                label=label,
                color=self.colors[res_id + 1],
                linestyle=self.linestyle[i % len(self.linestyle)],
                **self.plot_default_kargs
            )
            self._line.append(line)

    def draw_current(self):
        self._current, = self.canvas.axes.plot(
            [self._current_timestamp, self._current_timestamp],
            self.canvas.axes.get_ylim(),
            # label=self.label_timestamp,
            color="grey",
            linestyle="dashed",
            lw=1.,
        )

    def draw_max(self, res_id):
        results = self.results[res_id]
        reach = results.river.reach(self._current_reach_id)

        if reach.geometry.number_profiles == 0:
            self._init = False
            return

        x = self._timestamps
        y = []

        if res_id == 2:
            label = "Δ " + self.label_discharge_max
        else:
            label = self.label_discharge_max

        for ts in x:
            ts_y = -9999
            for profile in reach.profiles:
                q = profile.get_ts_key(ts, "Q")
                ts_y = max(ts_y, q)
            y.append(ts_y)

        m, = self.canvas.axes.plot(
            x, y,
            label=label,
            color=self.color_plot_highlight,
            linestyle='dotted',
            **self.plot_default_kargs
        )
        self._line_max.append(m)

    def set_reach(self, reach_id):
        self._current_reach_id = reach_id
        self._current_profile_id = [0]
        self.draw()

    def set_profile(self, profile_id):
        self._current_profile_id = profile_id
        # self.update()
        self.draw()

    def set_result(self, res_id):
        self._current_res_id = res_id
        self.draw()

    def set_timestamp(self, timestamp):
        self._current_timestamp = timestamp
        self.update_current()
        self.update_idle()

    def update(self):
        if not self._init:
            self.draw()

        self.update_data()
        self.update_current()
        self.update_idle()

    def update_data(self):
        for res, res_id in enumerate(self._current_res_id):
            results = self.results[res_id]
            reach = results.river.reach(self._current_reach_id)
            profile = reach.profile(self._current_profile_id)

            x = self._timestamps
            y = profile.get_key("Q")

            self._line[res].set_data(x, y)

    def update_current(self):
        y = self._current.get_ydata()
        self._current.set_data(
            [self._current_timestamp, self._current_timestamp],
            y
        )