import gc
import os
import copy
from typing import Tuple
import matplotlib.cm as cm
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import DroneScanner.utils.utils as utils
from matplotlib.ticker import AutoMinorLocator


class Waterfall:
    """
    Класс для вывода графика типа "Водопад"

    Атрибуты:
        _freqs: Список частот.
        _delay: Количество чтений, после которых обновляем водопад.
        _size: Размер водопада по оси Y.
        _counter_til_init: Счетчик от _size до 0. Отвечает за инициализацию сolors и valuest_to_plot.
        _skip_update_counter:
        _points_scale: Масштаб точек графика в зависимости от количества частот.
        debug_flag: Флаг для отладки класса.
        _save_plots:
        _cur_plot_file_idx:
        _max_plot_files:
        _dir_to_save:
        colors: Массив цветов точек графика.
        values_to_plot: Значения по оси Y.

    """

    def __init__(self, freqs: list, delay: int, size: int, debug_flag: bool, save_plots=False, max_plot_files=1000,
                 dir_to_save='/home/orangepi/plots'):
        """
        Инициализирует новый водопад.

        :param freqs: Список частот.
        :param delay: Количество чтений, после которых обновляем водопад.
        :param size: Размер водопада по оси Y.
        :param debug_flag: Флаг для отладки класса.
        :param save_plots:
        :param max_plot_files:
        :param dir_to_save:
        """
        self._freqs = freqs
        self._delay = delay
        self._size = size
        self._counter_til_init = size - 1
        self._skip_update_counter = 1
        self._points_scale = 500 / len(self._freqs)
        self.debug_flag = debug_flag
        self._save_plots = save_plots
        self._cur_plot_file_idx = 0
        self._max_plot_files = max_plot_files
        self._dir_to_save = dir_to_save
        self.colors = []
        self.values_to_plot = []

        if not os.path.exists(self._dir_to_save):
            os.makedirs(self._dir_to_save)

    def interpolate_color(self, start_color: tuple, end_color: tuple, factor: float) -> tuple:
        """
        Интерполяция цвета между start_color и end_color в зависимости от factor.
        Чем больше factor, тем ближе end_color, меньше - ближе к start_color.
        :param start_color: Тройка (R, G, B), где R, G, B из 0..255.
        :param end_color: Тройка (R, G, B), где R, G, B из 0..255.
        :param factor: Число от 0 до 1.
        :return: Тройка (a, b, c), где 0 <= a,b,c => 1
        """
        return tuple((start + (end - start) * factor) / 255 for start, end in zip(start_color, end_color))

    def decorate(self, data: list) -> list:
        """
        По строчке с датой строим строчку соответствующих цветов в зависимости от элементов из data.
        :param data: Строка длиной len(freqs) с характеристиками сигнала (предобработанным сигналом).
        :return: Строка длиной len(freqs) с тройками, соответствующие некоторым цветам из (R, G, B).
        """
        green = (0, 255, 0)
        red = (255, 0, 0)
        colored_data = []

        for elem in data:
            colored_data.append(self.interpolate_color(green, red, elem))

        return colored_data

    def transform_value_structure(self) -> Tuple[list, list, list]:
        """
        Приведение одномерного freqs и двумерных values_to_plot и colors в одномерные длиной len(freqs) x size
        для построения по ним водопада.
        :return: Координаты точек по х и y и соответствующие цвета, в которые нужно покрасить точки.
        """
        x = []
        y = []
        z = []
        for i in range(utils.get_num_columns_of_array(self.values_to_plot)):
            for j in range(utils.get_num_rows_of_array(self.values_to_plot) - 1, -1, -1):
                x.append(self._freqs[i])
                y.append(self.values_to_plot[j][i])
                z.append(self.colors[j][i])
        return x, y, z

    def plot(self) -> None:
        """
        Построение водопада при помощи plt.scatter по x,y, с=z, s=_points_scale и автомасштабируемой сеткой.
        :return: None.
        """
        x, y, colors = self.transform_value_structure()
        # colors_hex = [mcolors.to_hex(color) for color in colors]
        plt.scatter(x, y, c=colors, s=self._points_scale, marker='s')

        # Добавление цветовой шкалы.
        plt.colorbar(cm.ScalarMappable(norm=mcolors.Normalize(vmin=0, vmax=1), cmap='RdYlGn_r'), label='Значение')

        # Настройка ограничения по ОY
        plt.ylim(-self._size / 10, (self._size - 1) + (self._size / 10))

        # Настройка заголовков
        plt.xlabel('Частоты')
        plt.ylabel('Итерации')
        plt.title('Waterfall')

        # Добавление автомасштабируемой сетки.
        plt.grid(which='major', color='gray', linestyle='-', linewidth=0.5)
        plt.minorticks_on()
        plt.gca().xaxis.set_minor_locator(AutoMinorLocator())
        plt.gca().yaxis.set_minor_locator(AutoMinorLocator())
        plt.grid(which='minor', color='gray', linestyle=':', linewidth=0.5)

        if self._save_plots:
            plt.savefig(f'{self._dir_to_save}/waterfall_{self._cur_plot_file_idx}.png')
            if self._cur_plot_file_idx == self._max_plot_files:
                self._cur_plot_file_idx = 0
            else:
                self._cur_plot_file_idx += 1
        else:
            plt.show()

        plt.pause(0.001)
        plt.close()
        gc.collect()

    def update(self, data: list) -> None:
        """
        Конструирование массивов values_to_plot (значения по Y) и colors (цвета точек) с последующим
        обновлением colors (см. документацию).
        :param data:
        :return: None.
        """
        if self._skip_update_counter == self._delay:
            data_data = copy.deepcopy(data)
            colored_data = self.decorate(data_data)

            if self._counter_til_init != -1:
                self.values_to_plot.append([self._counter_til_init for _ in range(len(self._freqs))])
                self._counter_til_init -= 1
            else:
                self.colors.pop()

            self.colors.insert(0, colored_data)
            self.plot()
            self._skip_update_counter = 1

        else:
            self._skip_update_counter += 1
