// Copyright (2019-2023) Paul Scherrer Institute

#include <cstring>

#include <bitshuffle/bitshuffle.h>

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

template <class T> void FillVector(std::vector<char> &v, T fill_value) {
    auto ptr = (T *) v.data();
    for (int i = 0; i < v.size() / sizeof(T); i++)
        ptr[i] = fill_value;
}

FrameTransformation::FrameTransformation(const DiffractionExperiment &in_experiment) :
        experiment(in_experiment),
        pixel_depth(experiment.GetPixelDepth()), pixel_signed(experiment.IsPixelSigned()),
        compressor(in_experiment.GetCompressionAlgorithmEnum()) {

    precompression_buffer.resize(experiment.GetPixelsNum() * pixel_depth);

    if (pixel_depth == 2) {
        if (pixel_signed)
            FillVector<int16_t>(precompression_buffer, INT16_MIN);
        else
            FillVector<uint16_t>(precompression_buffer, UINT16_MAX);
    } else if (pixel_depth == 4) {
        if (pixel_signed)
            FillVector<int32_t>(precompression_buffer, INT32_MIN);
        else
            FillVector<uint32_t>(precompression_buffer, UINT32_MAX);
    }
}

size_t FrameTransformation::SaveCompressedImage(void *output) {
    return compressor.Compress((char *) output, precompression_buffer.data(),
                               experiment.GetPixelsNum(), pixel_depth);
}

void FrameTransformation::ProcessModule(const void *input, uint16_t module_number, int data_stream) {
    size_t module_number_abs = experiment.GetFirstModuleOfDataStream(data_stream) + module_number;

    if (experiment.GetDetectorMode() != DetectorMode::Conversion) {
        size_t mod_size = RAW_MODULE_SIZE * pixel_depth;
        memcpy(precompression_buffer.data() + module_number_abs * mod_size, input, mod_size);
    } else {
        if (pixel_depth == 2) {
            if (pixel_signed)
                TransferModuleAdjustMultipixels((int16_t *) precompression_buffer.data(), (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
                TransferModuleAdjustMultipixels((uint16_t *) precompression_buffer.data(), (uint16_t *) input,
                                                experiment.GetModuleSlowDirectionStep(module_number_abs),
                                                static_cast<uint16_t>(0),
                                                static_cast<uint16_t>(UINT16_MAX - 1),
                                                experiment.GetModuleFastDirectionStep(module_number_abs),
                                                experiment.GetPixel0OfModule(module_number_abs));
        } else if (pixel_depth == 4) {
            if (pixel_signed)
                TransferModuleAdjustMultipixels((int32_t *) precompression_buffer.data(), (int32_t *) input,
                                                experiment.GetModuleSlowDirectionStep(module_number_abs),
                                                static_cast<int32_t>(INT32_MIN),
                                                static_cast<int32_t>(INT32_MAX),
                                                experiment.GetModuleFastDirectionStep(module_number_abs),
                                                experiment.GetPixel0OfModule(module_number_abs));
            else
                TransferModuleAdjustMultipixels((uint32_t *) precompression_buffer.data(), (uint32_t *) input,
                                                experiment.GetModuleSlowDirectionStep(module_number_abs),
                                                static_cast<uint32_t>(0),
                                                static_cast<uint32_t>(UINT32_MAX - 1),
                                                experiment.GetModuleFastDirectionStep(module_number_abs),
                                                experiment.GetPixel0OfModule(module_number_abs));
        }
    }
}

const void *FrameTransformation::GetPreviewImage() const {
    return precompression_buffer.data();
}