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

#include "TestImagePusher.h"
#include "../tests/CheckImageOutput.h"
#include "JFJochCompressor.h"
#include "JFJochDecompress.h"
#include "../frame_serialize/CBORStream2Deserializer.h"

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

void TestImagePusher::StartDataCollection(StartMessage& message) {
    std::unique_lock ul(m);

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

bool TestImagePusher::SendCalibration(const CompressedImage &message) {
    std::unique_lock ul(m);

    if (!is_running)
        correct_sequence = false;
    return true;
}

bool TestImagePusher::EndDataCollection(const EndMessage& message) {
    std::unique_lock ul(m);

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

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

    frame_counter++;
    if (image_number == image_id) {
        auto deserialized = CBORStream2Deserialize(image_data, image_size);
        if (deserialized->data_message) {
            receiver_generated_image.resize(deserialized->data_message->image.GetCompressedSize());
            memcpy(receiver_generated_image.data(),
                   deserialized->data_message->image.GetCompressed(),
                   deserialized->data_message->image.GetCompressedSize());
        }
    }
    return true;
}

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

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

size_t TestImagePusher::GetCounter() const {
    std::unique_lock 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<uint8_t> decompressed_image_8;
        std::vector<uint16_t> decompressed_image_16;
        std::vector<uint32_t> decompressed_image_32;

        // Image decompression
        try {
            switch (x.GetByteDepthImage()) {
                case 1:
                    JFJochDecompress(decompressed_image_8, x.GetCompressionAlgorithm(),
                                     receiver_generated_image, x.GetPixelsNum());
                    break;
                case 2:
                    JFJochDecompress(decompressed_image_16, x.GetCompressionAlgorithm(),
                                     receiver_generated_image, x.GetPixelsNum());
                    break;
                case 4:
                    JFJochDecompress(decompressed_image_32, x.GetCompressionAlgorithm(),
                                     receiver_generated_image, x.GetPixelsNum());
                    break;
                default:
                    throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Pixel depth unsupported");
            }

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

        bool check_image = x.GetDetectorSetup().GetDetectorType() != DetectorType::EIGER
            || x.GetBitDepthReadout() == 16;

        if (no_errors && check_image) {
            // Check output

            size_t storage_cell = 0;
            if (x.GetStorageCellNumber() > 1)
                storage_cell = image_id % x.GetStorageCellNumber();

            double result = 0;
            if (x.GetByteDepthImage() == 2) {
                if (x.IsPixelSigned())
                    result = CheckImageOutput(x, calibration,
                                              raw_reference_image.data(),
                                              (int16_t *) decompressed_image_16.data(),
                                              storage_cell);
                else
                    result = CheckImageOutput(x, calibration,
                                              raw_reference_image.data(),
                                              decompressed_image_16.data(),
                                              storage_cell);
            } else if (x.GetByteDepthImage() == 4) {
                if (x.IsPixelSigned())
                    result = CheckImageOutput(x, calibration,
                                              raw_reference_image.data(),
                                              (int32_t *) decompressed_image_32.data(),
                                              storage_cell);
                else
                    result = CheckImageOutput(x, calibration,
                                              raw_reference_image.data(),
                                              decompressed_image_32.data(),
                                              storage_cell);
            } else if (x.GetByteDepthImage() == 1) {
                if (x.IsPixelSigned())
                    result = CheckImageOutput(x, calibration,
                                              raw_reference_image.data(),
                                              (int8_t *) decompressed_image_8.data(),
                                              storage_cell);
                else
                    result = CheckImageOutput(x, calibration,
                                              raw_reference_image.data(),
                                              decompressed_image_8.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);
        }
    }
    if (no_errors && data_message && data_message->bkg_estimate.has_value())
        logger.Info("Background estimate {} Spot count {}",
            data_message->bkg_estimate.value(),
            data_message->spots.size());
    return no_errors;
}

std::string TestImagePusher::PrintSetup() const {
    return "TestImagePusher: One image is validated against expected data, nothing is saved";
}