import time
import warnings

import matplotlib.pyplot as plt
from matplotlib.rcsetup import cycler

from kolibri.config import TaskType
from kolibri.utils import constants
from kolibri.utils.data_structures import Deque


class PerformanceVisualizer():
    """ This class is responsible for maintaining and updating the plot modules
    for the evaluators.
    
    It uses `matplotlib.pyplot` modules to create the main plot, which
    depending on the options passed to it as parameter, will create multiple 
    subplots to better display all requested metrics.
    """

    def __init__(self, task_type, n_wait, dataset_name, metrics, n_models, model_names, data_dict):
        super().__init__()

        # Default values
        self._sample_ids = []
        self._is_legend_set = False
        self._frame_cnt = 0
        self._plot_trackers = {}
        self._text_annotations = []
        self._last_draw_timestamp = 0

        # Configuration
        self._data_dict = data_dict
        self.n_wait = n_wait
        self.dataset_name = dataset_name
        self.n_models = n_models

        # Validate inputs
        if task_type is None or task_type == constants.UNDEFINED:
            raise ValueError('Task type for visualizer object is undefined.')
        else:
            if task_type in [TaskType.CLASSIFICATION, TaskType.REGRESSION, TaskType.MULTI_TARGET_CLASSIFICATION,
                             TaskType.MULTI_TARGET_REGRESSION]:
                self.task_type = task_type
            else:
                raise ValueError('Invalid task type: {}'.format(task_type))

        if model_names is None:
            self.model_names = ['M{}'.format(i) for i in range(n_models)]
        else:
            if isinstance(model_names, list):
                if len(model_names) != n_models:
                    raise ValueError("Number of model names {} does not match the number of models {}.".
                                     format(len(model_names), n_models))
                else:
                    self.model_names = model_names
            else:
                raise ValueError("model_names must be a list.")

        if metrics is not None:
            if len(metrics) < 1:
                raise ValueError('The metrics list is empty.')
            else:
                self.metrics = metrics
        else:
            raise ValueError('Invalid metrics {}'.format(metrics))

        if constants.DATA_POINTS in metrics and n_models > 1:
            raise ValueError("Can not use multiple models with texts points visualization.")

        # Proceed with configuration
        self.__configure()

    def on_new_train_step(self, sample_id, data_buffer):
        """ Whenever the EvaluationVisualiser should
        be aware of some new texts, the caller will invoke this function,
        passing the new texts buffer.
        
        Parameters
        ----------
        sample_id: int
            The current sample id.

        data_buffer: EvaluationDataBuffer
            A buffer containing evaluation texts for a single training / visualization step.
            
        Raises
        ------
        ValueError: If an exception is raised during the draw operation.
         
        """

        try:
            current_time = time.time()
            self._clear_annotations()
            self._update_plots(sample_id, data_buffer)

            # To mitigate re-drawing overhead for fast models use frame counter (default = 5 frames).
            # To avoid slow refresh rate in slow models use a time limit (default = 1 sec).
            if (self._frame_cnt < 5) or (current_time - self._last_draw_timestamp > 1):
                plt.subplots_adjust(right=0.72)  # Adjust subplots to include metrics annotations
                self.fig.canvas.draw()
                plt.pause(1e-9)
                self._frame_cnt = 0
                self._last_draw_timestamp = current_time
            else:
                self._frame_cnt += 1
        except BaseException as exception:
            raise ValueError('Failed when trying to draw plot. Exception: {} | Type: {}'.
                             format(exception, type(exception).__name__))

    def __configure(self):
        """  This function will verify which subplots should be create. Initializing
        all relevant objects to keep track of the plotting points.

        """
        font_size_small = 8
        font_size_medium = 10
        font_size_large = 12

        plt.rc('font', size=font_size_small)  # controls default text sizes
        plt.rc('axes', titlesize=font_size_medium)  # font size of the axes title
        plt.rc('axes', labelsize=font_size_small)  # font size of the x and y y_values
        plt.rc('xtick', labelsize=font_size_small)  # font size of the tick y_values
        plt.rc('ytick', labelsize=font_size_small)  # font size of the tick y_values
        plt.rc('legend', fontsize=font_size_small)  # legend font size
        plt.rc('figure', titlesize=font_size_large)  # font size of the figure title

        warnings.filterwarnings("ignore", ".*GUI is implemented.*")
        warnings.filterwarnings("ignore", ".*left==right.*")
        warnings.filterwarnings("ignore", ".*Passing 1d.*")

        self._sample_ids = []

        plt.ion()
        self.fig = plt.figure(figsize=(9, 5))
        self.fig.suptitle(self.dataset_name)
        base = 11 + len(self.metrics) * 100  # 3-digit integer describing the position of the subplot.
        self.fig.canvas.set_window_title('kolibri')

        # Subplots handler
        for metric_id in self.metrics:
            data_ids = self._data_dict[metric_id]
            self._plot_trackers[metric_id] = PlotDataTracker(data_ids)
            plot_tracker = self._plot_trackers[metric_id]
            plot_tracker.sub_plot_obj = self.fig.add_subplot(base)
            base += 1
            if metric_id == constants.TRUE_VS_PREDICTED:
                handle = []
                plot_tracker.sub_plot_obj.set_prop_cycle(cycler('color', ['c', 'm', 'y', 'k']))
                for data_id in data_ids:
                    if data_id == constants.Y_TRUE:
                        # True texts
                        plot_tracker.data[data_id] = []
                        label = 'True value'
                        line_style = '--'
                        line_obj = plot_tracker.line_objs
                        if self.task_type == TaskType.CLASSIFICATION:
                            line_obj[data_id], = plot_tracker.sub_plot_obj.step(self._sample_ids,
                                                                                plot_tracker.data[data_id],
                                                                                label=label, linestyle=line_style)
                        else:
                            line_obj[data_id], = plot_tracker.sub_plot_obj.plot(self._sample_ids,
                                                                                plot_tracker.data[data_id],
                                                                                label=label, linestyle=line_style)
                        handle.append(line_obj[data_id])
                    else:
                        # Predicted texts
                        plot_tracker.data[data_id] = [[] for _ in range(self.n_models)]
                        plot_tracker.line_objs[data_id] = [None for _ in range(self.n_models)]
                        line_obj = plot_tracker.line_objs[data_id]
                        for i in range(self.n_models):
                            label = 'Predicted {}'.format(self.model_names[i])
                            line_style = '--'
                            if self.task_type == TaskType.CLASSIFICATION:
                                line_obj[i], = plot_tracker.sub_plot_obj.step(self._sample_ids,
                                                                              plot_tracker.data[data_id][i],
                                                                              label=label, linestyle=line_style)
                            else:
                                line_obj[i], = plot_tracker.sub_plot_obj.plot(self._sample_ids,
                                                                              plot_tracker.data[data_id][i],
                                                                              label=label, linestyle=line_style)
                            handle.append(line_obj[i])
                plot_tracker.sub_plot_obj.legend(handles=handle, loc=2, bbox_to_anchor=(1.01, 1.))
                plot_tracker.sub_plot_obj.set_title('True vs Predicted')
                plot_tracker.sub_plot_obj.set_ylabel('y')

            elif metric_id == constants.DATA_POINTS:
                plot_tracker.data['buffer_size'] = 100
                plot_tracker.data['X'] = Deque(plot_tracker.data['buffer_size'])
                plot_tracker.data['target_values'] = None
                plot_tracker.data['predictions'] = Deque(plot_tracker.data['buffer_size'])
                plot_tracker.data['clusters'] = []
                plot_tracker.data['clusters_initialized'] = False

            else:
                # Default case, 'mean' and 'current' performance
                handle = []
                sorted_data_ids = data_ids.copy()
                sorted_data_ids.sort()  # For better usage of the color cycle, start with 'current' texts
                for data_id in sorted_data_ids:
                    plot_tracker.data[data_id] = [[] for _ in range(self.n_models)]
                    plot_tracker.line_objs[data_id] = [None for _ in range(self.n_models)]
                    line_obj = plot_tracker.line_objs[data_id]
                    for i in range(self.n_models):
                        if data_id == constants.CURRENT:
                            label = '{}  (current, {} samples)'.format(self.model_names[i], self.n_wait)
                            line_style = '-'
                        else:
                            label = '{} (mean)'.format(self.model_names[i])
                            line_style = ':'
                        line_obj[i], = plot_tracker.sub_plot_obj.plot(self._sample_ids,
                                                                      plot_tracker.data[data_id][i],
                                                                      label=label,
                                                                      linestyle=line_style)
                        handle.append(line_obj[i])
                self._set_fig_legend(handle)

                if metric_id == constants.ACCURACY:
                    plot_tracker.sub_plot_obj.set_title('Accuracy')
                    plot_tracker.sub_plot_obj.set_ylabel('acc')
                elif metric_id == constants.HAMMING_SCORE:
                    plot_tracker.sub_plot_obj.set_title('Hamming score')
                    plot_tracker.sub_plot_obj.set_ylabel('hamming score')
                elif metric_id == constants.HAMMING_LOSS:
                    plot_tracker.sub_plot_obj.set_title('Hamming loss')
                    plot_tracker.sub_plot_obj.set_ylabel('hamming loss')
                elif metric_id == constants.EXACT_MATCH:
                    plot_tracker.sub_plot_obj.set_title('Exact Match')
                    plot_tracker.sub_plot_obj.set_ylabel('exact match')
                elif metric_id == constants.J_INDEX:
                    plot_tracker.sub_plot_obj.set_title('Jaccard Index')
                    plot_tracker.sub_plot_obj.set_ylabel('j-index')
                elif metric_id == constants.MSE:
                    plot_tracker.sub_plot_obj.set_title('Mean Squared Error')
                    plot_tracker.sub_plot_obj.set_ylabel('mse')
                elif metric_id == constants.MAE:
                    plot_tracker.sub_plot_obj.set_title('Mean Absolute Error')
                    plot_tracker.sub_plot_obj.set_ylabel('mae')
                elif metric_id == constants.AMSE:
                    plot_tracker.sub_plot_obj.set_title('Average Mean Squared Error')
                    plot_tracker.sub_plot_obj.set_ylabel('amse')
                elif metric_id == constants.AMAE:
                    plot_tracker.sub_plot_obj.set_title('Average Mean Absolute Error')
                    plot_tracker.sub_plot_obj.set_ylabel('amae')
                elif metric_id == constants.ARMSE:
                    plot_tracker.sub_plot_obj.set_title('Average Root Mean Squared Error')
                    plot_tracker.sub_plot_obj.set_ylabel('armse')
                elif metric_id == constants.DATA_POINTS:
                    plot_tracker.sub_plot_obj.set_title('')
                    plot_tracker.sub_plot_obj.set_xlabel('Feature x')
                    plot_tracker.sub_plot_obj.set_ylabel('Feature y')
                else:
                    plot_tracker.sub_plot_obj.set_title('Unknown metric')
                    plot_tracker.sub_plot_obj.set_ylabel('')

        if constants.DATA_POINTS not in self.metrics:
            plt.xlabel('Samples')

        self.fig.subplots_adjust(hspace=.5)
        self.fig.tight_layout(rect=[0, .04, 1, 0.98], pad=2.6, w_pad=0.5, h_pad=1.0)

    def _set_fig_legend(self, handles=None):
        if not self._is_legend_set:
            self.fig.legend(handles=handles, ncol=self.n_models, bbox_to_anchor=(0.02, 0.0), loc="lower left")
            self._is_legend_set = True

    def _update_plots(self, sample_id, data_buffer):
        self._sample_ids.append(sample_id)
        for metric_id, data_ids in data_buffer.data_dict.items():
            update_xy_limits = True
            y_min = 0.0
            y_max = 1.0
            pad = 0.1  # Default padding to set above and bellow plots
            plot_tracker = self._plot_trackers[metric_id]
            if metric_id == constants.TRUE_VS_PREDICTED:
                # Process true values
                data_id = constants.Y_TRUE
                plot_tracker.data[data_id].append(data_buffer.get_data(metric_id=metric_id, data_id=data_id))
                plot_tracker.line_objs[data_id].set_data(self._sample_ids, plot_tracker.data[data_id])
                # Process predicted values
                data_id = constants.Y_PRED
                data = data_buffer.get_data(metric_id=metric_id, data_id=data_id)
                for i in range(self.n_models):
                    plot_tracker.data[data_id][i].append(data[i])
                    plot_tracker.line_objs[data_id][i].set_data(self._sample_ids, plot_tracker.data[data_id][i])
                    y_min = min([plot_tracker.data[data_id][i][-1], plot_tracker.data[constants.Y_TRUE][-1], y_min])
                    y_max = max([plot_tracker.data[data_id][i][-1], plot_tracker.data[constants.Y_TRUE][-1], y_max])
            elif metric_id == constants.DATA_POINTS:
                update_xy_limits = False
                # Process features
                data_id = 'X'
                features_dict = data_buffer.get_data(metric_id=metric_id, data_id=data_id)
                feature_indices = list(features_dict.keys())
                feature_indices.sort()
                # Store tuple of feature values into the buffer, sorted by index
                plot_tracker.data[data_id].add_element([[features_dict[feature_indices[0]],
                                                         features_dict[feature_indices[1]]]])

                plot_tracker.sub_plot_obj.set_xlabel('Feature {}'.format(feature_indices[0]))
                plot_tracker.sub_plot_obj.set_ylabel('Feature {}'.format(feature_indices[1]))
                # TODO consider a fading/update strategy instead
                plot_tracker.sub_plot_obj.clear()
                X1 = plot_tracker.data[data_id].get_queue()[-1][0]
                X2 = plot_tracker.data[data_id].get_queue()[-1][1]

                # Process label values
                data_id = 'target_values'
                plot_tracker.data[data_id] = data_buffer.get_data(metric_id=metric_id, data_id=data_id)
                if not plot_tracker.data['clusters_initialized']:
                    for j in range(len(plot_tracker.data[data_id])):
                        plot_tracker.data['clusters'].append(Deque(plot_tracker.data['buffer_size']))

                # Process predictions
                data_id = 'predictions'
                plot_tracker.data[data_id].add_element([data_buffer.get_data(metric_id=metric_id, data_id=data_id)])

                for k, cluster in enumerate(plot_tracker.data['clusters']):
                    if plot_tracker.data[data_id].get_queue()[-1] == k:
                        plot_tracker.data['clusters'][k].add_element([(X1, X2)])
                        # TODO confirm buffer update inside the loop
                    if cluster.get_queue():
                        temp = cluster.get_queue()
                        plot_tracker.sub_plot_obj.scatter(*zip(*temp), label="Class {k}".format(k=k))
                plot_tracker.sub_plot_obj.legend(loc=2, bbox_to_anchor=(1.01, 1.))

            else:
                # Default case, 'mean' and 'current' performance
                for data_id in data_ids:
                    # Buffer texts
                    data = data_buffer.get_data(metric_id=metric_id, data_id=data_id)
                    for i in range(self.n_models):
                        plot_tracker.data[data_id][i].append(data[i])
                        plot_tracker.line_objs[data_id][i].set_data(self._sample_ids, plot_tracker.data[data_id][i])
                # Process texts
                for i in range(self.n_models):
                    # Update annotations
                    self._update_annotations(i, plot_tracker.sub_plot_obj, self.model_names[i],
                                             plot_tracker.data[constants.MEAN][i][-1],
                                             plot_tracker.data[constants.CURRENT][i][-1])
                    # Update plot limits
                    if metric_id in [constants.MSE, constants.MAE, constants.AMSE, constants.AMAE, constants.ARMSE]:
                        y_min = -1
                        y_max = max([plot_tracker.data[constants.MEAN][i][-1],
                                     plot_tracker.data[constants.CURRENT][i][-1], y_max])
                        pad = 0.5 * y_max  # Padding bellow and above thresholds
            if update_xy_limits:
                plot_tracker.sub_plot_obj.set_ylim((y_min - pad, y_max + pad))
                plot_tracker.sub_plot_obj.set_xlim(0, self._sample_ids[-1])

    def _clear_annotations(self):
        """ Clear annotations, so next frame is correctly rendered. """
        for i in range(len(self._text_annotations)):
            self._text_annotations[i].remove()
        self._text_annotations = []

    def _update_annotations(self, idx, subplot, model_name, mean_value, current_value):
        xy_pos_default = (1.02, .90)  # Default xy position for metric annotations
        shift_y = 10 * (idx + 1)  # y axis shift for plot annotations
        xy_pos = xy_pos_default
        if idx == 0:
            self._text_annotations.append(subplot.annotate('{: <12} | {: ^16} | {: ^16}'.
                                                           format('Model', 'Mean', 'Current'),
                                                           xy=xy_pos, xycoords='axes fraction'))
        self._text_annotations.append(subplot.annotate('{: <10.10s}'.format(model_name[:6]),
                                                       xy=xy_pos, xycoords='axes fraction',
                                                       xytext=(0, -shift_y), textcoords='offset points'))
        self._text_annotations.append(subplot.annotate('{: ^16.4f}  {: ^16.4f}'.format(mean_value, current_value),
                                                       xy=xy_pos, xycoords='axes fraction',
                                                       xytext=(50, -shift_y), textcoords='offset points'))

    @staticmethod
    def hold():
        plt.show(block=True)

    def get_info(self):
        pass


class PlotDataTracker(object):
    """ A class to track relevant texts for plots corresponding to selected metrics.
    Data buffers and line objects are accessible via the corresponding data_id.

    """

    def __init__(self, data_ids: list):
        self.data_ids = None
        self.data = {}
        self.line_objs = {}
        self.sub_plot_obj = None
        self._validate(data_ids)
        self._configure()

    def _validate(self, data_ids):
        if isinstance(data_ids, list):
            if len(data_ids) > 0:
                self.data_ids = data_ids
            else:
                raise ValueError('data_ids is empty')
        else:
            raise TypeError('data_ids must be a list, received: {}'.format(type(data_ids)))

    def _configure(self):
        for data_id in self.data_ids:
            self.data[data_id] = None
            self.line_objs[data_id] = None
