Pamhyr2/src/View/Results/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
import numpy as np
from matplotlib import collections

from PyQt5.QtCore import (
    QCoreApplication, Qt, QItemSelectionModel,
    QItemSelection, QItemSelectionRange
)
from PyQt5.QtWidgets import QApplication, QTableView

_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
        self._tablemodel = parent._table["profile"]
        self._table = parent.find(QTableView, f"tableView_profile")

    def onpick(self, event):
        if event.mouseevent.inaxes != self.canvas.axes:
            return
        if event.mouseevent.button.value != 1:
            return

        modifiers = QApplication.keyboardModifiers()
        if modifiers not in [Qt.ControlModifier,
                             Qt.NoModifier,
                             Qt.ShiftModifier]:
            return

        ind, point = self._closest_section(event)
        if self._table is None:
            return
        self._select_in_table([ind])
        self._table.blockSignals(False)

        return

    def _closest_section(self, event):
        axes = self.canvas.axes
        mx = event.mouseevent.xdata
        my = event.mouseevent.ydata
        bx, by = axes.get_xlim(), axes.get_ylim()
        ratio = (bx[0] - bx[1]) / (by[0] - by[1])

        segments = event.artist.get_segments()
        ind = event.ind

        points = []
        for i in ind:
            points = points + [[i, j] for j in segments[i]]

        def dist_mouse(point):
            x, y = point[1]
            d2 = (((mx - x) / ratio) ** 2) + ((my - y) ** 2)
            return d2

        closest = min(
            points, key=dist_mouse
        )

        return closest

    def _select_in_table(self, ind):
        if self._table is None:
            return
        self._table.setFocus()
        selection = self._table.selectionModel()
        index = QItemSelection()
        if len(ind) == 0:
            return
        for i in ind:
            index.append(QItemSelectionRange(self._table.model().index(i, 0)))
        selection.select(
            index,
            QItemSelectionModel.Rows |
            QItemSelectionModel.ClearAndSelect |
            QItemSelectionModel.Select
        )

        if len(ind) > 0:
            self._table.scrollTo(self._table.model().index(ind[-1], 0))

    @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 = []
        for xy in zip(reach.geometry.get_x(),
                      reach.geometry.get_y()):
            self.line_xy.append(np.column_stack(xy))

        self._colors, self._style = self.color_hightlight()
        self.line_xy_collection = collections.LineCollection(
            self.line_xy,
            colors=self._colors,
            linestyle=self._style,
            picker=10
        )
        self.canvas.axes.add_collection(self.line_xy_collection)

    def color_hightlight(self):
        reach = self.results.river.reach(self._current_reach_id)
        rows = [self._current_profile_id]
        colors = [self.color_plot
                  for row in range(reach.geometry.number_profiles)]
        style = ["-" for row in range(reach.geometry.number_profiles)]
        if len(rows) > 0:
            for row in rows:
                colors[row] = self.color_plot_current
            if rows[0] > 0:
                colors[rows[0]-1] = self.color_plot_previous
                style[rows[0]-1] = "--"
            if rows[-1] < reach.geometry.number_profiles-1:
                colors[rows[-1]+1] = self.color_plot_next
                style[rows[-1]+1] = "--"
        return colors, style

    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"
            )

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

            if self.is_overflow_point(profile, pt_right):
                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)