// Copyright (2019-2023) Paul Scherrer Institute

#include "TestImagePusher.h"
#include "../tests/FPGAUnitTest.h"
#include "JFJochCompressor.h"
#include "../compression/JFJochDecompress.h"
#include "../frame_serialize/JFJochFrameDeserializer.h"

TestImagePusher::TestImagePusher(int64_t image_number) {
    image_id = image_number;
}

void TestImagePusher::StartDataCollection(const uint8_t *image_data, size_t image_size, int64_t data_file_count) {
    std::unique_lock<std::mutex> ul(m);

    if (is_running)
        correct_sequence = false;
    else
        is_running = true;
}

void TestImagePusher::EndDataCollection(const EndMessage& message) {
    std::unique_lock<std::mutex> ul(m);

    if (!is_running)
        correct_sequence = false;
    else
        is_running = false;
}

void TestImagePusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) {
    std::unique_lock<std::mutex> ul(m);

    frame_counter++;
    if (image_number == image_id) {

        JFJochFrameDeserializer deserializer;
        deserializer.Process(image_data, image_size);

        auto image_array = deserializer.GetDataMessage();
        receiver_generated_image.resize(image_array.image.size);
        memcpy(receiver_generated_image.data(), image_array.image.data, image_array.image.size);
    }
}

void TestImagePusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number,
                                ZeroCopyReturnValue *z) {
    std::unique_lock<std::mutex> ul(m);

    frame_counter++;
    if (image_number == image_id) {
        JFJochFrameDeserializer deserializer;
        deserializer.Process(image_data, image_size);

        auto image_array = deserializer.GetDataMessage();
        receiver_generated_image.resize(image_array.image.size);
        memcpy(receiver_generated_image.data(), image_array.image.data, image_array.image.size);
    }
    z->release();
}

bool TestImagePusher::CheckSequence() const {
    std::unique_lock<std::mutex> ul(m);
    return correct_sequence;
}

const std::vector<uint8_t> &TestImagePusher::GetImage() const {
    std::unique_lock<std::mutex> ul(m);
    return receiver_generated_image;
}

size_t TestImagePusher::GetCounter() const {
    std::unique_lock<std::mutex> ul(m);
    return frame_counter;
}

bool TestImagePusher::CheckImage(const DiffractionExperiment &x, const std::vector<uint16_t> &raw_reference_image,
                                 const JFCalibration &calibration,
                                 Logger &logger) {
    bool no_errors = true;

    if (receiver_generated_image.empty()) {
        logger.Error("Image empty");
        no_errors = false;
    } else {
        std::vector<uint16_t> decompressed_image_16;
        std::vector<uint32_t> decompressed_image_32;

        // Image decompression
        try {
            if (x.GetPixelDepth() == 2)
                JFJochDecompress(decompressed_image_16, x.GetCompressionAlgorithmEnum(),
                             receiver_generated_image, x.GetPixelsNum());
            else
                JFJochDecompress(decompressed_image_32, x.GetCompressionAlgorithmEnum(),
                                 receiver_generated_image, x.GetPixelsNum());

        } catch (const JFJochException &e) {
            logger.Error(e.what());
            no_errors = false;
        }
        if (no_errors) {
            // Check output

            if (x.GetDetectorMode() == DetectorMode::Conversion) {
                size_t storage_cell = 0;
                if (x.GetStorageCellNumber() > 1)
                    storage_cell = image_id % x.GetStorageCellNumber();

                double result = 0;
                if (x.GetPixelDepth() == 2) {
                    if (x.IsPixelSigned())
                        result = CheckConversionWithGeomTransform(x, calibration,
                                                              raw_reference_image.data(),
                                                              (int16_t *) decompressed_image_16.data(),
                                                              storage_cell);
                    else
                        result = CheckConversionWithGeomTransform(x, calibration,
                                                                  raw_reference_image.data(),
                                                                  (uint16_t *) decompressed_image_16.data(),
                                                                  storage_cell);
                } else if (x.GetPixelDepth() == 4) {
                    if (x.IsPixelSigned())
                        result = CheckConversionWithGeomTransform(x, calibration,
                                                              raw_reference_image.data(),
                                                              (int32_t *) decompressed_image_32.data(),
                                                              storage_cell);
                    else
                        result = CheckConversionWithGeomTransform(x, calibration,
                                                                  raw_reference_image.data(),
                                                                  (uint32_t *) decompressed_image_32.data(),
                                                                  storage_cell);
                }
                if (result > 0.5 * x.GetSummation()) {
                    logger.Error("Mean conversion error ({:.3f}) larger than threshold", result);
                    no_errors = false;
                } else
                    logger.Info("Mean conversion error: {:.3f}", result);
            } else if (x.GetDetectorMode() == DetectorMode::Raw) {
                size_t diff = 0;
                if (x.GetPixelDepth() == 2) {
                    for (int i = 0; i < x.GetPixelsNum(); i++) {
                        if (raw_reference_image[i] != decompressed_image_16[i])
                            diff++;
                    }
                } else {
                    for (int i = 0; i < x.GetPixelsNum(); i++) {
                        if (raw_reference_image[i] != decompressed_image_32[i])
                            diff++;
                    }
                }

                if (diff != 0) {
                    logger.Error("Raw data mismatch");
                    no_errors = false;
                }
            }
        }
    }
    return no_errors;
}