# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/visualization/visualization.basic_plots.ipynb (unless otherwise specified).

__all__ = ['Visual', 'BasicVisual', 'CountPlot', 'LongTailPlot', 'LongTailPlotv2', 'MatrixPlot', 'ConfusionMatrix',
           'EmbeddingPlot', 'ClusterPlot', 'TimelinePlot', 'ImageGridPlot']

# Cell
import matplotlib.pyplot as plt
import seaborn as sns

import pandas as pd
import numpy as np

# Cell
class Visual:
    def __init__(self):
        pass

# Cell
class BasicVisual(Visual):
    def __init__(self, figsize=(10,10), show=True, save=False, pallete='viridis',
                 style='whitegrid', rot=45, label_x='label_x', label_y='label_y',
                 title='Basic Plot', nrows=1, ncols=1):
        super(BasicVisual, self).__init__()
        self.nrows = nrows
        self.ncols = ncols
        self.figsize = figsize
        self.show = show
        self.pallete = pallete
        self.style = style
        self.rot = rot
        self.label_x = label_x
        self.label_y = label_y
        self.title = title

        self._plot()

    def plot(self):
        raise NotImplementedError

    def _plot(self):
        self.fig, self.ax = plt.subplots(nrows=self.nrows, ncols=self.ncols,
                                         figsize=self.figsize)
        sns.set_palette(self.pallete)
        sns.set_style(self.style)
        plt.xticks(rotation=self.rot)
        plt.xlabel(self.label_x)
        plt.ylabel(self.label_y)
        plt.title(self.title)

        self.plot()

        plt.show()


# Cell
class CountPlot(BasicVisual):
    def __init__(self, data_x, data_y, add_pct=True, **kwargs):
        self.data_x = data_x
        self.data_y = data_y
        self.add_pct = add_pct
        super(CountPlot, self).__init__(**kwargs)

    def plot(self):
        sns.barplot(x=self.data_x, y=self.data_y, ax=self.ax)
        if self.add_pct:
            self.pct_patch()

    def pct_patch(self):
        tot_pct = 0
        total = np.sum(self.data_y)
        for p in self.ax.patches:
            height = p.get_height()
            self.ax.text(p.get_x()+p.get_width()/2.,
                    height + 1,
                    '{:1.2f}%'.format(height*100/total),
                    ha="center")
            tot_pct+= height*100/total

# Cell
class LongTailPlot(BasicVisual):
    def __init__(self, item_ids, user_ids, **kwargs):
        kwargs['nrows'] = 1
        kwargs['ncols'] = 2
        self.df = pd.DataFrame.from_dict({
            'user_id':user_ids,
            'item_id':item_ids,
        })
        super(LongTailPlot, self).__init__(**kwargs)

    def plot(self):
        item_rate_count = self.df.groupby("item_id")["user_id"].nunique().sort_values(ascending=False)
        self.ax[0].bar(x=item_rate_count.index, height=item_rate_count.values, width=1.0, align="edge")
        self.ax[0].set_xticks([])
        self.ax[0].set(title="long tail of rating frequency",
                    xlabel="item ordered by decreasing frequency",
                    ylabel="#ratings")

        count = item_rate_count.value_counts()
        sns.scatterplot(x=np.log(count.index), y=np.log(count.values), ax=self.ax[1])
        self.ax[1].set(title="log-log plot", xlabel="#ratings (log scale)", ylabel="#items (log scale)")

# Cell
class LongTailPlotv2(BasicVisual):
    def __init__(self, item_ids, percentage=None, x_labels=True, **kwargs):
        """
        Plots the long tail for a user-item interaction dataset.
        ----------
        item_ids:
            list of item ids
        percentage: float, default=None
            percent of volume to consider as the head (percent as a decimal)
            (if default=None no line will be plotted)
        x_tables bool, default=True
            if True, plot x-axis tick labels
            if False, no x-axis tick labels will be plotted.
        Returns:
        -------
            A long tail plot
        """
        self.df = pd.DataFrame.from_dict({
            'item_id':item_ids,
        })
        self.percentage = percentage
        self.x_labels = x_labels
        super(LongTailPlotv2, self).__init__(**kwargs)

    def plot(self):
        from matplotlib.lines import Line2D

        #calculate cumulative volumes
        volume_df = self.df['item_id'].value_counts().reset_index()
        volume_df.columns = ['item_id', "volume"]
        volume_df['item_id'] = volume_df['item_id'].astype(str)
        volume_df['cumulative_volume'] = volume_df['volume'].cumsum()
        volume_df['percent_of_total_volume'] = volume_df['cumulative_volume']/volume_df['volume'].sum()

        #line plot of cumulative volume
        x = range(0,len(volume_df))
        sns.lineplot(x=x, y=volume_df["volume"].values, color="black", ax=self.ax)
        plt.xticks(x)

        #set labels
        if self.percentage != None:
            #plot vertical line at the tail location
            head = volume_df[volume_df["percent_of_total_volume"] <= self.percentage]
            tail = volume_df[volume_df["percent_of_total_volume"] > self.percentage]
            items_in_head = len(head)
            items_in_tail = len(tail)
            plt.axvline(x=items_in_head, color="red",  linestyle='--')

            # fill area under plot
            head = head.append(tail.head(1))
            x1 = head.index.values
            y1 = head['volume']
            x2 = tail.index.values
            y2 = tail['volume']
            self.ax.fill_between(x1, y1, color="blue", alpha=0.2)
            self.ax.fill_between(x2, y2,  color="orange", alpha=0.2)

            #create legend
            legend_elements = [Line2D([0], [0], marker='o', color='w', label=str(items_in_head)+': items in the head', markerfacecolor='blue', markersize=5),
                            Line2D([0], [0], marker='o', color='w', label=str(items_in_tail)+': items in the tail', markerfacecolor='orange', markersize=5)]
            self.ax.legend(handles=legend_elements, loc=1)

        else:
            x1 = volume_df['item_id']
            y1 = volume_df['volume']
            self.ax.fill_between(x1, y1, color="blue", alpha=0.3)
        if self.x_labels == False:
            plt.xticks([], [])
            self.ax.set(xticklabels=[])
        else:
            self.ax.set_xticklabels(labels = volume_df['item_id'], rotation = 45, ha="right")

# Cell
class MatrixPlot(BasicVisual):
    def __init__(self, rating_matrix, interpolation=None, vmax=1, vmin=0, **kwargs):
        self.rating_matrix = rating_matrix
        self.interpolation=interpolation
        self.vmax=vmax
        self.vmin=vmin
        super(MatrixPlot, self).__init__(**kwargs)

    def plot(self):
        plt.imshow(self.rating_matrix, interpolation=self.interpolation,
                   vmax=self.vmax, vmin=self.vmin)

# Cell
class ConfusionMatrix(BasicVisual):
    def __init__(self, y, yhat, **kwargs):
        self.y = y
        self.yhat = yhat
        super(ConfusionMatrix, self).__init__(**kwargs)

    def plot(self):
        from itertools import product
        from sklearn.metrics import confusion_matrix

        cm = confusion_matrix(self.y, self.yhat, labels=[1,0])
        cm = np.round(cm.astype('float') / cm.sum(axis=1)[:, np.newaxis],4)*100

        fmt = ".2f"
        _ = cm.max() / 2. # TODO: Unused argument
        descriptions = np.array([["True Positive", "False Negative"], ["False Positive", "True Negatives"]])
        colors = np.array([["green", "red"], ["red", "green"]])
        plt.imshow([[0,0],[0,0]], interpolation='nearest', cmap=plt.cm.Greys)
        for i, j in product(range(cm.shape[0]), range(cm.shape[1])):
                plt.text(j, i, format(cm[i, j], fmt)+'%\n' + descriptions[i, j],
                        horizontalalignment="center",
                        color=colors[i,j])
        plt.axhline(y=0.5, xmin=0, xmax=1, color="black", linewidth=0.75)
        plt.axvline(x=0.5, ymin=0, ymax=1, color="black", linewidth=0.75)
        plt.ylabel('True')
        plt.xlabel('Predicted')
        plt.title("Confusion Matrix")
        plt.xticks([0,1], [1,0], rotation=45)
        plt.yticks([0,1], [1,0])

# Cell
class EmbeddingPlot(BasicVisual):
    def __init__(self, vectors, labels=None, **kwargs):
        self.vectors = vectors
        self.labels = labels
        kwargs['figsize'] = (12, len(vectors))
        super(EmbeddingPlot, self).__init__(**kwargs)

    def plot(self):
        sns.heatmap(self.vectors, vmax=2, vmin=-2, ax=self.ax, cmap=self.pallete)
        if self.labels:
            self.ax.set_yticklabels(self.labels, rotation=self.rot)
            self.ax.tick_params(axis='both', which='major', labelsize=30)
        plt.tick_params(axis='x',          # changes apply to the x-axis
                        which='both',      # both major and minor ticks are affected
                        bottom=False,      # ticks along the bottom edge are off
                        top=False,         # ticks along the top edge are off
                        labelbottom=False) # labels along the bottom edge are off
        # From https://github.com/mwaskom/seaborn/issues/1773
        # fix for mpl bug that cuts off top/bottom of seaborn viz
        b, t = plt.ylim() # discover the values for bottom and top
        b += 0.5 # Add 0.5 to the bottom
        t -= 0.5 # Subtract 0.5 from the top
        plt.ylim(b, t) # update the ylim(bottom, top) values

# Cell
class ClusterPlot(BasicVisual):
    def __init__(self, labels, embedding_clusters, word_clusters, **kwargs):
        self.labels = labels
        self.embedding_clusters = embedding_clusters
        self.word_clusters = word_clusters
        super(ClusterPlot, self).__init__(**kwargs)

    def plot(self):
        for label, embeddings, words in zip(self.labels, self.embedding_clusters, self.word_clusters):
            x = embeddings[:,0]
            y = embeddings[:,1]
            plt.scatter(x, y, label=label)
            for i, word in enumerate(words):
                plt.annotate(word, alpha=0.5, xy=(x[i], y[i]), xytext=(5, 2),
                            textcoords='offset points', ha='right', va='bottom', size=8)
        plt.legend(loc=4)

# Cell
class TimelinePlot(BasicVisual):
    def __init__(self, data, xticks_type='Month', xticks_interval=3, **kwargs):
        self.combined_data = data
        self.xticks_type = xticks_type
        self.xticks_interval = xticks_interval
        super(TimelinePlot, self).__init__(**kwargs)

    def plot(self):
        from datetime import datetime
        import matplotlib.dates as mdates
        import matplotlib.patches as mpatches

        # Group data by date
        if self.xticks_type=='Day':
            self.combined_data['date'] = self.combined_data['date'].dt.strftime('%Y-%m-%d')
            self.combined_data['date'] = self.combined_data['date'].apply(lambda x: datetime.strptime(x, "%Y-%m-%d"))
        else:
            self.combined_data['date'] = self.combined_data['date'].dt.strftime('%B %Y')
            self.combined_data['date'] = self.combined_data['date'].apply(lambda x: datetime.strptime(x, "%B %Y"))

        self.combined_data = self.combined_data.groupby(['type', 'date', 'color']).size().rename('count').reset_index().sort_values('date')
        self.combined_data['name'] = self.combined_data['count'].map(str)+ ' '+self.combined_data['type']
        names = self.combined_data['name'].tolist()
        colors = self.combined_data['color'].tolist()
        dates = self.combined_data['date'].tolist()

        # Plot
        levels = np.array([-9, 9, -7, 7, -5, 5, -3, 3, -1, 1])

        # Create the base line
        start = min(dates)
        stop = max(dates)
        self.ax.plot((start, stop), (0, 0), 'k', alpha=.5)

        # Create annotations
        for ii, (iname, idate, icol) in enumerate(zip(names, dates, colors)):
            level = levels[ii % len(levels)]
            vert = 'top' if level < 0 else 'bottom'
            self.ax.scatter(idate, 0, s=100, facecolor=icol, edgecolor='k', zorder=9999)
            self.ax.plot((idate, idate), (0, level), c=icol, alpha=1.0, lw=2)
            self.ax.text(idate, level, iname, horizontalalignment='center', verticalalignment=vert, fontsize=12, backgroundcolor=icol)
        # Set the xticks formatting
        if self.xticks_type=='Month':
            self.ax.get_xaxis().set_major_locator(mdates.MonthLocator(interval=self.xticks_interval))
            self.ax.get_xaxis().set_major_formatter(mdates.DateFormatter("%B %Y"))
        elif self.xticks_type=='Day':
            self.ax.get_xaxis().set_major_locator(mdates.DayLocator(interval=self.xticks_interval))
            self.ax.get_xaxis().set_major_formatter(mdates.DateFormatter("%Y-%m-%d"))
        elif self.xticks_type=='Year':
            self.ax.get_xaxis().set_major_locator(mdates.YearLocator())
            self.ax.get_xaxis().set_major_formatter(mdates.DateFormatter("%Y"))
        self.fig.autofmt_xdate()
        #Legend
        legend = []
        for index, row in self.combined_data[['type', 'color']].drop_duplicates().iterrows():
            legend += [mpatches.Patch(color=row['color'], label=row['type'])]
        plt.legend(handles=legend, loc='center left', bbox_to_anchor=(1, 0.5))

        # Remove components for a cleaner look
        plt.setp((self.ax.get_yticklabels() + self.ax.get_yticklines() + list(self.ax.spines.values())), visible=False)

# Cell
class ImageGridPlot(BasicVisual):
    def __init__(self, paths, labels, columns=5, **kwargs):
        self.paths = paths
        self.labels = labels
        self.columns = columns
        kwargs['figsize'] = (15,int(0.625 * len(paths)))
        # kwargs['nrows'] = len(paths) / columns + 1
        # kwargs['ncols'] = (15,int(0.625 * len(paths)))
        super(ImageGridPlot, self).__init__(**kwargs)

    def plot(self):
        from skimage import io

        for i, path in enumerate(self.paths):
            plt.subplot(len(self.paths) / self.columns + 1, self.columns, i + 1)
            image = io.imread(path)
            plt.imshow(image)
            plt.xticks(color = 'w', fontsize = 0.1)
            plt.yticks(color = 'w', fontsize = 0.1)
            plt.xlabel(self.labels[i], fontsize = 12)
            plt.tight_layout(h_pad=0.4, w_pad=0)
            plt.subplots_adjust(wspace=None, hspace=None)