// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include <cstring>

#include <bitshuffle/bitshuffle.h>

#include "FrameTransformation.h"
#include "RawToConvertedGeometry.h"
#include "JFJochException.h"

FrameTransformation::FrameTransformation(const DiffractionExperiment &in_experiment) :
        experiment(in_experiment), summation(experiment.GetSummation()),
        pixel_depth(experiment.GetPixelDepth()), compressor(in_experiment.GetCompressionAlgorithmEnum()),
        compression_algorithm(in_experiment.GetCompressionAlgorithmEnum()),
        line_shift((experiment.IsUpsideDown() ? -1 : 1) * experiment.GetXPixelsNum()),
        binning_2x2(experiment.GetBinning2x2()) {

    if ((experiment.GetDetectorMode() == DetectorMode::Conversion) && (summation > 1)) {
        for (int i = 0; i < experiment.GetModulesNum(); i++)
            summation_buffer.emplace_back(RAW_MODULE_SIZE);
    }

    precompression_buffer.resize(experiment.GetPixelsNum() * pixel_depth);
    if (pixel_depth == 4)
        image16bit.resize(experiment.GetPixelsNum(), 0);

    if (experiment.GetApplyPixelMaskInFPGA()) {
        // Mask gaps
        if (pixel_depth == 2) {
            auto ptr = (int16_t *) precompression_buffer.data();
            for (int i = 0; i < experiment.GetPixelsNum(); i++)
                ptr[i] = INT16_MIN;
        } else {
            auto ptr = (uint32_t *) precompression_buffer.data();
            for (int i = 0; i < experiment.GetPixelsNum(); i++)
                ptr[i] = INT32_MIN;
        }
    }
}

FrameTransformation& FrameTransformation::SetOutput(void *output) {
    standard_output = (char *) output;
    return *this;
}

template <class Td> void AddToFramesSum(Td *destination, const int16_t *source) {
    for (int i = 0; i < RAW_MODULE_SIZE; i++) {
        if ((source[i] == INT16_MIN) || (destination[i] == INT32_MIN))
            destination[i] = INT32_MIN;
        else if ((source[i] == INT16_MAX) || (destination[i] == INT32_MAX))
            destination[i] = INT32_MAX;
        else destination[i] += source[i];
    }
}

void FrameTransformation::PackSummation() {
    for (int m = 0; m < experiment.GetModulesNum(); m++) {
        void *output = precompression_buffer.data() + sizeof(int32_t) * experiment.GetPixel0OfModule(m);
        if (binning_2x2)
            TransferModuleAdjustMultipixelsBin2x2((int32_t *) output,
                                                  (int32_t *) summation_buffer[m].data(),
                                                  line_shift,
                                                  static_cast<int32_t>(INT32_MIN),
                                                  static_cast<int32_t>(INT32_MAX));
        else
            TransferModuleAdjustMultipixels((int32_t *) output,
                                            (int32_t *) summation_buffer[m].data(),
                                            line_shift,
                                            static_cast<int32_t>(INT32_MIN),
                                            static_cast<int32_t>(INT32_MAX));

        for (auto &i: summation_buffer[m])
            i = 0;
    }

    if (pixel_depth == 4) {
        // Generate 16-bit preview image
        auto arr = (int32_t *) precompression_buffer.data();

        for (int pxl = 0; pxl < experiment.GetPixelsNum(); pxl++) {
            if (arr[pxl] >= INT16_MAX)
                image16bit[pxl] = INT16_MAX;
            else if (arr[pxl] <= INT16_MIN)
                image16bit[pxl] = INT16_MIN;
            else
                image16bit[pxl] = static_cast<int16_t>(arr[pxl]);
        }
    }
}

size_t FrameTransformation::PackStandardOutput() {
    if (summation > 1)
        PackSummation();

    return compressor.Compress(standard_output, precompression_buffer.data(),
                               experiment.GetPixelsNum(), pixel_depth);
}

void FrameTransformation::ProcessModule(const int16_t *input, size_t frame_number, uint16_t module_number,
                                        int data_stream) {
    if (standard_output == nullptr)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Default stream output not initialized");

    size_t module_number_abs = experiment.GetFirstModuleOfDataStream(data_stream) + module_number;

    if (summation == 1) {
        int16_t *output;

        output = ((int16_t *) precompression_buffer.data()) + experiment.GetPixel0OfModule(module_number_abs);

        if (experiment.GetDetectorMode() != DetectorMode::Conversion)
            memcpy(output, input, RAW_MODULE_SIZE * experiment.GetPixelDepth());
        else if (binning_2x2)
            TransferModuleAdjustMultipixelsBin2x2((int16_t *) output, (int16_t *) input, line_shift,
                                                  static_cast<int16_t>(INT16_MIN),
                                                  static_cast<int16_t>(INT16_MAX));
        else
            TransferModuleAdjustMultipixels((int16_t *) output, (int16_t *) input, line_shift,
                                            static_cast<int16_t>(INT16_MIN),
                                            static_cast<int16_t>(INT16_MAX));
    } else {
        AddToFramesSum(summation_buffer[module_number_abs].data(), input);
    }
}

int16_t *FrameTransformation::GetPreview16BitImage() {
    if (pixel_depth == 2)
        return (int16_t *) precompression_buffer.data();
    else
        return image16bit.data();
}

void FrameTransformation::ProcessModule(JFConversion &conv, const int16_t *input, size_t frame_number,
                                        uint16_t module_number, int data_stream) {
    if (standard_output == nullptr)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Default stream output not initialized");

    size_t module_number_abs = experiment.GetFirstModuleOfDataStream(data_stream) + module_number;

    auto output = ((int16_t *) precompression_buffer.data()) + experiment.GetPixel0OfModule(module_number_abs);

    if (experiment.GetDetectorMode() != DetectorMode::Conversion)
        memcpy(output, input, RAW_MODULE_SIZE * experiment.GetPixelDepth());
    else
        conv.ConvertAdjustGeom((int16_t *) output, (uint16_t *) input, line_shift);
}