from numpy import log10, arange, floor, ceil
from pyqtgraph import PlotWidget, PlotItem


__all__ = ['NMRPlotWidget', 'logTickValues']


class NMRPlotWidget(PlotWidget):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        for orient in ['top', 'bottom', 'left', 'right']:
            # BAD HACK!!! but seems to work, see function for explanation
            self.plotItem.getAxis(orient).logTickValues = logTickValues


class NMRPlotItem(PlotItem):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        for orient in ['top', 'bottom', 'left', 'right']:
            self.plotItem.getAxis(orient).logTickValues = logTickValues


def logTickValues(minVal, maxVal, size, stdTicks):
    # TODO FIND A BETTER SOLUTION!!!
    # Sometimes minVal and maxVal are not log-scaled values and the loop from v1 to v2 is humongous,
    # The minor list then fills the RAM completely and freezes everything
    # Until there is a better solution, we overwrite this function for every AxesItem
    # and do not draw minor ticks at all if there are too many

    # start with the tick spacing given by tickValues().
    # Any level whose spacing is < 1 needs to be converted to log scale
    ticks = []
    for (spacing, t) in stdTicks:
        if spacing >= 1.0:
            ticks.append((spacing, t))

    if len(ticks) < 3:
        v1 = int(floor(minVal))
        v2 = int(ceil(maxVal))
        # major = list(range(v1+1, v2))
        minor = []

        if v2 - v1 < 400:
            for v in range(v1, v2):
                minor.extend(v + log10(arange(1, 10)))
            minor = [x for x in minor if minVal < x < maxVal]
            ticks.append((None, minor))
    return ticks