# PlotRKC.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
import numpy as np

from functools import reduce

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

from PyQt5.QtCore import (
    QCoreApplication
)

logger = logging.getLogger()


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

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

        self.label_x = self._trad["unit_rk"]
        self.label_y = self._trad["unit_elevation"]

        self.label_bottom = self._trad["label_bottom"]
        self.label_water = self._trad["label_water"]
        self.label_water_max = self._trad["label_water_max"]

        self._isometric_axis = False

        self._auto_relim_update = True
        self._autoscale_update = False

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

        self.draw_bottom(reach)
        self.draw_water_elevation(reach)
        self.draw_water_elevation_max(reach)
        self.draw_water_elevation_overflow(reach)
        self.draw_current(reach)
        self.draw_profiles_hs(reach)

        # self.enable_legend()

        self.idle()
        self._init = True

    def draw_bottom(self, reach):
        if reach.has_bedload():
            self.draw_bottom_with_bedload(reach)
        else:
            self.draw_bottom_geometry(reach)

    def draw_bottom_with_bedload(self, reach):
        results = self.results[self._current_res_id]
        rk = reach.geometry.get_rk()

        table = results.get("table")["zfd"]
        ts = results.get_timestamp_id(self._current_timestamp)

        zfd = list(
            map(
                lambda p: table[ts, p.global_index],
                reach.profiles
            )
        )

        self.line_bottom, = self.canvas.axes.plot(
            rk, zfd,
            linestyle="solid", lw=1.,
            color=self.color_plot_river_bottom,
        )

        self._river_bottom = zfd

    def draw_profiles_hs(self, reach):
        results = self.results[self._current_res_id]
        lhs = filter(
            lambda hs: hs._input_reach.reach is reach.geometry,
            filter(
                lambda hs: hs._input_reach is not None,
                results.study.river.hydraulic_structures.lst
            )
        )

        for hs in lhs:
            x = hs.input_section.rk
            z_min = reach.geometry.get_z_min()
            z_max = reach.geometry.get_z_max()

            self.canvas.axes.plot(
                [x, x],
                [min(z_min), max(z_max)],
                linestyle="--",
                lw=1.,
                color=self.color_plot_previous,
            )

            self.canvas.axes.annotate(
                " > " + hs.name,
                (x, max(z_max)),
                horizontalalignment='left',
                verticalalignment='top',
                annotation_clip=True,
                fontsize=9, color=self.color_plot_previous,
            )

    def draw_bottom_geometry(self, reach):
        rk = reach.geometry.get_rk()
        z_min = reach.geometry.get_z_min()
        z_max = reach.geometry.get_z_max()

        self.line_rk_zmin = self.canvas.axes.plot(
            rk, z_min,
            color=self.color_plot_river_bottom,
            lw=1.
        )

        self._river_bottom = z_min

    def draw_water_elevation(self, reach):
        if len(reach.geometry.profiles) != 0:
            result = self.results[self._current_res_id]
            rk = reach.geometry.get_rk()
            z_min = reach.geometry.get_z_min()

            table = result.get("table")["Z"]
            ts = result.get_timestamp_id(self._current_timestamp)

            water_z = list(
                map(
                    lambda p: table[
                        ts, p.global_index
                    ],
                    reach.profiles
                )
            )

            self.water = self.canvas.axes.plot(
                rk, water_z,
                lw=1., color=self.color_plot_river_water,
            )

            self.water_fill = self.canvas.axes.fill_between(
                rk, self._river_bottom, water_z,
                where=[b <= z for b, z in zip(self._river_bottom, water_z)],
                color=self.color_plot_river_water_zone,
                alpha=0.7,
                interpolate=True
            )

    def draw_water_elevation_max(self, reach):
        if len(reach.geometry.profiles) != 0:
            result = self.results[self._current_res_id]
            rk = reach.geometry.get_rk()
            z_min = reach.geometry.get_z_min()

            table = result.get("table")["Z"]
            ts = result.get_timestamp_id(self._current_timestamp)

            water_z = list(
                map(
                    lambda p: np.max(
                        table[
                            :, p.global_index
                        ]
                    ),
                    reach.profiles
                )
            )

            self.canvas.axes.plot(
                rk, water_z, lw=1.,
                color=self.color_plot_river_water,
                linestyle='dotted',
            )

    def draw_additional_data(self, res_id):
        results = self.results[res_id]

        self._add_data_lines = []

        for data in results.get("additional_data"):
            data = data._data

            tx, ty = data['type_x'], data['type_y']
            x, y = data['x'], data['y']
            legend = data['legend']
            unit = data['unit']

            if tx == 'water_elevation' and ty == 'time':
                line, = self.canvas.axes.plot(
                    x, y, marker="+",
                    label=legend + ' ' + unit
                )

                self._add_data_lines.append(line)

        self.canvas.draw_idle()

    def draw_current(self, reach):
        rk = reach.geometry.get_rk()
        z_min = reach.geometry.get_z_min()
        z_max = reach.geometry.get_z_max()

        self.profile, = self.canvas.axes.plot(
            [
                rk[self._current_profile_id],
                rk[self._current_profile_id]
            ],
            [
                z_max[self._current_profile_id],
                z_min[self._current_profile_id]
            ],
            color=self.color_plot,
            lw=1.
        )

    def draw_water_elevation_overflow(self, reach):
        overflow = []
        result = self.results[self._current_res_id]
        table = result.get("table")["Z"]

        for profile in reach.profiles:
            z_max = np.max(table[:, profile.global_index])
            z_max_ts = 0
            for i, ts in enumerate(self._timestamps):
                z = table[i, profile.global_index]
                if z == z_max:
                    z_max_ts = ts
                    break

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

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

        for profile, z in overflow:
            self.canvas.axes.plot(
                profile.rk, z,
                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 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_current()

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

        results = self.results[self._current_res_id]
        reach = results.river.reach(self._current_reach_id)
        if reach.has_bedload():
            self.update_bottom_with_bedload()

        self.update_water_elevation()

        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_water_elevation(self):
        result = self.results[self._current_res_id]
        reach = result.river.reach(self._current_reach_id)
        rk = reach.geometry.get_rk()
        z_min = reach.geometry.get_z_min()
        table = result.get("table")["Z"]
        ts = result.get_timestamp_id(self._current_timestamp)

        water_z = list(
            map(
                lambda p: table[
                    ts, p.global_index
                ],
                reach.profiles
            )
        )

        self.water[0].set_data(
            rk, water_z
        )

        self.water_fill.remove()
        self.water_fill = self.canvas.axes.fill_between(
            rk, self._river_bottom, water_z,
            where=[rb <= wz for rb, wz in zip(self._river_bottom, water_z)],
            color=self.color_plot_river_water_zone,
            alpha=0.7, interpolate=True
        )

    def update_additional_data(self):
        self.draw_additional_data(0)
        self.update_idle()

    def update_current(self):
        results = self.results[self._current_res_id]
        reach = results.river.reach(self._current_reach_id)
        rk = reach.geometry.get_rk()
        z_min = reach.geometry.get_z_min()
        z_max = reach.geometry.get_z_max()
        cid = self._current_profile_id

        self.profile.set_data(
            [rk[cid], rk[cid]],
            [z_max[cid], z_min[cid]]
        )
        self.canvas.figure.canvas.draw_idle()

    def update_bottom_with_bedload(self):
        results = self.results[self._current_res_id]
        reach = results.river.reach(self._current_reach_id)
        rk = reach.geometry.get_rk()
        # z = self.sl_compute_current_z(reach)

        table = results.get("table")["zfd"]
        ts = results.get_timestamp_id(self._current_timestamp)

        zfd = list(map(
            lambda p: table[ts, p.global_index],
            reach.profiles
        ))

        self.line_bottom.remove()

        self.line_bottom, = self.canvas.axes.plot(
            rk, zfd,
            linestyle="solid", lw=1.,
            color=self.color_plot_river_bottom,
        )

        self._river_bottom = zfd
