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

#include "JFJochReceiverTest.h"
#include "JFJochReceiverService.h"
#include "../image_pusher/ZMQStream2Pusher.h"
#include "../image_pusher/TestImagePusher.h"

#define STORAGE_CELL_FOR_TEST 11

static JFCalibration GeneratePedestalCalibration(const DiffractionExperiment &x) {
    JFCalibration ret(x);

    std::vector<float> pedestal_g0( RAW_MODULE_SIZE), pedestal_g1(RAW_MODULE_SIZE), pedestal_g2(RAW_MODULE_SIZE);

    for (int s = 0; s < x.GetStorageCellNumber(); s++) {
        for (int m = 0; m < x.GetModulesNum(); m++) {
            for (int i = 0; i < RAW_MODULE_SIZE; i++) {
                pedestal_g0[i] = 3000 + 100 * s + 10 * m + i % 50;
                if (x.IsFixedGainG1())
                    pedestal_g1[i] = 2500 + 100 * s + 10 * m + i % 50;
                else
                    pedestal_g1[i] = 15000 - 100 * s + 10 * m + i % 50;
                pedestal_g2[i] = 14000 - 100 * s + 10 * m + i % 50;

            }
            ret.Pedestal(m, 0, s).LoadPedestal(pedestal_g0);
            ret.Pedestal(m, 1, s).LoadPedestal(pedestal_g1);
            ret.Pedestal(m, 2, s).LoadPedestal(pedestal_g2);
        }
    }

    return ret;
}

bool JFJochReceiverTest(JFJochReceiverOutput &output, Logger &logger,
                        AcquisitionDeviceGroup &aq_devices,
                        const DiffractionExperiment &x,
                        const PixelMask &pixel_mask,
                        uint16_t nthreads,
                        const std::string &numa_policy,
                        size_t send_buf_size_MiB,
                        bool quick_integrate) {

    std::vector<uint16_t> raw_expected_image(RAW_MODULE_SIZE * x.GetModulesNum());

    for (int i = 0; i < raw_expected_image.size(); i++)
        raw_expected_image[i] = i % 65536;

    return JFJochReceiverTest(output, logger, aq_devices, x, pixel_mask, raw_expected_image, nthreads, numa_policy, send_buf_size_MiB,
                              quick_integrate);
}

bool JFJochReceiverTest(JFJochReceiverOutput &output, Logger &logger,
                        AcquisitionDeviceGroup &aq_devices,
                        const DiffractionExperiment &x,
                        const PixelMask &pixel_mask,
                        const std::vector<uint16_t> &raw_expected_image,
                        uint16_t nthreads,
                        const std::string &numa_policy,
                        size_t send_buf_size_MiB,
                        bool quick_integrate) {
    std::vector<uint16_t> raw_expected_image_with_mask = raw_expected_image;
    if (x.IsApplyPixelMask()) {
        for (int i = 0; i < raw_expected_image_with_mask.size(); i++)
            if (pixel_mask.GetMaskRaw((x)).at(i) != 0)
                raw_expected_image_with_mask[i] = UINT16_MAX;
    }

    JFCalibration calib = GeneratePedestalCalibration(x);

    int64_t image_number = x.GetImageNum() - 1;

    if (x.GetStorageCellNumber() > 1)
        image_number = STORAGE_CELL_FOR_TEST;
    if (image_number < 0)
        image_number = 0;

    TestImagePusher pusher(image_number);

    JFJochReceiverService service(aq_devices, logger, pusher, send_buf_size_MiB);
    service.NumThreads(nthreads).NUMAPolicy(numa_policy);
    service.LoadInternalGeneratorImage(x, raw_expected_image, 0);
    service.Indexing(x.GetIndexingSettings());

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    settings.signal_to_noise_threshold = 2.5;
    settings.photon_count_threshold = 5;
    settings.min_pix_per_spot = 1;
    settings.max_pix_per_spot = 200;
    settings.quick_integration = quick_integrate;

    service.SetSpotFindingSettings(settings);

    service.Start(x, pixel_mask, &calib);
    output = service.Stop();

    bool no_errors = true;

    if (x.GetImageNum() > 0) {
        no_errors = pusher.CheckImage(x, raw_expected_image_with_mask, calib, logger);

        if (output.efficiency != 1.0) {
            logger.Error("Not all frames were received");
            no_errors = false;
        }

        if (!pusher.CheckSequence()) {
            logger.Error("Wrong sequence of operations for pusher");
            no_errors = false;
        }
        if (pusher.GetCounter() != x.GetImageNum()) {
            logger.Error("Wrong frame number from pusher");
            no_errors = false;
        }
    }

    return no_errors;
}