from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

def error_buckets(gold, pred, X=None):
    """Group items by error buckets
    
    Args:
        gold: a torch.Tensor of gold labels (ints)
        pred: a torch.Tensor of predictions (ints)
        X: an iterable of items
    Returns:
        buckets: A dict of items where buckets[i,j] is a list of items with
            predicted label i and true label j. If X is None, return indices
            instead.

    For a binary problem with (1=positive, 2=negative):
        buckets[1,1] = true positives
        buckets[1,2] = false positives
        buckets[2,1] = false negatives
        buckets[2,2] = true negatives
    """
    buckets = defaultdict(list)
    for i, (y, l) in enumerate(zip(gold.numpy(), pred.numpy())):
        buckets[y,l].append(X[i] if X is not None else i)
    return buckets


def confusion_matrix(gold, pred, null_pred=False, null_gold=False, 
    normalize=False, pretty=False):
    """A shortcut method for building a confusion matrix all at once.
    
    Args:
        gold: a torch.Tensor of gold labels (ints)
        pred: a torch.Tensor of predictions (ints)
        normalize: if True, divide counts by the total number of items
        pretty: if True, pretty-print the matrix
    """    
    conf = ConfusionMatrix(null_pred=null_pred, null_gold=null_gold)
    conf.add(gold, pred)
    mat = conf.compile()
    
    if normalize:
        mat = mat / len(gold)

    if pretty:
        conf.display()

    return mat

class ConfusionMatrix(object):
    """
    An iteratively built abstention-aware confusion matrix with pretty printing

    Assumed axes are true label on top, predictions on the side.
    """
    def __init__(self, null_pred=False, null_gold=False):
        """
        Args:
            null_pred: If True, show the row corresponding to null predictions
            null_gold: If True, show the col corresponding to null gold labels

        """
        self.counter = Counter()
        self.mat = None
        self.null_pred = null_pred
        self.null_gold = null_gold

    def __repr__(self):
        if self.mat is None:
            self.compile()
        return str(self.mat)

    def add(self, gold, pred):
        """
        Args:
            gold: a torch.Tensor of gold labels (ints)
            pred: a torch.Tensor of predictions (ints)
        """
        self.counter.update(zip(gold.numpy(), pred.numpy()))
    
    def compile(self, trim=True):
        k = max([max(tup) for tup in self.counter.keys()]) + 1  # include 0

        mat = np.zeros((k, k), dtype=int)
        for (p, y), v in self.counter.items():
            mat[p, y] = v
        
        if trim and not self.null_pred:
            mat = mat[1:, :]
        if trim and not self.null_gold:
            mat = mat[:, 1:]

        self.mat = mat
        return mat

    def display(self, counts=True, indent=0, spacing=2, decimals=3, 
        mark_diag=True):
        mat = self.compile(trim=False)
        m, n = mat.shape
        tab = ' ' * spacing
        margin = ' ' * indent

        # Print headers
        s = margin + ' ' * (5 + spacing)
        for j in range(n):
            if j == 0 and not self.null_gold:
                continue
            s += f" y={j} " + tab
        print(s)

        # Print data
        for i in range(m):
            # Skip null predictions row if necessary
            if i == 0 and not self.null_pred:
                continue
            s = margin + f" l={i} " + tab
            for j in range(n):
                # Skip null gold if necessary
                if j == 0 and not self.null_gold:
                    continue
                else:
                    if i == j and mark_diag and not counts:
                        s = s[:-1] + '*'
                    if counts:
                        s += f"{mat[i,j]:^5d}" + tab
                    else:
                        s += f"{mat[i,j]/sum(mat[i,1:]):>5.3f}" + tab
            print(s)

############################################################
# Plotting
############################################################
def plot_probabilities_histogram(Y_p, title=None):
    """Plot a histogram from a numpy array of probabilities"""
    plt.hist(Y_p, bins=20)
    plt.xlim((0, 1.025))
    plt.xlabel("Probability")
    plt.ylabel("# Predictions")
    if isinstance(title, str):
        plt.title(title)
    plt.show()

def plot_predictions_histogram(Y_ph, Y, title=None):
    """Plot a histogram comparing hard predictions vs true labels by class
    
    Args:
        Y_ph: An [N] or [N, 1] np.array of predicted hard labels
        Y: An [N] or [N, 1] np.array of gold labels 
    """
    labels = list(set(Y).union(set(Y_ph)))
    edges = [x - 0.5 for x in range(min(labels), max(labels) + 2)]

    plt.hist(
        [Y_ph, Y], 
        bins=edges,
        label=['Predicted', 'Gold'],
    )
    ax = plt.gca()
    ax.set_xticks(labels)
    plt.xlabel("Label")
    plt.ylabel("# Predictions")
    plt.legend(loc='upper right')
    if isinstance(title, str):
        plt.title(title)
    plt.show()
@@ -192,37 +192,36 @@
      gold: a
- torch.Tensor
+n np.ndarray
  of gold lab
@@ -242,37 +242,36 @@
      pred: a
- torch.Tensor
+n np.ndarray
  of predicti
@@ -790,38 +790,22 @@
 zip(gold
-.numpy(), pred.numpy()
+, pred
 )):%0A    
@@ -1079,37 +1079,36 @@
      gold: a
- torch.Tensor
+n np.ndarray
  of gold lab
@@ -1129,37 +1129,36 @@
      pred: a
- torch.Tensor
+n np.ndarray
  of predicti
@@ -2287,36 +2287,34 @@
     gold: a 
-torch.Tensor
+np.ndarray
  of gold lab
@@ -2348,20 +2348,18 @@
 : a 
-torch.Tensor
+np.ndarray
  of 
@@ -2429,30 +2429,14 @@
 gold
-.numpy(), pred.numpy()
+, pred
 ))%0A 
@@ -4226,16 +4226,112 @@
 bilities
+%0A    %0A    Args:%0A        Y_p: An %5BN%5D or %5BN, 1%5D np.ndarray of probabilities (floats in %5B0,1%5D)%0A    
 %22%22%22%0A    
@@ -4681,32 +4681,34 @@
 N%5D or %5BN, 1%5D np.
+nd
 array of predict
@@ -4753,16 +4753,18 @@
 , 1%5D np.
+nd
 array of

Update analysis tools to expect numpy.ndarray instead of torch.Tensor