import numpy as np
import pyqtgraph as pg

class IndexedAxis(pg.AxisItem):
    """
    Axis where pixels are mapped using a numpy array
    """
    def __init__(self, orientation, mapping=None, **kw):
        super(IndexedAxis, self).__init__(orientation, **kw)
        self.mapping = mapping
        self.fmt = "{:.02f}"

    def tickStrings(self, values, scale, spacing):
        if self.mapping is None:
            return super(IndexedAxis, self).tickStrings(values, scale, spacing)
        # count values smaller than 0
        labels = []
        idx = np.array(values, dtype=np.int)-1
        left_pad = np.count_nonzero(idx < 0)
        right_pad = np.count_nonzero(idx >= self.mapping.size)
        idx = np.compress(np.logical_and(idx>=0, idx< self.mapping.size),
                          idx)
        labels.extend([""]*left_pad)
        labels.extend([self.fmt.format(v) for v in self.mapping[idx]])
        labels.extend([""]*right_pad)
        return labels