Window.py

# Window.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 os
import logging

from tools import trace, timer, logger_exception

from View.Tools.PamhyrWindow import PamhyrWindow

from PyQt5 import QtWidgets
from PyQt5.QtGui import (
    QKeySequence,
)

from PyQt5.QtCore import (
    Qt, QVariant, QAbstractTableModel,
    QCoreApplication, QModelIndex, pyqtSlot,
    QRect, QSettings, QItemSelectionModel,
)

from PyQt5.QtWidgets import (
    QDialogButtonBox, QPushButton, QLineEdit,
    QFileDialog, QTableView, QAbstractItemView,
    QUndoStack, QShortcut, QAction, QItemDelegate,
    QComboBox, QVBoxLayout, QHeaderView, QTabWidget,
)

from Modules import Modules

from View.InitialConditions.UndoCommand import (
    SetCommand, AddCommand, DelCommand,
    SortCommand, MoveCommand, InsertCommand,
    DuplicateCommand,
)

from View.InitialConditions.Table import (
    InitialConditionTableModel, ComboBoxDelegate,
)

from View.Tools.Plot.PamhyrCanvas import MplCanvas
from View.Tools.Plot.PamhyrToolbar import PamhyrPlotToolbar

from View.InitialConditions.PlotDRK import PlotDRK
from View.InitialConditions.PlotDischarge import PlotDischarge
from View.InitialConditions.translate import ICTranslate
from View.InitialConditions.DialogDepth import DepthDialog
from View.InitialConditions.DialogHeight import HeightDialog
from View.InitialConditions.DialogDischarge import DischargeDialog
from View.Results.ReadingResultsDialog import ReadingResultsDialog

from Solver.Mage import Mage8

_translate = QCoreApplication.translate

logger = logging.getLogger()


class InitialConditionsWindow(PamhyrWindow):
    _pamhyr_ui = "InitialConditions"
    _pamhyr_name = "Initial condition"

    def __init__(self, study=None, config=None, reach=None, parent=None):
        trad = ICTranslate()

        if reach is not None:
            self._reach = reach
        else:
            self._reach = study.river.current_reach()

        name = (
            trad[self._pamhyr_name] +
            " - " + study.name +
            " - " + self._reach.name
        )

        super(InitialConditionsWindow, self).__init__(
            title=name,
            study=study,
            config=config,
            trad=trad,
            parent=parent
        )

        # Add reach to hash computation data
        self._hash_data.append(self._reach)

        self._ics = study.river.initial_conditions.get(self._reach)

        self.setup_table()
        self.setup_plot()
        self.setup_connections()
        self.setub_dialogs()

        self.ui.setWindowTitle(self._title)

    def setup_table(self):
        table = self.find(QTableView, f"tableView")
        self._delegate_rk = ComboBoxDelegate(
            reach=self._reach,
            parent=self
        )

        self._table = InitialConditionTableModel(
            reach=self._reach,
            table_view=table,
            table_headers=self._trad.get_dict("table_headers"),
            editable_headers=["rk", "discharge", "elevation", "height"],
            delegates={"rk": self._delegate_rk},
            data=self._study,
            undo=self._undo_stack,
            trad=self._trad
        )

        table.setModel(self._table)
        table.setSelectionBehavior(QAbstractItemView.SelectRows)
        table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        table.setAlternatingRowColors(True)

    def setup_plot(self):
        self.canvas_1 = MplCanvas(width=5, height=4, dpi=100)
        self.canvas_1.setObjectName("canvas_1")
        self.toolbar_1 = PamhyrPlotToolbar(
            self.canvas_1, self
        )
        self.plot_layout_1 = self.find(QVBoxLayout, "verticalLayout_1")
        self.plot_layout_1.addWidget(self.toolbar_1)
        self.plot_layout_1.addWidget(self.canvas_1)

        self.plot_1 = PlotDRK(
            canvas=self.canvas_1,
            data=self._ics,
            trad=self._trad,
            toolbar=self.toolbar_1,
            parent=self
        )
        self.plot_1.draw()

        self.canvas_2 = MplCanvas(width=5, height=4, dpi=100)
        self.canvas_2.setObjectName("canvas_2")
        self.toolbar_2 = PamhyrPlotToolbar(
            self.canvas_2, self
        )
        self.plot_layout_2 = self.find(QVBoxLayout, "verticalLayout_2")
        self.plot_layout_2.addWidget(self.toolbar_2)
        self.plot_layout_2.addWidget(self.canvas_2)

        self.plot_2 = PlotDischarge(
            canvas=self.canvas_2,
            data=self._ics,
            trad=self._trad,
            toolbar=self.toolbar_2,
            parent=self
        )
        self.plot_2.draw()

    def setup_connections(self):
        self.find(QAction, "action_add").triggered.connect(self.add)
        self.find(QAction, "action_del").triggered.connect(self.delete)
        self.find(QAction, "action_sort").triggered.connect(self.sort)
        self.find(QAction, "action_import").triggered\
                                           .connect(self.import_from_file)

        self.find(QPushButton, "pushButton_generate_1").clicked.connect(
            self.generate_growing_constant_depth
        )

        self.find(QPushButton, "pushButton_generate_2").clicked.connect(
            self.generate_discharge
        )

        self.find(QPushButton, "pushButton_generate_3").clicked.connect(
            self.generate_height
        )

        self._table.dataChanged.connect(self._update_plot)

    def setub_dialogs(self):
        self.height_values = [0.0, 0.0, 0.0]
        self.height_option = True
        self.discharge_value = 0.0
        self.discharge_option = True
        self.depth_value = 0.0
        self.depth_option = True

    def index_selected_row(self):
        table = self.find(QTableView, f"tableView")
        rows = table.selectionModel()\
                    .selectedRows()

        if len(rows) == 0:
            return 0

        return rows[0].row()

    def update(self):
        self._update(propagate=False)

    def _update(self, propagate=True):
        self._update_plot()

        if propagate:
            self._propagate_update(key=Modules.INITIAL_CONDITION)

    def _update_plot(self):
        self.plot_1.draw()
        self.plot_2.draw()

    def _propagated_update(self, key=Modules(0)):
        if Modules.GEOMETRY not in key:
            return

        self.update()

    def index_selected_rows(self):
        table = self.find(QTableView, f"tableView")
        return list(
            # Delete duplicate
            set(
                map(
                    lambda i: i.row(),
                    table.selectedIndexes()
                )
            )
        )

    def add(self):
        rows = self.index_selected_rows()
        if len(self._ics) == 0 or len(rows) == 0:
            self._table.add(0)
        else:
            self._table.add(rows[0])

        self._update()

    def delete(self):
        rows = self.index_selected_rows()
        if len(rows) == 0:
            return

        self._table.delete(rows)
        self._update()

    def sort(self):
        self._table.sort(False)
        self._update()

    def import_from_file(self):
        options = QFileDialog.Options()
        settings = QSettings(QSettings.IniFormat,
                             QSettings.UserScope, 'MyOrg', )
        options |= QFileDialog.DontUseNativeDialog
        filename, _ = QtWidgets.QFileDialog.getOpenFileName(
            self,
            self._trad["open_file"],
            "",
            ";; ".join(["Mage (*.BIN)"]),
            options=options
        )

        if filename != "":
            size = os.stat(filename).st_size
            #self._table.import_geometry(0, filename)
            print(f"filename: {filename}")
            self._import_from_file(filename)

    def _import_from_file(self, file_name):
        solver = Mage8("dummy")
        name = os.path.basename(file_name)\
                      .replace(".BIN", "")

        def reading():
            try:
                self._tmp_results = solver.results(
                    self._study,
                    os.path.dirname(file_name),
                    name=name
                )
            except Exception as e:
                logger.error(f"Failed to open results")
                logger_exception(e)

        dlg = ReadingResultsDialog(
            reading_fn=reading,
            parent=self
        )
        dlg.exec_()
        results = self._tmp_results
        if "timestamps" in results._meta_data:
            self._import_from_results(results)

    def _import_from_results(self, results):
        logger.debug(f"import from results: {results}")
        row = self.index_selected_row()

        self._table.import_from_results(row, results)

    def move_up(self):
        row = self.index_selected_row()
        self._table.move_up(row)
        self._update()

    def move_down(self):
        row = self.index_selected_row()
        self._table.move_down(row)
        self._update()

    def _copy(self):
        rows = list(
            map(
                lambda row: row.row(),
                self.tableView.selectionModel().selectedRows()
            )
        )

        table = list(
            map(
                lambda eic: list(
                    map(
                        lambda k: eic[1][k],
                        ["rk", "discharge", "elevation"]
                    )
                ),
                filter(
                    lambda eic: eic[0] in rows,
                    enumerate(self._ics.lst())
                )
            )
        )

        self.copyTableIntoClipboard(table)

    def _paste(self):
        header, data = self.parseClipboardTable()

        if len(data) + len(header) == 0:
            return

        logger.debug(
            "IC: Paste: " +
            f"header = {header}, " +
            f"data = {data}"
        )

        try:
            row = self.index_selected_row()
            # self._table.paste(row, header, data)
            self._table.paste(row, [], data)
        except Exception as e:
            logger_exception(e)

        self._update()

    def _undo(self):
        self._table.undo()
        self._update()

    def _redo(self):
        self._table.redo()
        self._update()

    def generate_growing_constant_depth(self):
        if self._reach.reach.number_profiles > 0:
            dlg = DepthDialog(self.depth_value,
                              self.depth_option,
                              trad=self._trad,
                              parent=self)
            if dlg.exec():
                self.depth_value = dlg.value
                self.depth_option = dlg.option
                self._table.generate("growing",
                                     self.depth_value,
                                     self.depth_option)
                self._update()

    def generate_discharge(self):
        if self._reach.reach.number_profiles > 1:
            dlg = DischargeDialog(self.discharge_value,
                                  self.discharge_option,
                                  trad=self._trad,
                                  parent=self)
            if dlg.exec():
                self.discharge_value = dlg.value
                self.discharge_option = dlg.option
                self._table.generate("discharge",
                                     self.discharge_value,
                                     self.discharge_option)
                self._update()

    def generate_height(self):
        if self._reach.reach.number_profiles > 0:
            dlg = HeightDialog(self.height_values,
                               self.height_option,
                               trad=self._trad,
                               parent=self)
            if dlg.exec():
                self.height_values = dlg.values
                self.height_option = dlg.option
                self._table.generate("height",
                                     self.height_values,
                                     self.height_option)
                self._update()