// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#ifndef JFJOCH_COLORSCALE_H
#define JFJOCH_COLORSCALE_H

#include <vector>
#include <cstdint>
#include <algorithm>
#include <stdexcept>

struct rgb {
    uint8_t r;
    uint8_t g;
    uint8_t b;
};

float luminance(rgb input);

enum class ColorScaleEnum : int {
    Viridis = 0,
    Heat = 1,
    Indigo = 2,
    BW = 3
};

enum class ColorScaleSpecial {
    Gap,
    BadPixel
};

rgb rainbowColor(float t);

class ColorScale {
    const std::vector<rgb> viridis_colormap = {
            {68, 1, 84}, {71, 44, 123}, {59, 81, 139}, {44, 113, 142}, {33, 144, 141},
            {39, 173, 129}, {92, 200, 99}, {170, 220, 50}, {253, 231, 37}
    };

    const std::vector<rgb> heat_colormap = {
            {0, 0, 0}, {128, 0, 0}, {255, 0, 0}, {255, 128, 0}, {255, 255, 0}, {255, 255, 128}, {255, 255, 255}
    };

    const std::vector<rgb> white_to_indigo_colormap = {
            {255, 255, 255}, {216, 209, 255}, {177, 162, 255}, {138, 114, 255}, {99, 67, 255}, {60, 19, 255}, {0, 0, 128}
    };

    const std::vector<rgb> white_to_black_colormap = {
            {255, 255, 255}, {0, 0, 0}
    };

    ColorScaleEnum current = ColorScaleEnum::Indigo;

    rgb gap = {.r = 190, .g = 190, .b = 190}; // Gray
    rgb bad = {.r = 255, .g =   0, .b = 255}; // Magenta

    static rgb Apply(float input, const std::vector<rgb> &map);
public:
    void Select(ColorScaleEnum val);
    [[nodiscard]] rgb Apply(float input, float min = 0.0, float max = 1.0) const;
    [[nodiscard]] rgb Apply(ColorScaleSpecial input) const;

    ColorScale &Gap(rgb input);
    ColorScale &BadPixel(rgb input);
};

#endif //JFJOCH_COLORSCALE_H