# Plot.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
from datetime import datetime

import numpy as np

from numpy import sqrt
from numpy import asarray

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

from View.Results.CustomPlot.Translate import CustomPlotTranslate

logger = logging.getLogger()

unit = {
    "bed_elevation": "0-meter",
    "bed_elevation_envelop": "0-meter",
    "water_elevation": "0-meter",
    "water_elevation_envelop": "0-meter",
    "discharge": "1-m3s",
    "discharge_envelop": "1-m3s",
    "velocity": "2-ms",
    "velocity_envelop": "2-ms",
    "depth": "3-meter",
    "depth_envelop": "3-meter",
    "mean_depth": "3-meter",
    "froude": "4-dimensionless",
    "wet_area": "5-m2",
}


class CustomPlot(PamhyrPlot):
    def __init__(self, x, y, envelop, reach, profile, timestamp,
                 data=None, canvas=None, trad=None, res_id=[0],
                 toolbar=None, parent=None):
        super(CustomPlot, self).__init__(
            canvas=canvas,
            trad=CustomPlotTranslate(),
            data=data,
            toolbar=toolbar,
            parent=parent
        )

        self._x = x
        self._y = y
        self._envelop = envelop
        self._current_reach = reach
        self._current_profile_id = profile
        self._current_timestamp = timestamp
        self._current_res_id = res_id[0]
        self._parent = parent

        logger.debug(
            "Create custom plot for: " +
            f"{x} -> {','.join(y)}: " +
            f"reach={reach}, profile={profile}, " +
            f"timestamp={timestamp}"
        )

        self._y_axes = sorted(
            set(
                map(
                    lambda y: unit[y],
                    self._y
                )
            )
        )

        self._axes = {}
        self.lines = {}

    def draw_bottom_with_bedload(self, reach):
        results = self.data[self._current_res_id]
        table = results.get("table")["zfd"]
        ts = results.get_timestamp_id(self._current_timestamp)

        # rk = reach.geometry.get_rk()
        z = list(
            map(
                lambda p: table[ts, p.global_index],
                reach.profiles
            )
        )

        return z

    def get_ts_zmin(self, profile, res_id):
        results = self.data[res_id]
        reach = results.river.reach(self._current_reach)

        table = results.get("table")["zfd"]

        zfd = table[:, reach.profile(profile).global_index]
        return zfd

    def _draw_rk(self):
        results = self.data[self._current_res_id]
        reach = results.river.reach(self._current_reach)
        if self._current_res_id == 2:  # compare results
            reach1 = self.data[0].river.reach(self._current_reach)
            reach2 = self.data[1].river.reach(self._current_reach)
        rk = reach.geometry.get_rk()
        if reach.has_bedload():
            z_min = self.draw_bottom_with_bedload(reach)
        else:
            z_min = reach.geometry.get_z_min()

        table = results.get("table")
        id_ts = results.get_timestamp_id(self._current_timestamp)

        q = list(
            map(
                lambda p: table["Q"][id_ts, p.global_index],
                reach.profiles
            )
        )
        z = list(
            map(
                lambda p: table["Z"][id_ts, p.global_index],
                reach.profiles
            )
        )
        v = list(
            map(
                lambda p: table["V"][id_ts, p.global_index],
                reach.profiles
            )
        )

        shift = 0
        compt = 0
        for ax in sorted(self._axes):
            if compt == 0:
                self._axes[ax].spines['left'].set_position(('outward', shift))
                compt += 1
            else:
                self._axes[ax].spines['right'].set_position(('outward', shift))
                shift += 60

        if len(self.lines) != 0:
            self.lines = {}

        if "bed_elevation" in self._y:
            ax = self._axes[unit["bed_elevation"]]
            if self._current_res_id < 2:
                if reach.has_bedload():
                    dz = self.draw_bottom_with_bedload(reach)
                else:
                    dz = z_min

                line = ax.plot(
                    rk, dz,
                    color='grey', lw=1.,
                )
            else:
                if reach.has_bedload():
                    z_min1 = list(map(
                        lambda p: table["zfd"][id_ts, p.global_index],
                        reach1.profiles
                    ))
                    z_min2 = list(map(
                        lambda p: table["zfd"][id_ts, p.global_index],
                        reach2.profiles
                    ))
                else:
                    z_min1 = reach1.geometry.get_z_min()
                    z_min2 = reach2.geometry.get_z_min()
                dz = list(
                    map(
                        lambda x, y: x - y,
                        z_min1, z_min2
                        )
                    )
                line = ax.plot(
                    rk, dz,
                    color='grey', lw=1.,
                )

            self.lines["bed_elevation"] = line

            if (self._envelop and
                    reach.has_bedload() and
                    self._current_res_id < 2):
                ax = self._axes[unit["bed_elevation_envelop"]]

                e = list(
                        map(
                            lambda p: max(self.get_ts_zmin(
                                p, self._current_res_id)),
                            range(len(reach))
                        )
                    )
                line1 = ax.plot(
                    rk, e,
                    color='grey', lw=1.,
                    linestyle='dotted',
                )
                self.lines["bed_elevation_envelop"] = line1

                e = list(
                        map(
                            lambda p: min(self.get_ts_zmin(
                                p, self._current_res_id)),
                            range(len(reach))
                        )
                    )
                line2 = ax.plot(
                    rk, e,
                    color='grey', lw=1.,
                    linestyle='dotted',
                )
                # self.lines["bed_elevation_envelop"] = line2

        if "water_elevation" in self._y:
            ax = self._axes[unit["water_elevation"]]
            line = ax.plot(
                rk, z, lw=1.,
                color='blue',
            )
            self.lines["water_elevation"] = line

            if "bed_elevation" in self._y and self._current_res_id < 2:
                self.fill = ax.fill_between(
                    rk, z_min, z,
                    color='blue', alpha=0.5, interpolate=True
                )

            if self._envelop:
                ax = self._axes[unit["water_elevation_envelop"]]

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

                line1 = ax.plot(
                    rk, d, lw=1.,
                    color='blue',
                    linestyle='dotted',
                )
                self.lines["water_elevation_envelop"] = line1

                d = list(
                    map(
                        lambda p: np.min(table["Z"][:, p.global_index]),
                        reach.profiles
                    )
                )

                line2 = ax.plot(
                    rk, d, lw=1.,
                    color='blue',
                    linestyle='dotted',
                )
                # self.lines["water_elevation_envelop2"] = line2

        if "discharge" in self._y:
            ax = self._axes[unit["discharge"]]
            line = ax.plot(
                rk, q, lw=1.,
                color='r',
            )
            self.lines["discharge"] = line

            if self._envelop:
                ax = self._axes[unit["discharge_envelop"]]

                q1 = list(
                    map(
                        lambda p: np.max(table["Q"][:, p.global_index]),
                        reach.profiles
                    )
                )
                line1 = ax.plot(
                    rk, q1, lw=1.,
                    color='r',
                    linestyle='dotted',
                )
                self.lines["discharge_envelop"] = line1

                q2 = list(
                    map(
                        lambda p: np.min(table["Q"][:, p.global_index]),
                        reach.profiles
                    )
                )
                line2 = ax.plot(
                    rk, q2, lw=1.,
                    color='r',
                    linestyle='dotted',
                )
                # self.lines["discharge_envelop2"] = line2

        if "velocity" in self._y:
            ax = self._axes[unit["velocity"]]

            line = ax.plot(
                rk, v, lw=1.,
                color='g',
            )
            self.lines["velocity"] = line

            if self._envelop:

                velocities = list(
                    map(
                        lambda p: table["V"][:, p.global_index],
                        reach.profiles
                    )
                )

                ax = self._axes[unit["velocity_envelop"]]
                vmax = [max(v) for v in velocities]

                line1 = ax.plot(
                    rk, vmax, lw=1.,
                    color='g', linestyle='dotted',
                )
                self.lines["velocity_envelop"] = line1
                vmin = [min(v) for v in velocities]

                line2 = ax.plot(
                    rk, vmin, lw=1.,
                    color='g',
                    linestyle='dotted',
                )
                # self.lines["velocity_envelop2"] = line2

        if "depth" in self._y:
            ax = self._axes[unit["depth"]]
            if self._current_res_id < 2:
                d = list(
                    map(
                        lambda p: p.geometry.max_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach.profiles
                    )
                )
            else:
                d1 = list(
                    map(
                        lambda p: p.geometry.max_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach1.profiles
                    )
                )
                d2 = list(
                    map(
                        lambda p: p.geometry.max_water_depth(
                            table["Z"][id_ts, p.global_index]),
                        reach2.profiles
                    )
                )
                d = list(map(
                    lambda x, y: x - y,
                    d1, d2
                ))

            line = ax.plot(
                rk, d,
                color='brown', lw=1.,
            )
            self.lines["depth"] = line

            if self._envelop and self._current_res_id < 2:
                ax = self._axes[unit["depth_envelop"]]

                d = list(map(
                    lambda p1, p2: p1 - p2, map(
                        lambda p: np.max(table["Z"][:, p.global_index]),
                        reach.profiles
                    ), z_min)
                )
                line1 = ax.plot(
                    rk, d,
                    color='brown', lw=1.,
                    linestyle='dotted',
                )
                self.lines["depth_envelop"] = line1

                d = list(map(
                    lambda p1, p2: p1 - p2, map(
                        lambda p: np.min(table["Z"][:, p.global_index]),
                        reach.profiles
                    ), z_min)
                )
                line2 = ax.plot(
                    rk, d,
                    color='brown', lw=1.,
                    linestyle='dotted',
                )
                # self.lines["depth_envelop2"] = line2

        if "mean_depth" in self._y:
            ax = self._axes[unit["mean_depth"]]
            if self._current_res_id < 2:
                d = list(
                    map(
                        lambda p: p.geometry.mean_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach.profiles
                    )
                )
            else:
                d1 = list(
                    map(
                        lambda p: p.geometry.mean_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach1.profiles
                    )
                )
                d2 = list(
                    map(
                        lambda p: p.geometry.mean_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach2.profiles
                    )
                )
                d = list(map(
                    lambda x, y: x - y,
                    d1, d2
                ))

            line = ax.plot(
                rk, d,
                color='orange', lw=1.,
            )
            self.lines["mean_depth"] = line

        if "froude" in self._y:
            ax = self._axes[unit["froude"]]

            if self._current_res_id < 2:
                fr = list(
                    map(
                        lambda p:
                            table["V"][id_ts, p.global_index] /
                            sqrt(9.81 * (
                                p.geometry.wet_area(
                                    table["Z"][id_ts, p.global_index]
                                ) /
                                p.geometry.wet_width(
                                    table["Z"][id_ts, p.global_index]
                                )
                            )),
                        reach.profiles
                    )
                )
            else:
                fr1 = list(
                    map(
                        lambda p:
                            table["V"][id_ts, p.global_index] /
                            sqrt(9.81 * (
                                p.geometry.wet_area(
                                    table["Z"][id_ts, p.global_index]
                                ) /
                                p.geometry.wet_width(
                                    table["Z"][id_ts, p.global_index]
                                )
                            )),
                        reach1.profiles
                    )
                )
                fr2 = list(
                    map(
                        lambda p:
                            table["V"][id_ts, p.global_index] /
                            sqrt(9.81 * (
                                p.geometry.wet_area(
                                    table["Z"][id_ts, p.global_index]
                                ) /
                                p.geometry.wet_width(
                                    table["Z"][id_ts, p.global_index]
                                )
                            )),
                        reach2.profiles
                    )
                )
                fr = list(
                    map(
                        lambda x, y: x - y,
                        fr1, fr2
                        )
                    )

            line = ax.plot(
                rk, fr, color='black', linestyle='--', lw=1.,
            )
            self.lines["froude"] = line

        if "wet_area" in self._y:

            ax = self._axes[unit["wet_area"]]
            if self._current_res_id < 2:
                d = list(
                    map(
                        lambda p: p.geometry.wet_area(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach.profiles
                    )
                )
            else:
                d1 = list(
                    map(
                        lambda p: p.geometry.wet_area(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach1.profiles
                    )
                )
                d2 = list(
                    map(
                        lambda p: p.geometry.wet_area(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach2.profiles
                    )
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            line = ax.plot(
                rk, d,
                color='blue', linestyle='--', lw=1.,
            )
            self.lines["wet_area"] = line

        # Legend
        self.update_legend()

    def _redraw_rk(self):
        results = self.data[self._current_res_id]
        reach = results.river.reach(self._current_reach)
        if self._current_res_id == 2:  # compare results
            reach1 = self.data[0].river.reach(self._current_reach)
            reach2 = self.data[1].river.reach(self._current_reach)

        table = results.get("table")
        id_ts = results.get_timestamp_id(self._current_timestamp)

        rk = reach.geometry.get_rk()
        z_min = reach.geometry.get_z_min()
        if reach.has_bedload():
            z_min = self.draw_bottom_with_bedload(reach)
        else:
            z_min = reach.geometry.get_z_min()

        q = list(
            map(
                lambda p: table["Q"][id_ts, p.global_index],
                reach.profiles
            )
        )
        z = list(
            map(
                lambda p: table["Z"][id_ts, p.global_index],
                reach.profiles
            )
        )
        v = list(
            map(
                lambda p: table["Z"][id_ts, p.global_index],
                reach.profiles
            )
        )
        if "bed_elevation" in self._y:
            if self._current_res_id < 2:
                if reach.has_bedload():
                    dz = self.draw_bottom_with_bedload(reach)
                else:
                    dz = z_min
            else:
                if reach.has_bedload():
                    z_min1 = self.draw_bottom_with_bedload(reach1)
                    z_min2 = self.draw_bottom_with_bedload(reach2)
                else:
                    z_min1 = reach1.geometry.get_z_min()
                    z_min2 = reach2.geometry.get_z_min()
                dz = list(
                    map(
                        lambda x, y: x - y,
                        z_min1, z_min2
                        )
                    )
            self.lines["bed_elevation"][0].set_ydata(dz)

        if "water_elevation" in self._y:
            self.lines["water_elevation"][0].set_ydata(z)

            if "bed_elevation" in self._y and self._current_res_id < 2:
                ax = self._axes[unit["water_elevation"]]
                self.fill.remove()
                self.fill = ax.fill_between(
                    rk, z_min, z,
                    color='blue', alpha=0.5, interpolate=True
                )

        if "discharge" in self._y:
            self.lines["discharge"][0].set_ydata(q)

        if "velocity" in self._y:
            self.lines["velocity"][0].set_ydata(v)

        if "depth" in self._y:
            if self._current_res_id < 2:
                d = list(
                    map(
                        lambda p: p.geometry.max_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach.profiles
                    )
                )
            else:
                d1 = list(
                    map(
                        lambda p: p.geometry.max_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach1.profiles
                    )
                )
                d2 = list(
                    map(
                        lambda p: p.geometry.max_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach2.profiles
                    )
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            self.lines["depth"][0].set_ydata(d)

        if "mean_depth" in self._y:
            if self._current_res_id < 2:
                d = list(
                    map(
                        lambda p: p.geometry.mean_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach.profiles
                    )
                )
            else:
                d1 = list(
                    map(
                        lambda p: p.geometry.mean_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach1.profiles
                    )
                )
                d2 = list(
                    map(
                        lambda p: p.geometry.mean_water_depth(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach2.profiles
                    )
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            self.lines["mean_depth"][0].set_ydata(d)

        if "froude" in self._y:
            if self._current_res_id < 2:
                fr = list(
                    map(
                        lambda p: (
                            table["V"][id_ts, p.global_index] /
                            sqrt(9.81 * (
                                p.geometry.wet_area(
                                    table["Z"][id_ts, p.global_index]
                                ) /
                                p.geometry.wet_width(
                                    table["Z"][id_ts, p.global_index]
                                )
                            ))
                        ),
                        reach.profiles
                    )
                )
            else:
                fr1 = list(
                    map(
                        lambda p: (
                            table["V"][id_ts, p.global_index] /
                            sqrt(9.81 * (
                                p.geometry.wet_area(
                                    table["Z"][id_ts, p.global_index]
                                ) /
                                p.geometry.wet_width(
                                    table["Z"][id_ts, p.global_index]
                                )
                            ))
                        ),
                        reach1.profiles
                    )
                )
                fr2 = list(
                    map(
                        lambda p: (
                            table["V"][id_ts, p.global_index] /
                            sqrt(9.81 * (
                                p.geometry.wet_area(
                                    table["Z"][id_ts, p.global_index]
                                ) /
                                p.geometry.wet_width(
                                    table["Z"][id_ts, p.global_index]
                                )
                            ))
                        ),
                        reach2.profiles
                    )
                )
                fr = list(
                    map(
                        lambda x, y: x - y,
                        fr1, fr2
                        )
                    )

            self.lines["froude"][0].set_ydata(fr)

        if "wet_area" in self._y:
            if self._current_res_id < 2:
                d = list(
                    map(
                        lambda p: p.geometry.wet_area(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach.profiles
                    )
                )
            else:
                d1 = list(
                    map(
                        lambda p: p.geometry.wet_area(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach1.profiles
                    )
                )
                d2 = list(
                    map(
                        lambda p: p.geometry.wet_area(
                            table["Z"][id_ts, p.global_index]
                        ),
                        reach2.profiles
                    )
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            self.lines["wet_area"][0].set_ydata(d)

    def _customize_x_axes_time(self, ts, mode="time"):
        # Custom time display
        nb = len(ts)
        mod = int(nb / 5)
        mod = mod if mod > 0 else nb

        fx = list(
            map(
                lambda x: x[1],
                filter(
                    lambda x: x[0] % mod == 0,
                    enumerate(ts)
                )
            )
        )

        if mode == "time":
            t0 = datetime.fromtimestamp(0)
            xt = list(
                map(
                    lambda v: (
                        str(
                            datetime.fromtimestamp(v) - t0
                        ).split(",")[0]
                        .replace("days", self._trad["days"])
                        .replace("day", self._trad["day"])
                    ),
                    fx
                )
            )
        else:
            xt = list(
                map(
                    lambda v: str(datetime.fromtimestamp(v).date()),
                    fx
                )
            )

        self.canvas.axes.set_xticks(ticks=fx, labels=xt, rotation=45)
        self.canvas.axes.relim(visible_only=True)
        self.canvas.axes.autoscale_view()

    def _draw_time(self):
        results = self.data[self._current_res_id]
        reach = results.river.reach(self._current_reach)
        profile = reach.profile(self._current_profile_id)

        shift = 0
        compt = 0
        for ax in sorted(self._axes):
            if compt == 0:
                self._axes[ax].spines['left'].set_position(('outward', shift))
                compt += 1
            else:
                self._axes[ax].spines['right'].set_position(('outward', shift))
                shift += 60

        ts = self._parent._timestamps
        table = results.get("table")

        if self._current_res_id == 2:  # compare results
            reach1 = self.data[0].river.reach(self._current_reach)
            reach2 = self.data[1].river.reach(self._current_reach)
            profile1 = reach1.profile(self._current_profile_id)
            profile2 = reach2.profile(self._current_profile_id)

            q1 = table["Q"][:, profile1.global_index]
            z1 = table["Z"][:, profile1.global_index]
            v1 = table["V"][:, profile1.global_index]

            q2 = table["Q"][:, profile2.global_index]
            z2 = table["Z"][:, profile2.global_index]
            v2 = table["V"][:, profile2.global_index]

        q = table["Q"][:, profile.global_index]
        z = table["Z"][:, profile.global_index]
        v = table["V"][:, profile.global_index]

        z_min = profile.geometry.z_min()
        if self._current_res_id < 2:
            if reach.has_bedload():
                ts_z_min = self.get_ts_zmin(
                    self._current_profile_id, self._current_res_id)
            else:
                ts_z_min = list(
                    map(
                        lambda ts:  z_min,
                        ts
                    )
                )

        self.lines = {}
        if "bed_elevation" in self._y:
            # Z min is constant in time

            ax = self._axes[unit["bed_elevation"]]

            if self._current_res_id == 2:
                if reach.has_bedload():
                    ts_z_min1 = self.get_ts_zmin(self._current_profile_id, 0)
                    ts_z_min2 = self.get_ts_zmin(self._current_profile_id, 1)
                    ts_z_min = list(
                        map(
                            lambda x, y: x - y,
                            ts_z_min1, ts_z_min2
                            )
                        )
                else:
                    z_min1 = profile1.geometry.z_min()
                    z_min2 = profile2.geometry.z_min()
                    ts_z_min = list(
                        map(
                            lambda ts:  z_min1 - z_min2,
                            ts
                        )
                    )

            line = ax.plot(
                ts, ts_z_min,
                color='grey', lw=1.
            )
            self.lines["bed_elevation"] = line

        if "water_elevation" in self._y:

            ax = self._axes[unit["water_elevation"]]
            line = ax.plot(
                ts, z, lw=1.,
                color='b',
            )
            self.lines["water_elevation"] = line

            if "bed_elevation" in self._y and self._current_res_id < 2:

                self.fill = ax.fill_between(
                    ts, ts_z_min, z,
                    color='blue', alpha=0.5, interpolate=True
                )

        if "discharge" in self._y:

            ax = self._axes[unit["discharge"]]
            line = ax.plot(
                ts, q, lw=1.,
                color='r',
            )
            self.lines["discharge"] = line

        if "velocity" in self._y:

            ax = self._axes[unit["velocity"]]

            line = ax.plot(
                ts, v, lw=1.,
                color='g',
            )
            self.lines["velocity"] = line

        if "depth" in self._y:

            ax = self._axes[unit["depth"]]
            if self._current_res_id < 2:
                d = list(
                    map(lambda z: profile.geometry.max_water_depth(z), z)
                )
            else:
                d1 = list(
                    map(lambda z: profile1.geometry.max_water_depth(z), z1)
                )
                d2 = list(
                    map(lambda z: profile2.geometry.max_water_depth(z), z2)
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            line = ax.plot(
                ts, d,
                color='brown', lw=1.,
            )
            self.lines["depth"] = line

        if "mean_depth" in self._y:

            ax = self._axes[unit["mean_depth"]]
            if self._current_res_id < 2:
                d = list(
                    map(lambda z: profile.geometry.mean_water_depth(z), z)
                )
            else:
                d1 = list(
                    map(lambda z: profile1.geometry.mean_water_depth(z), z1)
                )
                d2 = list(
                    map(lambda z: profile2.geometry.mean_water_depth(z), z2)
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            line = ax.plot(
                ts, d,
                color='orange', lw=1.,
            )
            self.lines["mean_depth"] = line

        if "froude" in self._y:

            ax = self._axes[unit["froude"]]
            if self._current_res_id < 2:
                d = list(
                    map(lambda z, v:
                        v /
                        sqrt(9.81 * (
                            profile.geometry.wet_area(z) /
                            profile.geometry.wet_width(z))
                        ), z, v)
                    )
            else:
                d1 = list(
                    map(lambda z, v:
                        v /
                        sqrt(9.81 * (
                            profile1.geometry.wet_area(z) /
                            profile1.geometry.wet_width(z))
                        ), z1, v1)
                    )
                d2 = list(
                    map(lambda z, v:
                        v /
                        sqrt(9.81 * (
                            profile2.geometry.wet_area(z) /
                            profile2.geometry.wet_width(z))
                        ), z2, v2)
                    )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            line = ax.plot(
                ts, d, color='black', linestyle='--', lw=1.,
            )
            self.lines["froude"] = line

        if "wet_area" in self._y:

            ax = self._axes[unit["wet_area"]]
            if self._current_res_id < 2:
                d = list(
                    map(lambda z: profile.geometry.wet_area(z), z)
                )
            else:
                d1 = list(
                    map(lambda z: profile1.geometry.wet_area(z), z1)
                )
                d2 = list(
                    map(lambda z: profile2.geometry.wet_area(z), z2)
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            line = ax.plot(
                ts, d, color='blue', linestyle='--', lw=1.,
            )
            self.lines["wet_area"] = line

        self._customize_x_axes_time(ts)

        # Legend
        self.update_legend()

    def _redraw_time(self):
        results = self.data[self._current_res_id]
        reach = results.river.reach(self._current_reach)
        profile = reach.profile(self._current_profile_id)
        ts = list(results.get("timestamps"))
        ts.sort()

        table = results.get("table")

        if self._current_res_id == 2:  # compare results
            reach1 = self.data[0].river.reach(self._current_reach)
            reach2 = self.data[1].river.reach(self._current_reach)
            profile1 = reach1.profile(self._current_profile_id)
            profile2 = reach2.profile(self._current_profile_id)

            q1 = table["Q"][:, profile1.global_index]
            z1 = table["Z"][:, profile1.global_index]
            v1 = table["V"][:, profile1.global_index]

            q2 = table["Q"][:, profile2.global_index]
            z2 = table["Z"][:, profile2.global_index]
            v2 = table["V"][:, profile2.global_index]

        q = table["Q"][:, profile.global_index]
        z = table["Z"][:, profile.global_index]
        v = table["V"][:, profile.global_index]

        if self._current_res_id < 2:
            if reach.has_bedload():
                ts_z_min = self.get_ts_zmin(
                    self._current_profile_id, self._current_res_id)
            else:
                z_min = profile.geometry.z_min()
                ts_z_min = list(
                    map(
                        lambda ts:  z_min,
                        ts
                    )
                )
        if "bed_elevation" in self._y:
            if self._current_res_id == 2:
                if reach.has_bedload():
                    ts_z_min1 = self.get_ts_zmin(self._current_profile_id, 0)
                    ts_z_min2 = self.get_ts_zmin(self._current_profile_id, 1)
                    ts_z_min = list(
                        map(
                            lambda x, y: x - y,
                            ts_z_min1, ts_z_min2
                            )
                        )
                else:
                    z_min1 = profile1.geometry.z_min()
                    z_min2 = profile2.geometry.z_min()
                    ts_z_min = list(
                        map(
                            lambda ts:  z_min1 - z_min2,
                            ts
                        )
                    )
            self.lines["bed_elevation"][0].set_ydata(ts_z_min)

        if "water_elevation" in self._y:
            self.lines["water_elevation"][0].set_ydata(z)

            if "bed_elevation" in self._y and self._current_res_id < 2:
                ax = self._axes[unit["bed_elevation"]]
                self.fill.remove()
                self.fill = ax.fill_between(
                    ts, ts_z_min, z,
                    color='blue', alpha=0.5, interpolate=True
                )

        if "discharge" in self._y:
            self.lines["discharge"][0].set_ydata(q)

        if "velocity" in self._y:
            self.lines["velocity"][0].set_ydata(v)

        if "depth" in self._y:
            if self._current_res_id < 2:
                d = list(
                    map(lambda z: profile.geometry.max_water_depth(z), z)
                )
            else:
                d1 = list(
                    map(lambda z: profile1.geometry.max_water_depth(z), z1)
                )
                d2 = list(
                    map(lambda z: profile2.geometry.max_water_depth(z), z2)
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            self.lines["depth"][0].set_ydata(d)

        if "mean_depth" in self._y:
            if self._current_res_id < 2:
                d = list(
                    map(lambda z: profile.geometry.mean_water_depth(z), z)
                )
            else:
                d1 = list(
                    map(lambda z: profile1.geometry.mean_water_depth(z), z1)
                )
                d2 = list(
                    map(lambda z: profile2.geometry.mean_water_depth(z), z2)
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            self.lines["mean_depth"][0].set_ydata(d)

        if "froude" in self._y:
            if self._current_res_id < 2:
                d = list(
                    map(lambda z, v:
                        v /
                        sqrt(9.81 * (
                            profile.geometry.wet_area(z) /
                            profile.geometry.wet_width(z))
                        ), z, v)
                    )
            else:
                d1 = list(
                    map(lambda z, v:
                        v /
                        sqrt(9.81 * (
                            profile1.geometry.wet_area(z) /
                            profile1.geometry.wet_width(z))
                        ), z1, v1)
                    )
                d2 = list(
                    map(lambda z, v:
                        v /
                        sqrt(9.81 * (
                            profile2.geometry.wet_area(z) /
                            profile2.geometry.wet_width(z))
                        ), z2, v2)
                    )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            self.lines["froude"][0].set_ydata(d)

        if "wet_area" in self._y:
            if self._current_res_id < 2:
                d = list(
                    map(lambda z: profile.geometry.wet_area(z), z)
                )
            else:
                d1 = list(
                    map(lambda z: profile1.geometry.wet_area(z), z1)
                )
                d2 = list(
                    map(lambda z: profile2.geometry.wet_area(z), z2)
                )
                d = list(
                    map(
                        lambda x, y: x - y,
                        d1, d2
                        )
                    )

            self.lines["wet_area"][0].set_ydata(d)

        self._current.set_visible(False)
        self.canvas.axes.relim(visible_only=True)
        self.canvas.axes.autoscale_view()
        self._current.set_visible(True)

    def draw(self):
        self.draw_static()

    def draw_update(self):
        if self._x == "rk":
            self._redraw_rk()
        elif self._x == "time":
            self._redraw_time()

    @timer
    def draw_static(self):
        self.canvas.axes.cla()
        self.canvas.axes.grid(color='grey', linestyle='--', linewidth=0.5)

        if self.data is None:
            return

        self._y_axes = sorted(
            set(
                map(
                    lambda y: unit[y],
                    self._y
                )
            )
        )

        self.canvas.axes.set_xlabel(
            self._trad[self._x],
            color='black', fontsize=10
        )

        if self._current_res_id < 2:
            self.canvas.axes.set_ylabel(
                self._trad[self._y_axes[0]],
                color='black', fontsize=10
            )
        else:
            self.canvas.axes.set_ylabel(
                "Δ " + self._trad[self._y_axes[0]],
                color='black', fontsize=10
            )

        self._axes[self._y_axes[0]] = self.canvas.axes
        if len(self.lines) != 0:
            self.lines.clear()
        for axes in self._y_axes[1:]:
            if axes in self._axes.keys():
                for ll in self._axes[axes].lines:
                    ll.remove()
            else:
                ax_new = self.canvas.axes.twinx()
                self._axes[axes] = ax_new
            if self._current_res_id < 2:
                self._axes[axes].set_ylabel(
                    self._trad[axes],
                    color='black', fontsize=10
                )
            else:
                self._axes[axes].set_ylabel(
                    "Δ " + self._trad[axes],
                    color='black', fontsize=10
                )

        if self._x == "rk":
            self._draw_rk()
        elif self._x == "time":
            self._draw_time()
        if self._x == "rk":
            reach = self.data[self._current_res_id].river.reach(
                self._current_reach)
            profile = reach.profile(self._current_profile_id)
            x = profile.rk
        elif self._x == "time":
            x = self._current_timestamp

        self._current, = self.canvas.axes.plot(
            [x, x],
            self.canvas.axes.get_ylim(),
            # label=self.label_timestamp,
            color='grey',
            linestyle="dashed",
            lw=1.,
        )

        if self.toolbar is not None:
            self.toolbar.update()
        self.canvas.draw_idle()

    @timer
    def update(self):
        self.draw_update()
        self.draw_current()
        # self.draw_static()
        return

    def set_reach(self, reach_id):
        self._current_reach = reach_id
        self._current_profile_id = 0

        self.draw()

    def set_profile(self, profile_id):
        self._current_profile_id = profile_id

        if self._x != "rk":
            self.update()
        else:
            self.draw_current()

    def set_result(self, res_id):
        self._current_res_id = res_id[0]
        self.draw()

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

        if self._x != "time":
            self.update()
        else:
            self.draw_current()

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

    def draw_current(self):
        if self._x == "rk":
            reach = self.data[self._current_res_id].river.reach(
                self._current_reach)
            profile = reach.profile(self._current_profile_id)
            x = profile.rk
        elif self._x == "time":
            x = self._current_timestamp
        self._current.set_data([x, x], self.canvas.axes.get_ylim())
        self.canvas.draw_idle()

    def update_legend(self):
        lns = reduce(
            lambda acc, line: acc + line,
            map(lambda line: self.lines[line], self.lines),
            []
        )
        labs = list(map(lambda line: self._trad[line], self.lines))
        self.canvas.axes.legend(lns, labs, loc="best")

    def add_imported_plot(self, data):
        if (data["type_x"] == self._x and
                data["type_y"] in self._y):
            while data["legend"] in self.lines:
                data["legend"] += '*'
            self._trad._dict[data["legend"]] = data["legend"] + data["unit"]
            self.lines[data["legend"]] = self.canvas.axes.plot(data["x"],
                                                               data["y"],
                                                               marker="+",
                                                               linestyle="--")
            self.update_legend()
            self.idle()
