from functools import wraps

from ..Qt import QtWidgets


def update_indexes(func):

    @wraps(func)
    def wrapped(self, *args, **kwargs):
        ret_val = func(self, *args, **kwargs)
        self.blockSignals(True)

        iterator = QtWidgets.QTreeWidgetItemIterator(self)
        i = j = 0
        while iterator.value():
            item = iterator.value()
            if item is not None:
                if item.parent() is None:
                    item.setText(1, 'g[{}]'.format(i))
                    i += 1
                    j = 0
                else:
                    item.setText(1, '.s[{}]'.format(j))
                    j += 1
                iterator += 1

        self.blockSignals(False)

        return ret_val

    return wrapped


def plot_update(func):

    @wraps(func)
    def wrapped(self, *args, **kwargs):
        ret_val = func(self, *args, **kwargs)
        m = self._data.mask
        _x = self._data.x
        _y = self._data.y

        self.plot_real.setData(x=_x[m], y=_y.real[m], name=self._data.name)
        if self.plot_imag is not None:
            self.plot_imag.setData(x=_x[m], y=_y.imag[m], name=self._data.name)

        if self.plot_error is not None:
            _y_err = self._data.y_err
            self.plot_error.setData(x=_x[m], y=_y.real[m], top=_y_err[m], bottom=_y_err[m])

        self.dataChanged.emit(self.id)

        return ret_val

    return wrapped