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

#include <cmath>
#include "ColorScale.h"
#include "JFJochException.h"

static inline float Clamp01(float x) {
    return (x < 0.0f) ? 0.0f : (x > 1.0f ? 1.0f : x);
}

// Gamma-mapped green (recommended gamma = 0.7)
static inline rgb GreenGamma(float f, float gamma = 0.7f) {
    f = Clamp01(f);
    const float g = std::pow(f, gamma);
    const uint8_t G = static_cast<uint8_t>(std::lround(255.0f * g));
    return {.r = 0, .g = G, .b = 0};
}

// Asinh-mapped green (recommended k = 8.0)
static inline rgb GreenAsinh(float f, float k = 8.0f) {
    f = Clamp01(f);
    const float g = std::asinh(k * f) / std::asinh(k);
    const uint8_t G = static_cast<uint8_t>(std::lround(255.0f * g));
    return {.r = 0, .g = G, .b = 0};
}

float luminance(rgb input) {
    return 0.2126f * input.r + 0.7152f * input.g + 0.0722f * input.b;
}

ColorScale::ColorScale() : lut_(kLutSize) {
    CalcLUT();
}

void ColorScale::Select(ColorScaleEnum val) {
    current = val;
    CalcLUT();
}

const std::vector<rgb> &ColorScale::LUTData() const {
    return lut_;
}

void ColorScale::CalcLUT() const {
    const std::vector<rgb>* map = nullptr;
    switch (current) {
        case ColorScaleEnum::Viridis: map = &viridis_colormap; break;
        case ColorScaleEnum::Heat: map = &heat_colormap;
            break;
        case ColorScaleEnum::Indigo: map = &white_to_indigo_colormap;
            break;
        case ColorScaleEnum::BW: map = &white_to_black_colormap;
            break;
        case ColorScaleEnum::WB: map = &black_to_white_colormap;
            break;
        case ColorScaleEnum::Green:
            map = &green_colormap;
            break;
        case ColorScaleEnum::Magma:
            map = &magma_colormap; break;
        case ColorScaleEnum::Inferno:
            map = &inferno_colormap; break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Color scale unknown");
    }

    for (size_t i = 0; i < kLutSize; ++i) {
        const float f = static_cast<float>(i) / static_cast<float>(kLutSize - 1);
        lut_[i] = Apply(f, *map);
    }
}

rgb ColorScale::Apply(float input, const std::vector<rgb> &map) {
    size_t num_colors = map.size();
    if (num_colors < 2) {
        throw std::invalid_argument("Colormap must have at least two colors.");
    }

    float scaled_value = input * (num_colors - 1);
    size_t lower_idx = static_cast<size_t>(scaled_value);
    size_t upper_idx = std::min(lower_idx + 1, num_colors - 1);
    float t = scaled_value - lower_idx; // Fraction for interpolation

    rgb lower = map[lower_idx];
    rgb upper = map[upper_idx];

    return {
            .r = static_cast<uint8_t>(lower.r + t * (upper.r - lower.r)),
            .g = static_cast<uint8_t>(lower.g + t * (upper.g - lower.g)),
            .b = static_cast<uint8_t>(lower.b + t * (upper.b - lower.b))
    };
}

rgb ColorScale::Apply(ColorScaleSpecial input) const {
    switch (input) {
        case ColorScaleSpecial::Gap:
            return gap;
        default:
        case ColorScaleSpecial::BadPixel:
            return bad;
    }
}

rgb ColorScale::Apply(float input, float min, float max) const {
    if (!std::isfinite(input))
        return gap;

    float f;
    if (input <= min)
        f = 0.0f;
    else if (input >= max)
        f = 1.0f;
    else
        f = (input - min) / (max - min);

    const size_t idx = static_cast<size_t>(f * static_cast<float>(kLutSize - 1));
    return lut_[idx];
}

rgb ColorScale::ApplyLUTIndex(size_t idx) const {
    if (idx >= kLutSize)
        return lut_[kLutSize-1];
    return lut_[idx];
}

ColorScale &ColorScale::Gap(rgb input) {
    gap = input;
    return *this;
}

ColorScale &ColorScale::BadPixel(rgb input) {
    bad = input;
    return *this;
}