# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

"""This package is responsible for graph-related operations using matplotlib."""

from collections import deque
from math import ceil

from matplotlib import animation, pyplot, style
from matplotlib.backends.backend_tkagg import (FigureCanvasTkAgg,
                                               NavigationToolbar2Tk,)
from tkinter import Toplevel

from ..common import err_symbol, update_period

class DataPlotter():
    def __init__(self):
        style.use("ggplot")
        self.dyn_data = []
        self.dyn_period = 120
        self.line_color = "#39A2DB"
        self.line_style = "-"

    def set_dyn_data(self, data):
        self.dyn_data = data

    def create_figure(self, subfig_nums, cols_per_row):
        fig = pyplot.Figure(figsize=(16, 12), dpi=100)
        subfigs = [None] * subfig_nums
        rows = ceil(subfig_nums / cols_per_row)

        for i in range(subfig_nums):
            subfigs[i] = fig.add_subplot(rows, cols_per_row, i + 1)

        fig.set_tight_layout(True)

        # Create the graph window
        graph_window = Toplevel()

        # Create the tkinter canvas containing the matplotlib figure then draw
        canvas = FigureCanvasTkAgg(fig, master=graph_window)
        canvas.draw()

        # Create the matplotlib toolbar then place the toolbar on the window
        toolbar = NavigationToolbar2Tk(canvas, graph_window)
        toolbar.update()

        # Pack the canvas
        canvas.get_tk_widget().pack()

        return graph_window, fig, subfigs

    def plot_static(self, y_data_list, selected_channels, cols_per_row=1):
        subfig_nums = len(y_data_list)
        _, fig, subfigs = self.create_figure(subfig_nums, cols_per_row)

        for i in range(subfig_nums):
            y_data = y_data_list[i]
            x_data = range(len(y_data))  # Use index as data of x-axis
            subfigs[i].clear()
            subfigs[i].set_title(selected_channels[i])
            subfigs[i].plot(x_data,
                            y_data,
                            color=self.line_color,
                            linestyle=self.line_style)

    def plot_dynamic(self, selected_channels, stop_dyn_func, cols_per_row=1):
        subfig_nums = len(selected_channels)
        self.dyn_graph_w, fig, subfigs = self.create_figure(subfig_nums,
                                                            cols_per_row)
        self.dyn_graph_w.bind("<Destroy>",
                              lambda event, callback_func=stop_dyn_func:
                              self.on_dyn_graph_destroy(event,
                                                        callback_func))
        lines = self.init_dyn_data(subfigs, selected_channels)

        # Set animation callback function
        self.ani = animation.FuncAnimation(fig,
                                           self.update_dyn_data,
                                           interval=update_period,
                                           fargs=(lines,
                                                  subfigs,
                                                  selected_channels))

    def on_dyn_graph_destroy(self, event, callback_func):
        # <Destroy> event occurs on all widgets attached on the graph window,
        # we only need to handle when the entire window is destroyed.
        if event.widget == self.dyn_graph_w:
            callback_func()

    def stop_dyn_plot(self):
        self.ani.event_source.stop()

    def init_dyn_data(self, subfigs, selected_channels):
        subfig_nums = len(selected_channels)
        self.circular_buffer_list = [None] * subfig_nums
        lines = [None] * subfig_nums
        for i in range(subfig_nums):
            # Use doubly-ended queue with max length as circular buffer
            # Therfore, each element of circular_buffer_list is circular buffer
            self.circular_buffer_list[i] = deque(maxlen=self.dyn_period)
            subfigs[i].set_title(selected_channels[i])
            lines[i] = subfigs[i].plot([],
                                       [],
                                       color=self.line_color,
                                       linestyle=self.line_style)[0]

        return lines

    def update_dyn_data(self, frame_id, lines, subfigs, selected_channels):
        subfig_nums = len(selected_channels)
        for subfig_id, channel in zip(range(subfig_nums), selected_channels):
            # Append new data to the circular buffer if new data != None
            data_key = channel.split("-")
            if (len(data_key) == 3): # Workaround for "XXX-therm-POWER_ATTR"
                new_data = self.dyn_data[(f"{data_key[0]}-{data_key[1]}",
                                          data_key[2])]
            else:
                new_data = self.dyn_data[(data_key[0], data_key[1])]

            if new_data != err_symbol:
                self.circular_buffer_list[subfig_id].append(new_data)

        for i in range(subfig_nums):
            y_data = list(self.circular_buffer_list[i])
            y_max_lim = max(y_data)
            y_min_lim = min(y_data)
            y_lim_padding = (y_max_lim - y_min_lim) / 10
            if y_min_lim == y_max_lim:
                y_lim_padding = 1

            # Update data and x/y limits for plotting
            x_data = range(len(y_data))
            lines[i].set_data(x_data, y_data)
            subfigs[i].set_xlim(0, len(x_data), auto=True)
            subfigs[i].set_ylim(y_min_lim - y_lim_padding,
                                y_max_lim + y_lim_padding,
                                auto=True)

        return lines

# Instantiate a singleton
data_plotter = DataPlotter()