// Copyright (2019-2024) Paul Scherrer Institute

#include "PreviewImage.h"

#include <cmath>

#include "WriteJPEG.h"
#include "WriteTIFF.h"
#include "../common/JFJochException.h"

PreviewImage::PreviewImage(const DiffractionExperiment &experiment) :
        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(experiment.GetPixelsNum() * experiment.GetPixelDepth()),
        roi_mask(experiment.ROI().GetMask()) {}

void PreviewImage::UpdateImage(const void *in_uncompressed_image,
                               const std::vector<SpotToSave> &in_spots) {
    std::unique_lock<std::mutex> ul(m);
    memcpy(uncompressed_image.data(), in_uncompressed_image, xpixel * ypixel * pixel_depth_bytes);
    spots = in_spots;
}

#define JFJOCH_JPEG_BASE_R (0x3f)
#define JFJOCH_JPEG_BASE_G (0x51)
#define JFJOCH_JPEG_BASE_B (0xb5)

void colormap(std::vector<unsigned char>& ret, float v, size_t pixel) {
    if ((v < 0.0) || (v > 1.0)) {
        ret[pixel * 3] = 0xBE;
        ret[pixel * 3+1] = 0xBE;
        ret[pixel * 3+2] = 0xBE;
    } else {
        ret[pixel * 3] = 255 - std::lround((255-JFJOCH_JPEG_BASE_R) * v);
        ret[pixel * 3 + 1] = 255 - std::lround((255-JFJOCH_JPEG_BASE_G) * v);
        ret[pixel * 3 + 2] = 255 - std::lround((255-JFJOCH_JPEG_BASE_B) * v);
    }
}

void roi(std::vector<unsigned char>& ret, size_t pixel) {
    ret[pixel * 3] = 0x7B;
    ret[pixel * 3+1] = 0x1F;
    ret[pixel * 3+2] = 0xA2;
}

void feature(std::vector<unsigned char>& ret, size_t pixel) {
    ret[pixel * 3] = 0xE9;
    ret[pixel * 3+1] = 0x1E;
    ret[pixel * 3+2] = 0x63;
}

void spot(std::vector<unsigned char>& ret, int64_t pixel, bool indexed) {
    if (indexed) {
        ret[pixel * 3] = 0xE9;
        ret[pixel * 3 + 1] = 0x1E;
        ret[pixel * 3 + 2] = 0x63;
    } else {
        ret[pixel * 3] = 0xCD;
        ret[pixel * 3 + 1] = 0xDC;
        ret[pixel * 3 + 2] = 0x39;
    }
}


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

std::string PreviewImage::GenerateJPEG(const PreviewJPEGSettings &settings) const {
    if (!pixel_is_signed)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Unsigned images not allowed");

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

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

    AddBeamCenter(v);

    return WriteJPEGToMem(v, xpixel, ypixel, settings.jpeg_quality);
}

std::string PreviewImage::GenerateTIFF() const {
    std::unique_lock<std::mutex> ul(m);
    std::string s = WriteTIFFToString(const_cast<uint8_t *>(uncompressed_image.data()),
                                      xpixel, ypixel, pixel_depth_bytes, pixel_is_signed);
    return s;
}

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

std::string PreviewImage::GenerateTIFFDioptas() const {
    std::unique_lock<std::mutex> ul(m);

    std::vector<uint16_t> vec;
    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);

    return WriteTIFFToString(vec.data(), xpixel, ypixel, 2, false);
}

std::string PreviewImage::GenerateTestJPEG(const PreviewJPEGSettings &settings) {
    std::vector<unsigned char> v;
    size_t xpixel = 1030 * 3 + 2 * 8;
    size_t ypixel =  514 * 6 + 5 * 36;

    std::vector<int16_t> uncompressed_image(xpixel * ypixel, 5);

    for (int y = 0; y < ypixel; y++) {
        for (int gap = 1; gap <= 2; gap++) {
            for (int pixel = 0; pixel < 8; pixel++)
                uncompressed_image[y * xpixel + 1030 * gap + 8 * (gap - 1) + pixel] = INT16_MIN;
        }
    }

    for (int gap = 1; gap <= 5; gap++) {
        for (int pixel = 0; pixel < 36; pixel++) {
            for (int x = 0; x < xpixel; x++)
                uncompressed_image[(gap * 514 + 36 * (gap - 1) + pixel) * xpixel + x] = INT16_MIN;
        }
    }

    v = GenerateRGB<int16_t>((int16_t *) uncompressed_image.data(), xpixel * ypixel, settings.saturation_value,
                             INT16_MIN);

    size_t beam_x_int = std::lround(1200);
    size_t beam_y_int = std::lround(800);

    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++) {
            feature(v, (beam_y_int + w) * xpixel + beam_x_int + i);
            feature(v, (beam_y_int + i) * xpixel + beam_x_int + w);
        }
    }

    return WriteJPEGToMem(v, xpixel, ypixel, settings.jpeg_quality);
}

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++) {
            feature(rgb_image, (beam_y_int + w) * xpixel + beam_x_int + i);
            feature(rgb_image, (beam_y_int + i) * xpixel + beam_x_int + w);
        }
    }
}

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;

        if ((spot_x_int - rectangle_width < 0)
            || (spot_x_int + rectangle_width >= xpixel)
            || (spot_y_int - rectangle_width < 0)
            || (spot_y_int + rectangle_width >= ypixel))
            continue;

        for (int z = rectangle_size; z < rectangle_size + rectangle_width; z++) {
            for (int w = -z; w <= z; w++) {
                spot(rgb_image, (spot_y_int + w) * xpixel + spot_x_int + z, s.indexed);
                spot(rgb_image, (spot_y_int + w) * xpixel + spot_x_int - z, s.indexed);
                spot(rgb_image, (spot_y_int + z) * xpixel + spot_x_int + w, s.indexed);
                spot(rgb_image, (spot_y_int - z) * xpixel + spot_x_int - w, s.indexed);
            }
        }
    }
}

void PreviewImage::AddROI(std::vector<uint8_t> &rgb_image) const {
    if (rgb_image.size() != roi_mask.size() * 3)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in array size");

    for (int i = 0; i < roi_mask.size(); i++) {
        if (roi_mask[i] != UINT16_MAX)
            roi(rgb_image, i);
    }
}