PlotXY.py

# PlotXY.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 tools import timer, trace
from View.Tools.PamhyrPlot import PamhyrPlot

from PyQt5.QtCore import (
    QCoreApplication
)

_translate = QCoreApplication.translate

logger = logging.getLogger()


class PlotXY(PamhyrPlot):
    def __init__(self, canvas=None, trad=None, toolbar=None,
                 results=None, reach_id=0, profile_id=0,
                 display_current=True, parent=None):
        super(PlotXY, self).__init__(
            canvas=canvas,
            trad=trad,
            data=results,
            toolbar=toolbar,
            parent=parent
        )

        self.line_xy = []
        self.line_gl = []
        self.overflow = []

        self._timestamps = results.get("timestamps")
        self._current_timestamp = max(self._timestamps)
        self._current_reach_id = reach_id
        self._current_profile_id = profile_id

        self.label_x = _translate("Results", "X (m)")
        self.label_y = _translate("Results", "Y (m)")

        self._isometric_axis = True

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

    @results.setter
    def results(self, results):
        self.data = results
        self._current_timestamp = max(results.get("timestamps"))

    @timer
    def draw(self, highlight=None):
        self.init_axes()

        if self.results is None:
            return

        reach = self.results.river.reach(self._current_reach_id)

        self.draw_profiles(reach)
        self.draw_water_elevation(reach)
        self.draw_water_elevation_max(reach)
        self.draw_guide_lines(reach)
        self.draw_current(reach)

        self.idle()
        self._init = True

    def draw_profiles(self, reach):
        if reach.geometry.number_profiles == 0:
            self._init = False
            return

        self.line_xy = [
            self.canvas.axes.plot(
                x, y,
                color=self.color_plot_river_bottom,
                **self.plot_default_kargs
            )
            for x, y, kp in zip(
                reach.geometry.get_x(),
                reach.geometry.get_y(),
                reach.geometry.get_kp()
            )
        ]

    def draw_guide_lines(self, reach):
        x_complete = reach.geometry.get_guidelines_x()
        y_complete = reach.geometry.get_guidelines_y()

        self.line_gl = [
            self.canvas.axes.plot(
                x, y,
            )
            for x, y in zip(x_complete, y_complete)
        ]

    def draw_current(self, reach):
        profile = reach.profile(self._current_profile_id)

        self.plot_selected, = self.canvas.axes.plot(
            profile.geometry.x(),
            profile.geometry.y(),
            color=self.color_plot,
            **self.plot_default_kargs
        )


    def draw_water_elevation_max(self, reach):
        l_x, l_y, r_x, r_y = [], [], [], []
        overflow = []

        for profile in reach.profiles:
            z_max = max(profile.get_key("Z"))
            z_max_ts = 0
            for ts in self._timestamps:
                z = profile.get_ts_key(ts, "Z")
                if z == z_max:
                    z_max_ts = ts
                    break

            pt_left, pt_right = profile.get_ts_key(z_max_ts, "water_limits")

            l_x.append(pt_left.x)
            l_y.append(pt_left.y)
            r_x.append(pt_right.x)
            r_y.append(pt_right.y)

            if self.is_overflow_point(profile, pt_left):
                overflow.append(pt_left)

            if self.is_overflow_point(profile, pt_right):
                overflow.append(pt_right)

        self.water_max_left = self.canvas.axes.plot(
            l_x, l_y,
            color=self.color_plot_river_water,
            linestyle='dotted',
            lw=1.,
        )

        self.water_max_right = self.canvas.axes.plot(
            r_x, r_y,
            color=self.color_plot_river_water,
            linestyle='dotted',
            lw=1.,
        )

        for p in overflow:
            self.canvas.axes.plot(
                p.x, p.y,
                lw=1.,
                color=self.color_plot,
                markersize=3,
                marker='x'
            )


    def is_overflow_point(self, profile, point):
            left_limit = profile.geometry.point(0)
            right_limit = profile.geometry.point(
                profile.geometry.number_points - 1
            )

            return (
                point == left_limit
                or point == right_limit
            )


    def draw_water_elevation(self, reach):
        reach = self.results.river.reach(self._current_reach_id)
        poly_l_x, poly_l_y, poly_r_x, poly_r_y = [], [], [], []

        for profile in reach.profiles:
            water_z = profile.get_ts_key(
                self._current_timestamp, "Z"
            )
            pt_left, pt_right = profile.get_ts_key(
                self._current_timestamp,
                "water_limits"
            )

            poly_l_x.append(pt_left.x)
            poly_l_y.append(pt_left.y)
            poly_r_x.append(pt_right.x)
            poly_r_y.append(pt_right.y)

        poly_x = poly_l_x + list(reversed(poly_r_x))
        poly_y = poly_l_y + list(reversed(poly_r_y))

        self.water_fill = self.canvas.axes.fill(
            poly_x, poly_y,
            color=self.color_plot_river_water_zone,
            alpha=0.7
        )

    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_profile()
        self.update_idle()

    def set_timestamp(self, timestamp):
        self._current_timestamp = timestamp
        self.update()

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

        self.update_water_elevation()
        self.update_water_elevation_overflow()

        self.update_idle()

    def update_profile(self):
        reach = self.results.river.reach(self._current_reach_id)
        profile = reach.profile(self._current_profile_id)

        self.plot_selected.set_data(
            profile.geometry.x(),
            profile.geometry.y()
        )

    def update_water_elevation(self):
        reach = self.results.river.reach(self._current_reach_id)
        poly_l_x, poly_l_y, poly_r_x, poly_r_y = [], [], [], []

        for profile in reach.profiles:
            water_z = profile.get_ts_key(
                self._current_timestamp, "Z"
            )
            pt_left, pt_right = profile.get_ts_key(
                self._current_timestamp,
                "water_limits"
            )

            poly_l_x.append(pt_left.x)
            poly_l_y.append(pt_left.y)
            poly_r_x.append(pt_right.x)
            poly_r_y.append(pt_right.y)

        poly_x = poly_l_x + list(reversed(poly_r_x))
        poly_y = poly_l_y + list(reversed(poly_r_y))

        poly = []
        for i in range(len(poly_x)):
            poly.append([poly_x[i], poly_y[i]])

        self.water_fill[0].set_xy(poly)

    def update_water_elevation_overflow(self):
        reach = self.results.river.reach(self._current_reach_id)
        profile = reach.profile(self._current_profile_id)

        overflow = []

        for profile in reach.profiles:
            pt_left, pt_right = profile.get_ts_key(
                self._current_timestamp,
                "water_limits"
            )

            left_limit = profile.geometry.point(0)
            right_limit = profile.geometry.point(
                profile.geometry.number_points - 1
            )

            if pt_left == left_limit:
                overflow.append(pt_left)

            if pt_right == right_limit:
                overflow.append(pt_right)

        for plot in self.overflow:
            plot[0].remove()
            del plot[0]
        self.overflow = []

        for p in overflow:
            plot = self.canvas.axes.plot(
                p.x, p.y,
                lw=1.,
                color=self.color_plot,
                markersize=3,
                marker='o'
            )
            self.overflow.append(plot)