# PlotSedAdisDx.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 datetime import datetime
from math import inf

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

from PyQt5.QtCore import (
    QCoreApplication
)

_translate = QCoreApplication.translate

logger = logging.getLogger()


class PlotAdis(PamhyrPlot):
    def __init__(self, canvas=None, trad=None, toolbar=None,
                 results=None, reach_id=0, profile_id=0,
                 pol_id=[0], key="C", type_pol=[7], parent=None):
        super(PlotAdis, self).__init__(
            canvas=canvas,
            trad=trad,
            data=results,
            toolbar=toolbar,
            parent=parent
        )

        self._current_timestamp = max(results.get("timestamps"))
        self._current_reach_id = reach_id
        self._current_profile_id = profile_id
        self._current_pol_id = pol_id
        self._key = key
        self._type_pol = type_pol

        self.label_x = self._trad["unit_rk"]
        self._available_values_y = self._trad.get_dict("values_y_pol")

        self.label = {}
        self.label_max = {}
        self.label_min = {}
        self.label["C"] = _translate("Results", "Concentration")
        self.label_max["C"] = _translate("Results", "Max Concentration")
        self.label_min["C"] = _translate("Results", "Min Concentration")
        self.label["T"] = _translate("Results", "Temperature")
        self.label_max["T"] = _translate("Results", "Max Temperature")
        self.label_min["T"] = _translate("Results", "Min Temperature")
        if self._type_pol[pol_id[0]] == -1:  # Total
            self.label["M"] = _translate("Results", "Thickness")
            self.label_max["M"] = _translate("Results", "Max Thickness")
            self.label_min["M"] = _translate("Results", "Min Thickness")
            self.label_y = self._available_values_y["unit_thickness"]
        else:
            self.label["M"] = _translate("Results", "Mass")
            self.label_max["M"] = _translate("Results", "Max Mass")
            self.label_min["M"] = _translate("Results", "Min Mass")
            self.label_y = self._available_values_y["unit_"+self._key]
        self.val_id = {}
        self.val_id["C"] = 0
        self.val_id["T"] = 0
        self.val_id["G"] = 1
        self.val_id["M"] = 2
        self.val_id["D"] = 3

        self._lines = []

        self._isometric_axis = False

    @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)
        profile = reach.profile(self._current_profile_id)

        if reach.geometry.number_profiles == 0:
            self._init = False
            return

        self.draw_min_and_max()
        self.draw_data()
        self.draw_current()

        self.idle()
        self._init = True

    def draw_min_and_max(self):
        x, y, z = self.get_min_and_max()

        self._line_max, = self.canvas.axes.plot(
            x, y,
            label=self.label_max[self._key],
            color=self.color_plot_highlight,
            linestyle='dotted',
            lw=1.,
            markersize=0
        )

        self._line_min, = self.canvas.axes.plot(
            x, z,
            label=self.label_min[self._key],
            color='g',
            linestyle='dotted',
            lw=1.,
            markersize=0
        )

    def update_min_and_max(self):
        x, y, z = self.get_min_and_max()

        self._line_max.set_data(x, y)
        self._line_min.set_data(x, z)

    def set_reach(self, reach_id):
        self._current_reach_id = reach_id
        self._current_profile_id = 0
        self.draw()
        self.update_current()

    def update(self):
        if not self._init:
            self.draw()

        self.update_min_and_max()
        self.update_data()
        self.update_idle()

    def hide_current(self):
        self._current.set_visible(False)

    def show_current(self):
        self._current.set_visible(True)

    def set_pollutant(self, pol_id):
        for j in self._lines:
            j.remove()
        self._lines = []
        self._current_pol_id = pol_id
        self.update_min_and_max()
        self.draw_data()
        self.update_current()
        self.update_idle()


class PlotAdis_dx(PlotAdis):
    def __init__(self, canvas=None, trad=None, toolbar=None,
                 results=None, reach_id=0, profile_id=0,
                 pol_id=[0], key="C", type_pol=[7], parent=None):
        super(PlotAdis_dx, self).__init__(
            canvas=canvas,
            trad=trad,
            toolbar=toolbar,
            results=results,
            reach_id=reach_id,
            profile_id=profile_id,
            pol_id=pol_id,
            key=key,
            type_pol=type_pol,
            parent=parent
        )

    def draw_data(self):
        reach = self.results.river.reach(self._current_reach_id)
        profile = reach.profile(self._current_profile_id)
        x = reach.geometry.get_rk()

        self._lines = []
        for i in range(len(self._current_pol_id)):
            pol_id = self._current_pol_id[i]
            name = self.data.pollutants_list[pol_id]
            if self.val_id[self._key] > 0 and self._type_pol[pol_id] == 1:
                y = [0.0]*len(x)
            else:
                y = list(map(lambda p:
                             p.get_ts_key(
                                 self._current_timestamp, "pols"
                             )[pol_id][self.val_id[self._key]],
                             reach.profiles))
            label = self.label[self._key]
            if name != "TEM":
                label += (" "+name)
            self._lines.append(self.canvas.axes.plot(
                x, y,
                label=label,
                color=self.colors[pol_id % len(self.colors)],
                **self.plot_default_kargs
            )[0])

        self.enable_legend()

    def update_data(self):
        reach = self.results.river.reach(self._current_reach_id)
        profile = reach.profile(self._current_profile_id)

        x = reach.geometry.get_rk()

        for i in range(len(self._current_pol_id)):
            pol_id = self._current_pol_id[i]
            if self.val_id[self._key] > 0 and self._type_pol[pol_id] == 1:
                continue  # no mass
            else:
                y = list(map(
                    lambda p: p.get_ts_key(
                                self._current_timestamp, "pols"
                            )[pol_id][self.val_id[self._key]],
                    reach.profiles))

            self._lines[i].set_data(x, y)

        self.update_current()

    def draw_current(self):
        reach = self.results.river.reach(self._current_reach_id)
        rk = reach.geometry.get_rk()

        self._current, = self.canvas.axes.plot(
            [rk[self._current_profile_id], rk[self._current_profile_id]],
            self.canvas.axes.get_ylim(),
            color="grey",
            linestyle="dashed",
            lw=1.,
        )

    def update_current(self):
        self.hide_current()
        self.canvas.axes.relim(visible_only=True)
        self.canvas.axes.autoscale_view()
        self.show_current()
        reach = self.results.river.reach(self._current_reach_id)
        rk = reach.geometry.get_rk()
        self._current.set_data(
            [rk[self._current_profile_id], rk[self._current_profile_id]],
            self.canvas.axes.get_ylim()
        )

    def get_min_and_max(self):
        reach = self.results.river.reach(self._current_reach_id)
        profile = reach.profile(self._current_profile_id)
        x = reach.geometry.get_rk()

        y = [-inf]*reach.geometry.number_profiles
        z = [inf]*reach.geometry.number_profiles
        for pol_id in self._current_pol_id:
            if self.val_id[self._key] > 0 and self._type_pol[pol_id] == 1:
                for i in range(reach.geometry.number_profiles):
                    y[i] = max(0.0, y[i])
                    z[i] = min(0.0, z[i])
            else:
                val_id = self.val_id[self._key]
                for i, p in enumerate(reach.profiles):
                    rk_y = list(map(lambda data_el:
                                    data_el[pol_id][val_id],
                                    p.get_key("pols")
                                    ))
                    y[i] = max(np.max(rk_y), y[i])
                    z[i] = min(np.min(rk_y), z[i])

        return x, y, z

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

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


class PlotAdis_dt(PlotAdis):
    def __init__(self, canvas=None, trad=None, toolbar=None,
                 results=None, reach_id=0, profile_id=0,
                 pol_id=[0], key="C", type_pol=[7], parent=None):
        super(PlotAdis_dt, self).__init__(
            canvas=canvas,
            trad=trad,
            toolbar=toolbar,
            results=results,
            reach_id=reach_id,
            profile_id=profile_id,
            pol_id=pol_id,
            key=key,
            type_pol=type_pol,
            parent=parent
        )
        self.label_x = self._trad["unit_time_s"]

    def draw_data(self):
        reach = self.results.river.reach(self._current_reach_id)
        profile = reach.profile(self._current_profile_id)
        self.ts = list(self.results.get("timestamps"))
        self.ts.sort()

        self._lines = []
        x = self.ts
        for i in range(len(self._current_pol_id)):
            pol_id = self._current_pol_id[i]
            name = self.data.pollutants_list[pol_id]
            if self.val_id[self._key] > 0 and self._type_pol[pol_id] == 1:
                y = [0.0]*len(x)
            else:
                val_id = self.val_id[self._key]
                y = list(map(lambda data_el:
                             data_el[pol_id][val_id],
                             profile.get_key("pols")
                             ))

            self._lines.append(self.canvas.axes.plot(
                x, y,
                label=self.label[self._key]+" "+name,
                color=self.colors[pol_id % len(self.colors)],
                **self.plot_default_kargs
            )[0])

        self.set_ticks_time_formater()
        self.enable_legend()

    def update_data(self):
        reach = self.results.river.reach(self._current_reach_id)
        profile = reach.profile(self._current_profile_id)

        x = self.ts
        for i in range(len(self._current_pol_id)):
            pol_id = self._current_pol_id[i]
            if self.val_id[self._key] > 0 and self._type_pol[pol_id] == 1:
                continue  # no mass
            else:
                y = list(map(lambda data_el:
                             data_el[pol_id][self.val_id[self._key]],
                             profile.get_key("pols")
                             ))

            self._lines[i].set_data(x, y)

        self._current.set_data(
            [self._current_timestamp, self._current_timestamp],
            self.canvas.axes.get_ylim()
        )

        self.update_current()

    def draw_current(self):
        self._current, = self.canvas.axes.plot(
            [self._current_timestamp, self._current_timestamp],
            self.canvas.axes.get_ylim(),
            color="grey",
            linestyle="dashed",
            lw=1.,
        )

    def update_current(self):
        self.hide_current()
        self.canvas.axes.relim(visible_only=True)
        self.canvas.axes.autoscale_view()
        self.show_current()
        self._current.set_data(
            [self._current_timestamp, self._current_timestamp],
            self.canvas.axes.get_ylim()
        )

    def get_min_and_max(self):
        reach = self.results.river.reach(self._current_reach_id)
        self.ts = list(self.results.get("timestamps"))
        self.ts.sort()

        x = self.ts

        y = [-inf]*len(x)
        z = [inf]*len(x)
        for pol_id in self._current_pol_id:
            if self.val_id[self._key] > 0 and self._type_pol[pol_id] == 1:
                for i in range(len(x)):
                    y[i] = max(0.0, y[i])
                    z[i] = min(0.0, z[i])
            else:
                val_id = self.val_id[self._key]
                for i, ts in enumerate(x):
                    ts_y = list(map(lambda p:
                                    p.get_ts_key(
                                        ts, "pols"
                                    )[pol_id][val_id],
                                    reach.profiles
                                    ))
                    y[i] = max(np.max(ts_y), y[i])
                    z[i] = min(np.min(ts_y), z[i])

        return x, y, z

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

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