import numpy as np
from matplotlib import pyplot as plt
from spikewidgets.widgets.basewidget import BaseWidget


def plot_confusion_matrix(gt_comparison, count_text=True, unit_ticks=True,
                          ax=None, figure=None):
    """
    Plots sorting comparison confusion matrix.

    Parameters
    ----------
    gt_comparison: GroundTruthComparison
        The ground truth sorting comparison object
    count_text: bool
        If True counts are displayed as text
    unit_ticks: bool
        If True unit tick labels are displayed
    figure: matplotlib figure
        The figure to be used. If not given a figure is created
    ax: matplotlib axis
        The axis to be used. If not given an axis is created

    Returns
    -------
    W: ConfusionMatrixWidget
        The output widget
    """
    W = ConfusionMatrixWidget(
        gt_comparison=gt_comparison,
        count_text=count_text,
        unit_ticks=unit_ticks,
        figure=figure,
        ax=ax,
    )
    W.plot()
    return W


class ConfusionMatrixWidget(BaseWidget):
    def __init__(self, *, gt_comparison, count_text=True, unit_ticks=True,
                 figure=None, ax=None):
        BaseWidget.__init__(self, figure, ax)
        self._gtcomp = gt_comparison
        self._count_text = count_text
        self._unit_ticks = unit_ticks
        self.name = 'ConfusionMatrix'

    def plot(self):
        self._do_plot()

    def _do_plot(self):
        # a dataframe
        confusion_matrix = self._gtcomp.get_confusion_matrix()

        N1 = confusion_matrix.shape[0] - 1
        N2 = confusion_matrix.shape[1] - 1

        # Using matshow here just because it sets the ticks up nicely. imshow is faster.
        self.ax.matshow(confusion_matrix.values, cmap='Greens')

        if self._count_text:
            for (i, j), z in np.ndenumerate(confusion_matrix.values):
                if z != 0:
                    if z > np.max(confusion_matrix.values) / 2.:
                        self.ax.text(j, i, '{:d}'.format(z), ha='center', va='center', color='white')
                    else:
                        self.ax.text(j, i, '{:d}'.format(z), ha='center', va='center', color='black')

        self.ax.axhline(int(N1 - 1) + 0.5, color='black')
        self.ax.axvline(int(N2 - 1) + 0.5, color='black')

        # Major ticks
        self.ax.set_xticks(np.arange(0, N2 + 1))
        self.ax.set_yticks(np.arange(0, N1 + 1))
        self.ax.xaxis.tick_bottom()
        
        # Labels for major ticks
        if self._unit_ticks:
            self.ax.set_yticklabels(confusion_matrix.index, fontsize=12)
            self.ax.set_xticklabels(confusion_matrix.columns, fontsize=12)
        else:
            self.ax.set_xticklabels(np.append([''] * N2, 'FN'), fontsize=10)
            self.ax.set_yticklabels(np.append([''] * N1, 'FP'), fontsize=10)

        self.ax.set_xlabel(self._gtcomp.name_list[1], fontsize=20)
        self.ax.set_ylabel(self._gtcomp.name_list[0], fontsize=20)
        
        self.ax.set_xlim(-0.5, N2+0.5)
        self.ax.set_ylim(N1+0.5, -0.5, )
        
        

