##################################################################################
#  SENSIML CONFIDENTIAL                                                          #
#                                                                                #
#  Copyright (c) 2017-18  SensiML Corporation.                                   #
#                                                                                #
#  The source code contained or  described  herein and all documents related     #
#  to the  source  code ("Material")  are  owned by SensiML Corporation or its   #
#  suppliers or licensors. Title to the Material remains with SensiML Corpora-   #
#  tion  or  its  suppliers  and  licensors. The Material may contain trade      #
#  secrets and proprietary and confidential information of SensiML Corporation   #
#  and its suppliers and licensors, and is protected by worldwide copyright      #
#  and trade secret laws and treaty provisions. No part of the Material may      #
#  be used,  copied,  reproduced,  modified,  published,  uploaded,  posted,     #
#  transmitted, distributed,  or disclosed in any way without SensiML's prior    #
#  express written permission.                                                   #
#                                                                                #
#  No license under any patent, copyright,trade secret or other intellectual     #
#  property  right  is  granted  to  or  conferred upon you by disclosure or     #
#  delivery of the Materials, either expressly, by implication,  inducement,     #
#  estoppel or otherwise.Any license under such intellectual property rights     #
#  must be express and approved by SensiML in writing.                           #
#                                                                                #
#  Unless otherwise agreed by SensiML in writing, you may not remove or alter    #
#  this notice or any other notice embedded in Materials by SensiML or SensiML's #
#  suppliers or licensors in any way.                                            #
#                                                                                #
##################################################################################


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import NullFormatter
import matplotlib.patches as patches
from copy import deepcopy
from seaborn import boxplot, violinplot, set_style
from pandas import DataFrame, concat
import io


class Visualize(object):

    def __init__(self, feature_vector=None, model=None, label=None):
        """
                model: model result or knowledpack
                feature_vector (Dataframe): feature set that is results of generator_set or feature_selector
                label (String), name of class
        """
        self._feature_vector = None
        self._knowledgepack = None
        self._label = None

        self.knowledgepack = model
        self.label = label
        if self.label is None:
            print(
                "Label column is not specified for this knowledpack, specify as label=XXX")
        self.feature_vector = feature_vector

    @property
    def knowledgepack(self):
        """Auto generated unique identifier for the project object"""
        return self._knowledgepack

    @knowledgepack.setter
    def knowledgepack(self, value):
        if value is None:
            self._knowledgepack = None
        elif hasattr(value, 'knowledgepack'):
            self._knowledgepack = value.knowledgepack
        else:
            self._knowledgepack = value

    @property
    def label(self):
        """Auto generated unique identifier for the project object"""
        return self._label

    @label.setter
    def label(self, value):
        if value is None and self.knowledgepack:
            self._label = self.knowledgepack.query_summary.get(
                'label_column', None)
        else:
            self._label = value

    @property
    def feature_vector(self):
        """Auto generated unique identifier for the project object"""
        return self._feature_vector

    @feature_vector.setter
    def feature_vector(self, value):
        self._feature_vector = value[[
            i for i in sorted(value.columns.tolist()) if ('gen_' == i[:4])]+[self.label]]

    def plot_features(self, features_list=None,
                      plot_style='boxplot', figsize=(20, 5),
                      features_per_row=4):
        """
        Display the distributions of feature values across classes.

        Args:
            features (Dataframe): Feature set that is results of generator_set or feature_selector
            label (String): Name of class label
            features_list (List):  List of features to plot, if None plots all (default: None)
            plot_style (String):  'violinplot' / 'boxplot' to determine style of plot
            figsize (tuple): Size of figure to make
            features_per_row (int): Number of features to show in each row of the plot
        """

        # if features_list is not given use all featues
        if not(features_list):
            features_list = [
                i for i in self.feature_vector.columns.tolist() if ('gen_' == i[:4])]

        # reframe the feature dataframe so it can be used to input for violinplot or boxplot
        num_plots = len(features_list)//(features_per_row)
        if len(features_list) % features_per_row != 0:
            num_plots += 1
        figsize = (figsize[0], figsize[1]*num_plots)
        f, axarr = plt.subplots(num_plots, figsize=figsize)
        set_style('ticks')
        if not isinstance(axarr, np.ndarray):
            axarr = [axarr]

        for index in range(num_plots):
            df_list = [reframe(self.feature_vector, feature, self.label) for feature in
                       features_list[index*features_per_row:(index+1)*features_per_row]]

            if plot_style == 'violinplot':
                violinplot(x="Features", y="Sensor Values", hue=self.label,
                           data=concat(df_list), palette="muted", ax=axarr[index])
            elif plot_style == 'boxplot':
                boxplot(x="Features", y="Sensor Values", hue=self.label,
                        data=concat(df_list), palette="muted", ax=axarr[index])
            else:
                print('Error: plot_style is not in the library')
                return

            axarr[index].set_xticklabels(
                features_list[index *
                              features_per_row:(index+1)*features_per_row],
                fontsize=16)
            axarr[index].set_ylabel('Feature Values', fontsize=16)
            axarr[index].set_xlabel('')
            axarr[index].legend()
            axarr[index].set_ylim((0, 255))
            axarr[index].legend().set_visible(False)

        axarr[-1].set_xlabel('Features', fontsize=16)
        axarr[0].legend(loc='center left',
                        bbox_to_anchor=(1, 0.85), fontsize=16)
        axarr[0].set_title("Feature Vector Distribution", fontsize=16)

        plt.show()

    def neuron_feature_map(self,
                           feature1,
                           features_list,
                           label=None,
                           norm=None,
                           hist_figure=True,
                           neuron_figure=True,
                           vector_labels=['all'],
                           vector_alpha=0.3,
                           neuron_alpha=0.2,
                           color_list=['m', 'k', 'b', 'lime', 'r', 'c', 'y', 'brown', 'tan', 'g', 'pink', 'gray']):
        """
            This function display neuron map, feature map and feature distributions.

            Args:
                hist_figure = Bool, display historgram of vector for given features.
                neuron_figure = True,
                pivot_feature = String, it is y-axis feature that used to display neuron map, feature map and feature distributions
                versus_features = List<String>, x-axis features that used to display neuron map, feature map and feature distributions
                vector_labels (List): which features will be displayed  [ f1, f2, f3, ..., 'all', [] ], ['all'] display all features, [] does not display any features
                vector_alpha = 0.3, level of vector transparency
                neuron_alpha = 0.1, level of neuron transparency
                norm(int): takes from knowledpack, to overide 0=L1, 1=LSUP

            Returns: None

        """

        self.hist_figure = hist_figure
        self.pivot_feature = feature1
        self.vector_labels = vector_labels
        self.vector_alpha = vector_alpha
        self.neuron_alpha = neuron_alpha
        self.axis = 10  # axis
        self.color_list = color_list

        if isinstance(features_list, str):
            features_list = [features_list]

        self.columns_index = [None, None]
        self.columns_index[1] = self.feature_vector.columns.tolist().index(
            self.pivot_feature)

        for feature in features_list:
            self.columns_index[0] = self.feature_vector.columns.tolist().index(
                feature)

            if self.vector_labels == ['all']:
                self.vector_labels = np.unique(
                    self.feature_vector[self.label]).tolist()

            self._set_initial_figure()

            if neuron_figure and self.knowledgepack:
                if norm is None:
                    norm = self.knowledgepack.device_configuration.get(
                        'distance_mode', 0)
                else:
                    norm = norm
                self._plot_neurons(
                    self.knowledgepack.neuron_array, feature, norm)

            if (vector_labels != ['None']) | (vector_labels == []):
                self._plot_feature_vector(feature)

            self._set_final_figure_settings(feature)
            plt.show()

    def _plot_neurons(self, neurons, versus_feature, norm):
        num = 0
        for neuron in neurons:
            x_val = neuron['Vector'][self.columns_index[0]]
            y_val = neuron['Vector'][self.columns_index[1]]
            aif = neuron['AIF']

            if neuron['Category'] < 30000:  # if it is not degenerated neuron
                cat_label = self.knowledgepack.class_map[str(
                    neuron['Category'])]
                clr = self.color_list[self.label_indx_list.index(cat_label)]
                if norm == 0:
                    dim = len(neuron['Vector'])
                    new_aif = aif * np.sqrt(2)/dim
                    self.neuron_scatter.add_patch(patches.Rectangle((x_val-(aif*2/dim), y_val),
                                                                    2*new_aif,
                                                                    2*new_aif,
                                                                    color=clr,
                                                                    alpha=self.neuron_alpha,
                                                                    angle=-45.0))

                else:
                    self.neuron_scatter.add_patch(patches.Rectangle((x_val-aif, y_val-aif),
                                                                    2*aif,
                                                                    2*aif,
                                                                    color=clr,
                                                                    alpha=self.neuron_alpha,
                                                                    angle=0))

                self.neuron_scatter.scatter(
                    x_val, y_val, c=clr, marker='*', s=96)

    def _plot_hist(self, Xs, Ys, clr):
        binwidth = 5
        bins = np.arange(0, 256, binwidth)
        self.axHistx.hist(Xs, bins=bins, color=clr, alpha=self.vector_alpha)
        self.axHistx.set_xlim(self.neuron_scatter.get_xlim())

        self.axHisty.hist(Ys, bins=bins, color=clr,
                          alpha=self.vector_alpha, orientation='horizontal')
        self.axHisty.set_ylim(self.neuron_scatter.get_ylim())

    def _plot_feature_vector(self, versus_feature):
        for label_indx in self.vector_labels:
            temp_label_data = self.feature_vector[self.feature_vector[self.label] == label_indx].index.tolist(
            )
            temp = self.feature_vector.loc[temp_label_data[0]].tolist()
            cat_label = temp[-1]
            clr = self.color_list[self.label_indx_list.index(cat_label)]
            Xs = self.feature_vector.loc[temp_label_data,
                                         versus_feature].tolist()
            Ys = self.feature_vector.loc[temp_label_data,
                                         self.pivot_feature].tolist()

            self.neuron_scatter.scatter(Xs, Ys, c=clr, alpha=self.vector_alpha)

            if self.hist_figure:
                self._plot_hist(Xs, Ys, clr)

    def _set_initial_figure(self):
        plt.figure(1, figsize=(self.axis, self.axis))

        left, width = 0.1, 0.65
        bottom, height = 0.1, 0.65
        rect_scatter = [left, bottom, width, height]
        self.neuron_scatter = plt.axes(rect_scatter, aspect='equal')

        if self.hist_figure:
            # definitions for the axes
            bottom_h = left_h = left + width + 0.02
            rect_histx = [left, bottom_h, width, 0.2]
            rect_histy = [left_h, bottom, 0.2, height]
            nullfmt = NullFormatter()
            self.axHistx = plt.axes(rect_histx)
            self.axHistx.xaxis.set_major_formatter(nullfmt)
            self.axHisty = plt.axes(rect_histy)
            self.axHisty.yaxis.set_major_formatter(nullfmt)

        if self.knowledgepack:
            self.label_indx_list = np.unique([self.knowledgepack.class_map[str(i)] for i in
                                              np.sort(list(map(int, list(self.knowledgepack.class_map.keys()))))] +
                                             np.unique(self.feature_vector[self.label]).tolist()).tolist()
        else:
            self.label_indx_list = list(
                np.unique(self.feature_vector[self.label]))

        # this part creates dummy vectors with the unique number of labels for legent.
        for indx in range(len(self.label_indx_list)):
            self.neuron_scatter.scatter(
                0, 0, c=self.color_list[indx], label=self.label_indx_list[indx])

        self.neuron_scatter.set_xlim([0, 256])
        self.neuron_scatter.set_ylim([0, 256])

    def _set_final_figure_settings(self, versus_feature):

        self.neuron_scatter.set_xlabel(versus_feature, fontsize=12)
        self.neuron_scatter.set_ylabel(self.pivot_feature, fontsize=12)

        self.neuron_scatter.set_xticks(range(0, 256, 10))
        self.neuron_scatter.set_xticklabels(
            range(0, 256, 10), rotation=0, fontsize=8)

        self.neuron_scatter.set_yticks(range(0, 256, 10))
        self.neuron_scatter.set_yticklabels(
            range(0, 256, 10), rotation=0, fontsize=8)

        if self.hist_figure:
            self.axHistx.set_xticks(range(0, 256, 10))
            self.axHistx.yaxis.set_tick_params(labelsize=8)
            self.axHistx.set_title(
                self.pivot_feature + ' vs ' + versus_feature)

            self.axHisty.set_yticks(range(0, 256, 10))
            self.axHisty.xaxis.set_tick_params(labelsize=8)
            self.neuron_scatter.legend(
                loc='center left', bbox_to_anchor=(1.02, 1.26), fontsize=12)

        else:
            self.neuron_scatter.legend(
                loc='center left', bbox_to_anchor=(1, 0.945), fontsize=12)
            self.neuron_scatter.set_title(
                self.pivot_feature + ' vs ' + versus_feature)


def reframe(features, feature, label):
    # This funciton reframes the original features dataframe
    temp_list = DataFrame(columns=[label, 'Features', 'Sensor Values'])
    temp_list[label] = features[label].tolist()
    temp_list['Features'] = [feature] * len(features)
    temp_list['Sensor Values'] = features[feature].tolist()
    return temp_list
