// 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()),
        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 (binning_2x2)
        full_image_buffer.resize(experiment.GetPixelsNumFullImage() * pixel_depth);

    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++) {
        int32_t *output;
        if (binning_2x2)
            output = (int32_t *) full_image_buffer.data();
        else
            output = (int32_t *) precompression_buffer.data();

        TransferModuleAdjustMultipixels(output,
                                        (int32_t *) summation_buffer[m].data(),
                                        experiment.GetModuleSlowDirectionStep(m),
                                        static_cast<int32_t>(INT32_MIN),
                                        static_cast<int32_t>(INT32_MAX),
                                        experiment.GetModuleFastDirectionStep(m),
                                        experiment.GetPixel0OfModule(m));

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

    if (binning_2x2)
        Bin2x2_sum<int32_t>((int32_t *) precompression_buffer.data(),
                            (int32_t *) full_image_buffer.data(),
                            experiment.GetXPixelsNumFullImage(),
                            experiment.GetYPixelsNumFullImage(),
                            static_cast<int32_t>(experiment.GetUnderflow()),
                            static_cast<int32_t>(experiment.GetOverflow()));

    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();
    else {
        if (binning_2x2)
            Bin2x2_sum<int16_t>((int16_t *) precompression_buffer.data(),
                                (int16_t *) full_image_buffer.data(),
                                experiment.GetXPixelsNumFullImage(),
                                experiment.GetYPixelsNumFullImage(),
                                INT16_MIN, INT16_MAX);
    }
    return compressor.Compress(standard_output, precompression_buffer.data(),
                               experiment.GetPixelsNum(), pixel_depth);
}

void FrameTransformation::ProcessModule(const int16_t *input, 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;

        if (binning_2x2)
            output = (int16_t *) full_image_buffer.data();
        else
            output = (int16_t *) precompression_buffer.data();

        if (experiment.GetDetectorMode() != DetectorMode::Conversion)
            memcpy(output + RAW_MODULE_SIZE * module_number_abs,
                   input,
                   RAW_MODULE_SIZE * experiment.GetPixelDepth());
        else
            TransferModuleAdjustMultipixels(output, (int16_t *) input,
                                            experiment.GetModuleSlowDirectionStep(module_number_abs),
                                            static_cast<int16_t>(INT16_MIN),
                                            static_cast<int16_t>(INT16_MAX),
                                            experiment.GetModuleFastDirectionStep(module_number_abs),
                                            experiment.GetPixel0OfModule(module_number_abs));
    } 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, 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;

        if (binning_2x2)
            output = (int16_t *) full_image_buffer.data();
        else
            output = (int16_t *) precompression_buffer.data();

        if (experiment.GetDetectorMode() != DetectorMode::Conversion)
            memcpy(output + RAW_MODULE_SIZE * module_number_abs, input, RAW_MODULE_SIZE * experiment.GetPixelDepth());
        else
            conv.ConvertAdjustGeom((int16_t *) output, (uint16_t *) input,
                                   experiment.GetModuleSlowDirectionStep(module_number_abs),
                                   experiment.GetModuleFastDirectionStep(module_number_abs),
                                   experiment.GetPixel0OfModule(module_number_abs));
    } else {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Summation with CPU conversion not supported at the moment");
    }
}