// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <cstring>

#include "bitshuffle/bitshuffle.h"

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

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

FrameTransformation::FrameTransformation(const DiffractionExperiment &in_experiment) :
        experiment(in_experiment),
        image_mode(in_experiment.GetImageMode()),
        pixel_depth(in_experiment.GetByteDepthImage()),
        compressor(in_experiment.GetCompressionAlgorithm()),
        err_value(experiment.GetByteDepthImage() * RAW_MODULE_SIZE) {

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

    switch (image_mode) {
        case CompressedImageMode::Int8:
            FillVector<int8_t>(precompression_buffer, experiment.GetImageFillValue());
            FillVector<int8_t>(err_value, experiment.GetImageFillValue());
            break;
        case CompressedImageMode::Int16:
            FillVector<int16_t>(precompression_buffer, experiment.GetImageFillValue());
            FillVector<int16_t>(err_value, experiment.GetImageFillValue());
            break;
        case CompressedImageMode::Int32:
            FillVector<int32_t>(precompression_buffer, experiment.GetImageFillValue());
            FillVector<int32_t>(err_value, experiment.GetImageFillValue());
            break;
        case CompressedImageMode::Uint8:
            FillVector<uint8_t>(precompression_buffer, experiment.GetImageFillValue());
            FillVector<uint8_t>(err_value, experiment.GetImageFillValue());
            break;
        case CompressedImageMode::Uint16:
            FillVector<uint16_t>(precompression_buffer, experiment.GetImageFillValue());
            FillVector<uint16_t>(err_value, experiment.GetImageFillValue());
            break;
        case CompressedImageMode::Uint32:
            FillVector<uint32_t>(precompression_buffer, experiment.GetImageFillValue());
            FillVector<uint32_t>(err_value, experiment.GetImageFillValue());
            break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Image mode not supprted for FrameTransformation");
    }
}

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

CompressedImage FrameTransformation::GetCompressedImage() {
    const uint8_t *data;
    size_t size;
    if (experiment.GetCompressionAlgorithm() == CompressionAlgorithm::NO_COMPRESSION) {
        data = (uint8_t *) precompression_buffer.data();
        size = experiment.GetPixelsNum() * experiment.GetByteDepthImage();
    } else {
        compressed_buffer.resize(MaxCompressedSize(experiment.GetCompressionAlgorithm(),
                                                   experiment.GetPixelsNum(),
                                                   experiment.GetByteDepthImage()));
        data = (uint8_t *) compressed_buffer.data();
        size = CompressImage(compressed_buffer.data());
    }

    return CompressedImage(data, size,experiment.GetXPixelsNum(),
                           experiment.GetYPixelsNum(), image_mode,
                           experiment.GetCompressionAlgorithm());
}

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.IsGeometryTransformed()) {
        size_t mod_size = RAW_MODULE_SIZE * pixel_depth;
        memcpy(precompression_buffer.data() + module_number_abs * mod_size, input, mod_size);
    } else {
        auto pixel0 = experiment.GetPixel0OfModuleConv(module_number_abs);
        auto fast_step = experiment.GetModuleFastDirectionStep(module_number_abs);
        auto slow_step = experiment.GetModuleSlowDirectionStep(module_number_abs);

        switch (image_mode) {
            case CompressedImageMode::Int8:
                TransferModuleAdjustMultipixels((int8_t *) precompression_buffer.data(), (int8_t *) input,
                                                slow_step, static_cast<int8_t>(INT8_MIN), static_cast<int8_t>(INT8_MAX),
                                                fast_step, pixel0);
                break;
            case CompressedImageMode::Int16:
                TransferModuleAdjustMultipixels((int16_t *) precompression_buffer.data(), (int16_t *) input,
                                                slow_step, static_cast<int16_t>(INT16_MIN), static_cast<int16_t>(INT16_MAX),
                                                fast_step, pixel0);
                break;
            case CompressedImageMode::Int32:
                TransferModuleAdjustMultipixels((int32_t *) precompression_buffer.data(), (int32_t *) input,
                                                slow_step, static_cast<int32_t>(INT32_MIN), static_cast<int32_t>(INT32_MAX),
                                                fast_step, pixel0);

                break;
            case CompressedImageMode::Uint8:
                TransferModuleAdjustMultipixels((uint8_t *) precompression_buffer.data(), (uint8_t *) input,
                                                slow_step, static_cast<uint8_t>(0), static_cast<uint8_t>(UINT8_MAX - 1),
                                                fast_step, pixel0);
                break;
            case CompressedImageMode::Uint16:
                TransferModuleAdjustMultipixels((uint16_t *) precompression_buffer.data(), (uint16_t *) input,
                                                slow_step, static_cast<uint16_t>(0), static_cast<uint16_t>(UINT16_MAX - 1),
                                                fast_step, pixel0);
                break;
            case CompressedImageMode::Uint32:
                TransferModuleAdjustMultipixels((uint32_t *) precompression_buffer.data(), (uint32_t *) input,
                                                slow_step, static_cast<uint32_t>(0), static_cast<uint32_t>(UINT32_MAX - 1),
                                                fast_step, pixel0);
                break;
            default:
                throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Image mode not supprted for FrameTransformation");
        }
    }
}

void FrameTransformation::ProcessModule(const DeviceOutput *output, int data_stream) {
    ProcessModule(output->pixels, output->module_statistics.module_number, data_stream);
}

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

void FrameTransformation::FillNotCollectedModule(uint16_t module_number, int data_stream) {
    ProcessModule(err_value.data(), module_number, data_stream);
}