// 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 {
    Indigo = 0,
    Viridis = 1,
    Magma = 2,
    Inferno = 3,
    Heat = 4,
    BW = 5,
    WB = 6,
    Green = 7
};

enum class ColorScaleSpecial {
    Gap,
    BadPixel
};

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}
    };

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

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

    const std::vector<rgb> magma_colormap = {
        {0, 0, 4}, {28, 16, 68}, {79, 18, 123}, {129, 37, 129},
        {181, 54, 122}, {229, 80, 100}, {251, 135, 97}, {252, 191, 125}, {252, 253, 191}
    };

    const std::vector<rgb> inferno_colormap = {
        {0, 0, 4}, {31, 12, 72}, {85, 15, 109}, {136, 34, 106},
        {186, 54, 85}, {227, 89, 51}, {249, 140, 10}, {252, 195, 63}, {252, 255, 164}
    };

    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);

    static constexpr size_t kLutSize = 512;
    mutable std::vector<rgb> lut_;

    void CalcLUT() const;

public:
    ColorScale();

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

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

#endif //JFJOCH_COLORSCALE_H