# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

class ColormapPresets:

    DEFAULT = "Grayscale Ramp"

    GENERATORS = {
        # Simple color ramp from fully transparent white to semi-transparent gray
        DEFAULT: lambda size, src : ColormapPresets._generate_ramp(size, src),
        "Darker Grayscale Ramp": lambda size, src : ColormapPresets._generate_ramp(size, src, (1, 1, 1, 0), (0, 0, 0, 0.5)),

        # Same, but with one color
        "Red Ramp": lambda size, src : ColormapPresets._generate_ramp(size, src, (1, 0, 0, 0), (0.5, 0.0, 0.0, 0.5)),
        "Green Ramp": lambda size, src : ColormapPresets._generate_ramp(size, src, (0, 1, 0, 0), (0.0, 0.0, 0.0, 0.5)),
        "Blue Ramp": lambda size, src : ColormapPresets._generate_ramp(size, src, (0, 0, 1, 0), (0.0, 0.0, 0.0, 0.5)),

        # Three colors
        "White-Yellow-Red": lambda size, src : ColormapPresets._generate_three_colors(
            size, src, (1.0, 1.0, 1.0), (1.0, 1.0, 0.0), (1.0, 0.0, 0.0), 0.0, 0.5),
        "Red-Green-Blue": lambda size, src : ColormapPresets._generate_three_colors(
            size, src, (1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0), 0.0, 0.5),
        "Blue-White-Red": lambda size, src : ColormapPresets._generate_three_colors(
            size, src, (0.0, 0.0, 1.0), (1.0, 1.0, 1.0), (1.0, 0.0, 0.0), 0.0, 0.5),
        "Black-Red-Yellow": lambda size, src : ColormapPresets._generate_three_colors(
            size, src, (0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (1.0, 1.0, 0.0), 0.0, 0.5),
        "Black-Blue-White": lambda size, src : ColormapPresets._generate_three_colors(
            size, src, (0.0, 0.0, 0.0), (0.0, 0.0, 1.0), (1.0, 1.0, 1.0), 0.0, 0.5),
    }

    def _generate_ramp(colormap_size = 256, colormap_source = "", ramp_min = (1.0, 1.0, 1.0, 0.0), ramp_max = (0.5, 0.5, 0.5, 0.5)):
        if colormap_source == "rgbaPoints":
            return (
                [0, 1],                     # xPoints
                [ramp_min, ramp_max]        # rgbaPoints
                )

        c = [0, 0, 0, 0]
        colormap_data = [tuple(c)] * colormap_size
        for i in range(colormap_size):
            u = i / (colormap_size - 1)
            for j in range(4):
                c[j] = ramp_min[j] * (1.0 - u) + ramp_max[j] * u
            colormap_data[i] = tuple(c)

        return colormap_data

    def _generate_three_colors(colormap_size, colormap_source, color1, color2, color3, opacity_ramp_min = 0.0, opacity_ramp_max = 0.5):
        if colormap_source == "rgbaPoints":
            opacity_ramp_mid = (opacity_ramp_min + opacity_ramp_max) / 2
            return (
                [0, 0.5, 1],                                                                                # xPoints
                [color1 + (opacity_ramp_min,), color2 + (opacity_ramp_mid,), color3 + (opacity_ramp_max,)]  # rgbaPoints
                )

        c = [0, 0, 0, 0]
        colormap_data = [tuple(c)] * colormap_size
        for i in range(colormap_size):
            u = i / (colormap_size - 1)
            for j in range(3):
                if u < 0.5:
                    c[j] = color1[j] * (1.0 - u) + color2[j] * u
                else:
                    c[j] = color2[j] * (1.0 - u) + color3[j] * u

            c[3] = opacity_ramp_min * (1.0 - u) + opacity_ramp_max * u
            colormap_data[i] = tuple(c)

        return colormap_data

    def get_labels():
        return list(ColormapPresets.GENERATORS.keys())

    def generate(label, colormap_size = 256, colormap_source = ""):
        if label in ColormapPresets.GENERATORS:
            return ColormapPresets.GENERATORS[label](colormap_size, colormap_source)
        else:
            return None

    def has(label):
        return (label in ColormapPresets.GENERATORS)
