import time


def annotate_heatmap(im,  # the AxesImage to be labeled
                     # a 2D numpy array of shape (N, M) containing the standard deviations to be used to annotate
                     # the heatmap
                     std,
                     # a 2D python array containing Boolean values indicating whether each single cell contains
                     # data or not
                     mask,
                     # the format of the annotations inside the heatmap.  This should be either 'time' or 'speed'
                     ann_format='time',
                     # a pair of colors.  The first is used for values below a threshold, the second for those above
                     textcolors=("black", "white"),
                     # value in data units according to which the colors from textcolors are applied.
                     # If None (the default) uses the middle of the colormap as separation
                     threshold=None,
                     # all other arguments are forwarded to each call to `text` used to create the text labels
                     **textkw):
    """ A function to annotate a heatmap.

    Args:
        im: The AxesImage to be labeled
        std: A 2D numpy array of shape (N, M) containing the standard deviations to be used to annotate the heatmap
        mask: A 2D python array containing Boolean values indicating whether each single cell contains data or not
        ann_format: The format of the annotations inside the heatmap.  This should be either 'time' or 'speed'
                    (default: 'time')
        textcolors: A pair of colors.  The first is used for values below a threshold, the second for those above
                    (default: ("black", "white"))
        threshold: Value in data units according to which the colors from textcolors are applied.
                   If None (the default) uses the middle of the colormap as separation (default: None)
        **textkw: All other arguments are forwarded to each call to `text` used to create the text labels
    """

    # get data from 'im' AxesImage
    data = im.get_array()

    # normalize the threshold to the images color range
    if threshold is not None:
        threshold = im.norm(threshold)
    else:  # if threshold was not provided set threshold using the middle of the colormap as separation
        threshold = im.norm(data.max()) / 2.

    # Set default alignment to center, but allow it to be overwritten by textkw.
    kw = dict(horizontalalignment="center",
              verticalalignment="center")
    kw.update(textkw)

    # if the annotation is not of type 'time' nor 'speed', raise error
    if ann_format != 'time' and ann_format != 'speed':
        raise ValueError('Unknown format type.')

    # initialize texts list to be empty
    texts = []

    # loop over the data
    for i in range(data.shape[0]):
        for j in range(data.shape[1]):
            if mask[i][j]:  # skip data point if its mask value is true
                continue

            # update color depending on the current data point
            kw.update(color=textcolors[int(im.norm(data[i, j]) > threshold)])

            if ann_format == 'time':  # if the annotation is of type 'time'
                # create text of type 'time' for the current data point
                text = im.axes.text(j, i, '{}\n+/-{:4.0f}s'
                                    .format(time.strftime("%H:%M:%S",
                                                          time.gmtime(float(data[i, j]))), std[i, j]), **kw)
            else:  # if the annotation is of type 'speed'
                # create text of type 'speed' for the current data point
                text = im.axes.text(j, i, '{:6.3f}it/s\n+/-\n{:6.3f}it/s'.format(data[i, j],
                                                                                 std[i, j]), **kw)

            # append generated text to texts list
            texts.append(text)

    return texts