// Copyright (2019-2024) Paul Scherrer Institute

#include "PreviewImage.h"

#include <cmath>

#include "WriteJPEG.h"
#include "WriteTIFF.h"
#include "../common/JFJochException.h"
#include "../common/DiffractionGeometry.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 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 indigo = {.r = 0x3f, .g = 0x51, .b = 0xb5};
constexpr const static rgb gray = {.r = 0xbe, .g = 0xbe, .b = 0xbe};

PreviewImage::PreviewImage(std::chrono::microseconds period) : counter(period) {}

void colormap(std::vector<unsigned char>& ret, float v, size_t pixel) {
    if ((v < 0.0) || (v > 1.0)) {
        ret[pixel * 3] = gray.r;
        ret[pixel * 3+1] = gray.g;
        ret[pixel * 3+2] = gray.b;
    } else {
        ret[pixel * 3] = 255 - std::lround((255-indigo.r) * v);
        ret[pixel * 3 + 1] = 255 - std::lround((255-indigo.g) * v);
        ret[pixel * 3 + 2] = 255 - std::lround((255-indigo.b) * v);
    }
}

void PreviewImage::color_pixel(std::vector<unsigned char> &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) * 3] = color.r;
        ret[(in_ypixel * xpixel + in_xpixel) * 3 + 1] = color.g;
        ret[(in_ypixel * xpixel + in_xpixel) * 3 + 2] = color.b;
    }
}

void PreviewImage::spot(std::vector<unsigned char> &ret, int64_t in_xpixel, int64_t in_ypixel, bool indexed) const {
    if (indexed)
        color_pixel(ret, in_xpixel, in_ypixel, pink);
    else
        color_pixel(ret, in_xpixel, in_ypixel, amber);
}

void PreviewImage::roi(std::vector<unsigned char> &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]);
}

void PreviewImage::beam_center_mark(std::vector<unsigned char> &ret, int64_t in_xpixel, int64_t in_ypixel) const {
    color_pixel(ret, in_xpixel, in_ypixel, lime);
}

template<class T>
std::vector<unsigned char> GenerateRGB(const T* value, size_t npixel, uint32_t saturation_value, T special_value) {
    std::vector<unsigned char> ret(3*npixel, 0);
    for (int i = 0; i < npixel; i++) {
        if (value[i] == special_value) {
            colormap(ret, -1.0, i);
        } else if (value[i] >= static_cast<T>(saturation_value))
            colormap(ret, 1.0, i);
        else if (value[i] <= 0)
            colormap(ret, 0.0, i);
        else {
            colormap(ret, value[i] / static_cast<float>(saturation_value), i);
        }
    }
    return ret;
}

template <class T>
std::vector<uint16_t> GenerateDioptasPreview(const void* input, size_t xpixel, size_t ypixel, T special_value) {
    auto input_ptr = (T *) input;

    std::vector<uint16_t> vec(xpixel * ypixel);
    for (int y = 0; y < ypixel; y++) {
        for (int x = 0; x < xpixel; x++) {
            T val = input_ptr[(ypixel - 1 - y) * xpixel + x];
            if ((val == special_value) || (val >= INT16_MAX))
                vec[y * xpixel + x] = INT16_MAX;
            else if (val <= 0)
                vec [y * xpixel + x] = 0;
            else
                vec[y * xpixel + x] = val;
        }
    }
    return vec;
}

void PreviewImage::AddBeamCenter(std::vector<uint8_t> &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++) {
            beam_center_mark(rgb_image, beam_x_int + i, beam_y_int + w);
            beam_center_mark(rgb_image, beam_x_int + w, beam_y_int + i);
        }
    }
}

void PreviewImage::AddSpots(std::vector<uint8_t> &rgb_image) const {
    for (const auto &s: 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;

        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, s.indexed);
                spot(rgb_image, spot_x_int - z, spot_y_int + w, s.indexed);
                spot(rgb_image, spot_x_int + w, spot_y_int + z, s.indexed);
                spot(rgb_image, spot_x_int + w, spot_y_int - z, s.indexed);
            }
        }
    }
}

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

    for (const auto &box: experiment.ROI().GetROIBox()) {
        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().GetROICircle()) {
        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<uint8_t> &rgb_image, float d) const {
    int width = 3;
    float radius =  ResToPxl(experiment, 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) {
    std::unique_lock<std::mutex> ul(m);

    experiment = in_experiment;
    initialized = false;
    xpixel = experiment.GetXPixelsNum();
    ypixel = experiment.GetYPixelsNum();
    beam_x = experiment.GetBeamX_pxl();
    beam_y = experiment.GetBeamY_pxl();
    pixel_depth_bytes = experiment.GetPixelDepth();
    pixel_is_signed = experiment.IsPixelSigned();
    uncompressed_image.resize(experiment.GetPixelsNum() * experiment.GetPixelDepth());
    memset(uncompressed_image.data(), 0, experiment.GetPixelsNum() * experiment.GetPixelDepth());
}

void PreviewImage::Configure() {
    std::unique_lock<std::mutex> ul(m);

    initialized = false;
    xpixel = 0;
    ypixel = 0;
}

void PreviewImage::UpdateImage(const void *in_uncompressed_image,
                               const std::vector<SpotToSave> &in_spots) {
    if (counter.GeneratePreview()) {
        std::unique_lock<std::mutex> ul(m);

        if (xpixel * ypixel == 0)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Preview not configured");

        initialized = true;
        memcpy(uncompressed_image.data(), in_uncompressed_image, xpixel * ypixel * pixel_depth_bytes);
        spots = in_spots;
    }
}

std::string PreviewImage::GenerateJPEG(const PreviewJPEGSettings &settings) const {
    std::vector<unsigned char> v;
    size_t local_xpixel;
    size_t local_ypixel;

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

        local_xpixel = xpixel;
        local_ypixel = ypixel;

        if (!initialized)
            return {};

        if (!pixel_is_signed) {
            if (pixel_depth_bytes == 2)
                v = GenerateRGB<uint16_t>((uint16_t *) uncompressed_image.data(), xpixel * ypixel,
                                          settings.saturation_value, UINT16_MAX);
            else
                v = GenerateRGB<uint32_t>((uint32_t *) uncompressed_image.data(), xpixel * ypixel,
                                          settings.saturation_value, UINT32_MAX);
        } else {
            if (pixel_depth_bytes == 2)
                v = GenerateRGB<int16_t>((int16_t *) uncompressed_image.data(), xpixel * ypixel,
                                         settings.saturation_value, INT16_MIN);
            else
                v = GenerateRGB<int32_t>((int32_t *) uncompressed_image.data(), xpixel * ypixel,
                                         settings.saturation_value, INT32_MIN);
        }
        if (settings.show_spots)
            AddSpots(v);

        if (settings.show_roi)
            AddROI(v);

        if (settings.resolution_ring)
            AddResolutionRing(v, settings.resolution_ring.value());

        AddBeamCenter(v);
    }

    return WriteJPEGToMem(v, local_xpixel, local_ypixel, settings.jpeg_quality);
}


std::string PreviewImage::GenerateTIFF() const {
    std::unique_lock<std::mutex> ul(m);
    if (!initialized)
        return {};

    std::string s = WriteTIFFToString(const_cast<uint8_t *>(uncompressed_image.data()),
                                      xpixel, ypixel, pixel_depth_bytes, pixel_is_signed);
    return s;
}

std::string PreviewImage::GenerateTIFFDioptas() const {
    std::unique_lock<std::mutex> ul(m);
    if (!initialized)
        return {};

    std::vector<uint16_t> vec;
    if (pixel_is_signed) {
        if (pixel_depth_bytes == 2)
            vec = GenerateDioptasPreview<int16_t>(uncompressed_image.data(), xpixel, ypixel, INT16_MIN);
        else
            vec = GenerateDioptasPreview<int32_t>(uncompressed_image.data(), xpixel, ypixel, INT32_MIN);
    } else {
        if (pixel_depth_bytes == 2)
            vec = GenerateDioptasPreview<uint16_t>(uncompressed_image.data(), xpixel, ypixel, UINT16_MAX);
        else
            vec = GenerateDioptasPreview<uint32_t>(uncompressed_image.data(), xpixel, ypixel, UINT32_MAX);
    }
    return WriteTIFFToString(vec.data(), xpixel, ypixel, 2, false);
}