# PlotXY.py -- Pamhyr
# Copyright (C) 2023-2026  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

import numpy as np

from tools import timer, trace, logger_exception
from View.Tools.PamhyrPlot import PamhyrPlot
from matplotlib import collections

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

try:
    import rasterio
    import rasterio.control
    import rasterio.crs
    import rasterio.sample
    import rasterio.vrt
    import rasterio._features

    from rasterio.io import MemoryFile
    _rasterio_loaded = True
except Exception as e:
    print(f"Module 'rasterio' is not available: {e}")
    _rasterio_loaded = False

_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, res_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._plot_img = {}

        self._geotiff_img = []
        self._geotiff_bounds = []

        self._parent = parent
        self._timestamps = parent._timestamps
        self._current_timestamp = max(self._timestamps)
        self._current_reach_id = reach_id
        self._current_profile_id = profile_id
        self._current_res_id = res_id[0]

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

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

        if self.results is None:
            return

        if self.results[self._current_res_id] is None:
            return

        results = self.results[self._current_res_id]
        reach = results.river.reach(self._current_reach_id)
        reaches = results.river.reachs

        self.draw_geotiff()
        self.draw_profiles(reach, reaches)
        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, reaches):
        if reach.geometry.number_profiles == 0:
            self._init = False
            return

        self.line_xy = []
        # TODO uncomment to draw all the reaches
        # self.draw_other_profiles(reaches)
        for xy in zip(reach.geometry.get_x(),
                      reach.geometry.get_y()):
            self.line_xy.append(np.column_stack(xy))

        self._colors = self.color_plot_river_bottom
        self._style = "-"
        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 draw_other_profiles(self, reaches):
        for reach in reaches:
            for xy in zip(reach.geometry.get_x(),
                          reach.geometry.get_y()):
                self.line_xy.append(np.column_stack(xy))

            self.line_xy_collection = collections.LineCollection(
                self.line_xy,
                colors=self.color_plot_river_bottom,
            )
            self.canvas.axes.add_collection(self.line_xy_collection)

    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):
        result = self.results[self._current_res_id]

        l_x, l_y, r_x, r_y = [], [], [], []
        overflow = []

        for profile in reach.profiles:
            zs = result.get("table")["Z"][
                :, profile.global_index
            ]
            z_max = np.max(zs)
            z_max_ts = 0

            for i, ts in enumerate(self._timestamps):
                z = zs[i]
                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):
        results = self.results[self._current_res_id]
        reach = results.river.reach(self._current_reach_id)
        poly_l_x, poly_l_y, poly_r_x, poly_r_y = [], [], [], []

        for profile in reach.profiles:
            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 add_geotiff(self, img, bounds):
        self._geotiff_img.append(img)
        self._geotiff_bounds.append(bounds)

    def draw_add_geotiff(self):
        xlim = self.canvas.axes.get_xlim()
        ylim = self.canvas.axes.get_ylim()

        for i, img in enumerate(self._geotiff_img):
            bounds = self._geotiff_bounds[i]

            self.canvas.axes.imshow(img.transpose((1, 2, 0)),
                                    extent=bounds)

        self.idle()
        self.canvas.axes.set_xlim(xlim)
        self.canvas.axes.set_ylim(ylim)

    def draw_geotiff(self):
        if not _rasterio_loaded:
            return

        lst = self._data[0]._study.river._geotiff.lst

        for img in self._plot_img:
            try:
                self._plot_img[img].remove()
            except Exception as e:
                logger_exception(e)

        self._plot_img = {}

        for geotiff in lst:
            memfile = geotiff.memfile
            if memfile is None:
                return

            with memfile.open() as gt:
                img = gt.read()
                coords = geotiff.coordinates

                left = coords["left"]
                right = coords["right"]
                bottom = coords["bottom"]
                top = coords["top"]

                self._plot_img[geotiff] = self.canvas.axes.imshow(
                    img.transpose((1, 2, 0)),
                    extent=(left, right, bottom, top)
                )

            if not geotiff.is_enabled():
                self._plot_img[geotiff].set(alpha=0.5)

        self.draw_add_geotiff()
        self.idle()

    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_result(self, res_id):
        self._current_res_id = res_id[0]
        self.update()

    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_all(self):
        self._current_reach_id = self._parent._get_current_reach()
        self._current_profile_id = self._parent._get_current_profile()
        self._current_res_id = self._parent._get_current_results()[0]
        self._current_timestamp = self._parent._get_current_timestamp()
        self._init = False
        self.update()

    def update_profile(self):
        results = self.results[self._current_res_id]
        reach = 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):
        results = self.results[self._current_res_id]
        reach = results.river.reach(self._current_reach_id)
        poly_l_x, poly_l_y, poly_r_x, poly_r_y = [], [], [], []

        for profile in reach.profiles:
            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):
        results = self.results[self._current_res_id]
        reach = 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:
            try:
                plot[0].remove()
            except Exception as e:
                logger_exception(e)
            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)
