// Copyright (2019-2023) Paul Scherrer Institute

#include "JFJochReceiverTest.h"
#include "JFJochReceiverService.h"
#include "../common/ZMQImagePusher.h"
#include "../common/TestImagePusher.h"

#define STORAGE_CELL_FOR_TEST 11

JFJochReceiverOutput RunJFJochReceiverTest(AcquisitionDeviceGroup &aq_devices, ImagePusher &pusher,
                                           const DiffractionExperiment &x,
                                           Logger &logger, JFCalibration &calib,
                                           uint16_t nthreads, bool abort,
                                           ZMQPreviewPublisher *in_preview_writer,
                                           const std::string &numa_policy) {
    ZMQContext context;
    context.NumThreads(4);

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.PreviewPublisher(in_preview_writer).NumThreads(nthreads).NUMAPolicy(numa_policy);

    DataProcessingSettings settings;
    settings.signal_to_noise_threshold = 2.5;
    settings.photon_count_threshold = 5;
    settings.min_pix_per_spot = 3;
    settings.max_pix_per_spot = 200;
    settings.local_bkg_size = 5;
    service.SetDataProcessingSettings(settings);

    service.Start(x, &calib);

    if (x.GetImageNum() > 0) {
        // Initial communication
        if (abort) {
            std::this_thread::sleep_for(std::chrono::seconds(60));
            service.Cancel();
        }
    }

    return service.Stop();
}

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;
                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,
                        uint16_t nthreads, bool abort,
                        ZMQPreviewPublisher *in_preview_writer,
                        const std::string &numa_policy) {

    std::vector<uint16_t> raw_expected_image(RAW_MODULE_SIZE * x.GetModulesNum());
    for (int i = 0; i < x.GetDataStreamsNum(); i++) {
        uint32_t module0 = x.GetFirstModuleOfDataStream(i);
        auto int_gen_frame = aq_devices[i].GetInternalGeneratorFrame();
        if (int_gen_frame.size() < x.GetModulesNum(i) * RAW_MODULE_SIZE) {
            logger.Error("Wrong internal generator frame size");
            return false;
        }

        for (int m = 0; m < x.GetModulesNum(i); m++) {
                memcpy(raw_expected_image.data() + (module0 + m) * RAW_MODULE_SIZE,
                       int_gen_frame.data() + m * RAW_MODULE_SIZE,
                       RAW_MODULE_SIZE * sizeof(uint16_t));
        }
    }

    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);

    output = RunJFJochReceiverTest(aq_devices, pusher, x, logger, calib,  nthreads, abort,
                                   in_preview_writer, numa_policy);

    bool no_errors = true;

    if ((x.GetImageNum() > 0) && (!abort)) {
        no_errors = pusher.CheckImage(x, raw_expected_image, 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;
}