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

#include "PreviewImage.h"

#include <cmath>

#include "JFJochJPEG.h"
#include "JFJochTIFF.h"
#include "../common/JFJochException.h"
#include "../common/DiffractionGeometry.h"
#include "../frame_serialize/CBORStream2Deserializer.h"
#include "../compression/JFJochDecompress.h"

constexpr const static rgb lime = {.r = 0xcd, .g = 0xdc, .b = 0x39};
constexpr const static rgb pink = {.r = 0xe9, .g = 0x1e, .b = 0x63};

constexpr const static rgb purple = {.r = 0x7b, .g = 0x1f, .b = 0xA2};
constexpr const static rgb orange = {.r = 0xff, .g = 0x57, .b = 0x22};
constexpr const static rgb amber = {.r =0xff, .g = 0xc1, .b = 0x07};
constexpr const static rgb blue = {.r = 0x0d, .g = 0x47, .b = 0xa1};
constexpr const static rgb cyan = {.r = 0x00, .g = 0xff, .b = 0xff}; // "ice" color

constexpr const static rgb plotly[] = {{0x1f, 0x77, 0xb4},
                                       {0xff, 0x7f, 0x0e},
                                       {0x2c, 0xa0, 0x2c},
                                       {0xd6, 0x27, 0x28},
                                       {0x94, 0x67, 0xbd},
                                       {0x8c, 0x56, 0x4b},
                                       {0xe3, 0x77, 0xc2},
                                       {0x7f, 0x7f, 0x7f},
                                       {0xbd, 0xbd, 0x22},
                                       {0x17, 0xbe, 0xcf}};

constexpr const static rgb gray = {.r = 0xbe, .g = 0xbe, .b = 0xbe};

void PreviewImage::color_pixel(std::vector<rgb> &ret, int64_t in_xpixel, int64_t in_ypixel,const rgb &color) const {
    if ((in_xpixel >= 0) && (in_xpixel < xpixel) && (in_ypixel >= 0) && (in_ypixel < ypixel))
        ret[(in_ypixel * xpixel + in_xpixel)] = color;
}

void PreviewImage::spot(std::vector<rgb> &ret, int64_t in_xpixel, int64_t in_ypixel, const rgb &color) const {
        color_pixel(ret, in_xpixel, in_ypixel, color);
}

void PreviewImage::roi(std::vector<rgb> &ret, int64_t in_xpixel, int64_t in_ypixel, int64_t roi_number) const {
    color_pixel(ret, in_xpixel, in_ypixel, plotly[roi_number % 10]);
}

template<class T>
std::vector<rgb> PreviewImage::GenerateRGB(const uint8_t *value_8,
                                           int64_t special_value_64,
                                           int64_t sat_value_64,
                                           const ColorScale &scale,
                                           const PreviewImageSettings &settings) const {
    auto value = reinterpret_cast<const T *>(value_8);
    auto special_value = static_cast<T>(special_value_64);

    float background = settings.background_value.value_or(0.0);
    float foreground;
    if (settings.saturation_value.has_value())
        foreground = settings.saturation_value.value();
    else {
        // Auto-contrast procedure
        std::vector<int64_t> valid;
        valid.reserve(xpixel * ypixel);
        for (int i = 0; i < xpixel * ypixel; i++) {
            if ((value[i] != special_value)
                && (value[i] != sat_value_64)
                && (mask[i] != MaskDet)
                && (mask[i] != MaskGap)
                && (!settings.show_user_mask || (mask[i] != MaskUsr)))
                valid.push_back(static_cast<int64_t>(value[i]));
        }

        if (!valid.empty()) {
            const size_t m = valid.size();
            size_t ignore = std::max<size_t>(1, static_cast<size_t>(std::floor(m * auto_foreground_range)));
            if (ignore >= m) ignore = m - 1; // ensure at least one value remains
            const size_t rank = m - ignore - 1; // 0-based index for the desired value
            std::nth_element(valid.begin(), valid.begin() + rank, valid.end());
            foreground = static_cast<float>(valid[rank]);
        } else {
            // Fallback to something above background if no valid pixels remain
            foreground = background + 1.0f;
        }

    }

    std::vector<rgb> ret(xpixel * ypixel);
    for (int i = 0; i < xpixel*ypixel; i++) {
        if (mask[i] == MaskGap)
            ret[i] = scale.Apply(ColorScaleSpecial::Gap);
        else if ((value[i] == special_value)
                 || (mask[i] == MaskDet)
                 || (settings.show_user_mask && (mask[i] == MaskUsr))) {
            ret[i] = scale.Apply(ColorScaleSpecial::BadPixel);
        } else
            ret[i] = scale.Apply(value[i], background, foreground);
    }
    return ret;
}

void PreviewImage::AddBeamCenter(std::vector<rgb> &rgb_image) const {
    size_t beam_x_int = std::lround(beam_x);
    size_t beam_y_int = std::lround(beam_y);

    int crosshair_size = 30;
    int crosshair_width = 3;
    for (int w = -crosshair_width; w <= crosshair_width; w++) {
        for (int i = -crosshair_size; i <= crosshair_size; i++) {
            color_pixel(rgb_image, beam_x_int + i, beam_y_int + w, lime);
            color_pixel(rgb_image, beam_x_int + w, beam_y_int + i, lime);
        }
    }
}

void PreviewImage::AddSpots(std::vector<rgb> &rgb_image,
                            const std::vector<SpotToSave>& in_spots) const {
    for (const auto &s: in_spots) {
        int64_t spot_x_int = std::lround(s.x);
        int64_t spot_y_int = std::lround(s.y);

        int rectangle_size = 4;
        int rectangle_width = 3;

        rgb color = pink;
        if (s.indexed)
            color = amber;
        else if (s.ice_ring)
            color = cyan;

        for (int z = rectangle_size; z < rectangle_size + rectangle_width; z++) {
            for (int w = -z; w <= z; w++) {
                spot(rgb_image, spot_x_int + z, spot_y_int + w, color);
                spot(rgb_image, spot_x_int - z, spot_y_int + w, color);
                spot(rgb_image, spot_x_int + w, spot_y_int + z, color);
                spot(rgb_image, spot_x_int + w, spot_y_int - z, color);
            }
        }
    }
}

void PreviewImage::AddROI(std::vector<rgb> &rgb_image) const {
    int64_t roi_counter = 0;

    for (const auto &box: experiment.ROI().GetROIDefinition().boxes) {
        int rectangle_width = 5;

        for (auto x = box.GetXMin() - rectangle_width; x <= box.GetXMax() + rectangle_width; x++) {
            for (auto w = 1; w <= rectangle_width; w++) {
                roi(rgb_image, x, box.GetYMax() + w, roi_counter);
                roi(rgb_image, x, box.GetYMin() - w, roi_counter);
            }
        }

        for (auto y = box.GetYMin() - rectangle_width; y <= box.GetYMax() + rectangle_width; y++) {
            for (auto w = 1; w <= rectangle_width; w++) {
                roi(rgb_image, box.GetXMax() + w, y, roi_counter);
                roi(rgb_image, box.GetXMin() - w, y, roi_counter);
            }
        }
        roi_counter++;
    }

    for (const auto &circle: experiment.ROI().GetROIDefinition().circles) {
        int width = 5;

        for (int64_t y = std::floor(circle.GetY() - circle.GetRadius_pxl() - width);
             y <= std::ceil(circle.GetY() + circle.GetRadius_pxl() + width);
             y++) {
            for (int64_t x = std::floor(circle.GetX() - circle.GetRadius_pxl() - width);
                 x <= std::ceil(circle.GetX() + circle.GetRadius_pxl() + width);
                 x++) {
                    float dist = sqrtf((x - circle.GetX()) * (x - circle.GetX())
                        + (y - circle.GetY()) * (y - circle.GetY()));

                    if ((dist > circle.GetRadius_pxl()) && (dist < circle.GetRadius_pxl() + width))
                        roi(rgb_image, x, y, roi_counter);
            }
        }
        roi_counter++;
    }
}

void PreviewImage::AddResolutionRing(std::vector<rgb> &rgb_image, float d) const {
    DiffractionGeometry geom = experiment.GetDiffractionGeometry();
    int width = 3;
    float radius =  geom.ResToPxl(d);
    for (int64_t y = 0; y <= ypixel; y++) {
        for (int64_t x = 0; x <= xpixel; x++) {
            float dist = sqrtf((x - beam_x) * (x - beam_x) + (y - beam_y) * (y - beam_y));

            if ((dist > radius) && (dist < radius + width))
                color_pixel(rgb_image, x, y, orange);
        }
    }
}

void PreviewImage::Configure(const DiffractionExperiment &in_experiment, const PixelMask& pixel_mask) {
    std::unique_lock ul(m);

    experiment = in_experiment;
    xpixel = experiment.GetXPixelsNum();
    ypixel = experiment.GetYPixelsNum();
    beam_x = experiment.GetBeamX_pxl();
    beam_y = experiment.GetBeamY_pxl();
    pixel_depth_bytes = experiment.GetByteDepthImage();
    pixel_is_signed = experiment.IsPixelSigned();

    mask = std::vector<uint8_t>(experiment.GetPixelsNum(), 0);

    std::vector<uint32_t> mask_tmp = pixel_mask.GetMask(experiment);
    for (int i = 0; i < experiment.GetPixelsNum(); i++) {
        if (((mask_tmp[i] & (1 << PixelMask::ModuleGapPixelBit)) != 0)
            || ((mask_tmp[i] & (1 << PixelMask::ChipGapPixelBit)) != 0)
            || ((mask_tmp[i] & (1 << PixelMask::ModuleEdgePixelBit)) != 0))
            mask[i] = MaskGap;
        else if (mask_tmp[i] & 0xFEFE) // bits 1-7 and 9-15
            mask[i] = MaskDet;
        else if ((mask_tmp[i] & (1 << PixelMask::UserMaskedPixelBit)) != 0)
            mask[i] = MaskUsr;
    }
}

void PreviewImage::Configure() {
    std::unique_lock ul(m);
    xpixel = 0;
    ypixel = 0;
}

std::vector<rgb> PreviewImage::GenerateRGB(const PreviewImageSettings &settings, const DataMessage &msg) const {
    std::vector<rgb> v(msg.image.GetWidth() * msg.image.GetHeight());
    if (msg.image.GetUncompressedSize() == 0)
        return {};

    std::vector<uint8_t> tmp;
    const uint8_t* image_ptr = msg.image.GetUncompressedPtr(tmp);

    {
        // JPEG compression is outside the critical loop protected by m
        std::unique_lock ul(m);

        ColorScale scale;
        scale.Select(settings.scale);

        switch (msg.image.GetMode()) {
            case CompressedImageMode::Int8:
                v = GenerateRGB<int8_t>(image_ptr, INT8_MIN, INT8_MAX, scale, settings);
                break;
            case CompressedImageMode::Int16:
                v = GenerateRGB<int16_t>(image_ptr, INT16_MIN, INT16_MAX, scale, settings);
                break;
            case CompressedImageMode::Int32:
                v = GenerateRGB<int32_t>(image_ptr, INT32_MIN, INT32_MAX, scale, settings);
                break;
            case CompressedImageMode::Uint8:
                v = GenerateRGB<uint8_t>(image_ptr,UINT8_MAX, UINT8_MAX, scale, settings);
                break;
            case CompressedImageMode::Uint16:
                v = GenerateRGB<uint16_t>(image_ptr,UINT16_MAX, UINT16_MAX, scale, settings);
                break;
            case CompressedImageMode::Uint32:
                v = GenerateRGB<uint32_t>(image_ptr,UINT32_MAX, UINT32_MAX, scale, settings);
                break;
            default:
                throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mode not supported");
        }

        if (settings.show_spots)
            AddSpots(v, msg.spots);

        if (settings.show_roi)
            AddROI(v);

        if (settings.resolution_ring)
            AddResolutionRing(v, settings.resolution_ring.value());
        else if (settings.show_res_est && msg.resolution_estimate)
            AddResolutionRing(v, msg.resolution_estimate.value());

        if (settings.show_beam_center)
            AddBeamCenter(v);
    }
    return v;
}

std::string PreviewImage::GenerateImage(const PreviewImageSettings& settings, const DataMessage &msg) const {
    auto v = GenerateRGB(settings, msg);
    CompressedImage rgb_image(v, msg.image.GetWidth(), msg.image.GetHeight());
    switch (settings.format) {
        case PreviewImageFormat::JPEG:
            return WriteJPEGToMem(rgb_image, settings.jpeg_quality);
        case PreviewImageFormat::TIFF:
            return WriteTIFFToString(rgb_image);
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Preview image format not supported");
    }
}

std::string PreviewImage::GenerateImage(const PreviewImageSettings &settings, const std::vector<uint8_t> &cbor_format) {
    auto cbor = CBORStream2Deserialize(cbor_format);
    if (!cbor || !cbor->data_message)
        return {};

    return GenerateImage(settings, *cbor->data_message);
}

std::string PreviewImage::GenerateTIFF(const std::vector<uint8_t>& cbor_format) {
    auto cbor = CBORStream2Deserialize(cbor_format);
    if (!cbor || !cbor->data_message)
        return {};

    return WriteTIFFToString(cbor->data_message->image);
}