"""
Functions and classes that relate to colorscales and colorbars
"""
from matplotlib.colors import to_rgb
from paraview.simple import GetColorTransferFunction, GetScalarBar


def get_range(src):
    """
    Return the data range that defines the colorscale

    :rtype: tuple
    :return: (min, max) values of the src
    """
    return src.PointData.GetArray(0).GetRange(0)


def get_name(src):
    """
    Return the name of the data that defines the colorscale

    :rtype: str
    :return: name of the value
    """
    return src.PointData.GetArray(0).Name  # e.g. 'vs'


def colormap(src, colormap, use_above_range=0, use_below_range=0, invert=0,
             num_values=11):
    """
    Create a lookup table that defines the colorscale of the data

    :type colormap: str
    :param colormap: colormap from list of available ParaView colormaps or
        one generated by the user.
    :type use_above_range: str or 0
    :param use_above_range: RGB color that defines colors above given range,
        if 0, above range will not be used
    :param src:
    :return:
    """
    # ParaView requires ints in place of bools
    bool_to_int = {False: 0, True: 1}

    lut = GetColorTransferFunction(get_name(src))
    lut.ApplyPreset(colormap, True)
    lut.NumberOfTableValues = num_values

    # Deal with Above and Below range color
    if bool_to_int[bool(use_above_range)]:
        lut.UseAboveRangeColor = bool_to_int[bool(use_above_range)]
        lut.AboveRangeColor = use_above_range
    if bool_to_int[bool(use_below_range)]:
        lut.UseBelowRangeColor = bool_to_int[bool(use_below_range)]
        lut.BelowRangeColor = use_below_range

    if invert:
        lut.InvertTransferFunction()

    return lut


def colorbar(lut, title="", component_title="", orientation=None,
             position=None, range_label_format=None, label_format="%.2e",
             thickness=35, length=0.15, fontsize=12,
             fontcolor="k", title_justification="Left", text_position="right"):
    """
    Make a colorbar based on the lookup table

    :param lut:
    :return:
    """
    cbar = GetScalarBar(lut)

    cbar.Title = title
    cbar.TitleColor = to_rgb(fontcolor)
    cbar.LabelColor = to_rgb(fontcolor)
    cbar.TitleJustification = title_justification
    cbar.ComponentTitle = component_title

    # Location of title and labels w.r.t colorbar
    assert(text_position in ["top", "right", "bottom", "left"]), \
        "text_position must be in: 'top', 'right', 'bottom', 'left"
    if text_position in ["top", "right"]:
        cbar.TextPosition = "Ticks left/bottom, annotations right/top"
    elif text_position in ["bottom", "left"]:
        cbar.TextPosition = "Ticks right/top, annotations left/bottom"

    if orientation is not None:
        cbar.AutoOrient = 0
        cbar.Orientation = orientation.title()
    else:
        cbar.AutoOrient = 1

    if position:
        cbar.Position = position

    if range_label_format is not None:
        cbar.AddRangeLabels = 1
        cbar.RangeLabelFormat = range_label_format
    else:
        cbar.AddRangeLabels = 0

    if label_format:
        cbar.AutomaticLabelFormat = 0
        cbar.UseCustomLabels = 1
        cbar.LabelFormat = label_format

    cbar.LabelFontSize = fontsize
    cbar.TitleFontSize = fontsize

    cbar.ScalarBarThickness = thickness
    cbar.ScalarBarLength = length




