from __future__ import annotations

from contextlib import contextmanager
from numpy import linspace
from scipy.interpolate import interp1d

from ..Qt import QtGui, QtWidgets, QtCore


@contextmanager
def busy_cursor():
    try:
        cursor = QtGui.QCursor(QtCore.Qt.CursorShape.ForbiddenCursor)
        QtWidgets.QApplication.setOverrideCursor(cursor)
        yield

    finally:
        QtWidgets.QApplication.restoreOverrideCursor()


class RdBuCMap:
    # taken from Excel sheet from colorbrewer.org
    _rdbu = [
        (103, 0, 31),
        (178, 24, 43),
        (214, 96, 77),
        (244, 165, 130),
        (253, 219, 199),
        (247, 247, 247),
        (209, 229, 240),
        (146, 197, 222),
        (67, 147, 195),
        (33, 102, 172),
        (5, 48, 97)
    ]

    def __init__(self, vmin=-1., vmax=1.):
        self.min = vmin
        self.max = vmax

        self.spline = [interp1d(linspace(self.max, self.min, num=11), [rgb[i] for rgb in RdBuCMap._rdbu])
                       for i in range(3)]

    def color(self, val: float):
        if val > self.max:
            col = QtGui.QColor.fromRgb(*RdBuCMap._rdbu[0])
        elif val < self.min:
            col = QtGui.QColor.fromRgb(*RdBuCMap._rdbu[-1])
        else:
            col = QtGui.QColor.fromRgb(*(int(self.spline[i](val)) for i in range(3)))

        return col