Jungfraujoch/acquisition_device/AcquisitionDevice.cpp

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

#include <thread>
#include <fstream>
#include <cmath>

#include "../common/JFJochException.h"
#include "AcquisitionDevice.h"
#include "../common/NetworkAddressConvert.h"

AcquisitionDevice::AcquisitionDevice(uint16_t in_data_stream) {
    logger = nullptr;
    data_stream = in_data_stream;
}

uint64_t AcquisitionDevice::GeneratorModuleFrames(uint32_t hbm_size_bytes) {
    return 2 * static_cast<uint64_t>(hbm_size_bytes) / (RAW_MODULE_SIZE * sizeof(uint16_t));
}

// frame_generator.cpp refuses a run that fills its memory, and load_calibration.cpp lets one module frame past it
// through, so the host keeps the last one free.
void AcquisitionDevice::CheckInternalGeneratorFrame(size_t module_frame) const {
    if (module_frame + 1 >= internal_generator_module_frames)
        throw JFJochException(JFJochExceptionCategory::InputParameterAboveMax,
                              "Internal frame generator: the card memory holds "
                              + std::to_string(internal_generator_module_frames - 1)
                              + " module frames (modules x images), this needs "
                              + std::to_string(module_frame + 1));
}

void AcquisitionDevice::PrepareAction(const DiffractionExperiment &experiment) {
    if (experiment.GetModulesNum(data_stream) > max_modules)
        throw(JFJochException(JFJochExceptionCategory::InputParameterAboveMax,
                              "Number of modules exceeds max possible for FPGA"));

    if (experiment.IsUsingInternalPacketGen()) {
        int64_t module_frames = experiment.GetModulesNum(data_stream) * experiment.GetInternalPacketGeneratorImages();
        // A 32-bit EIGER frame takes two
        if ((experiment.GetDetectorType() == DetectorType::EIGER) && (experiment.GetBitDepthReadout() == 32))
            module_frames *= 2;
        CheckInternalGeneratorFrame(module_frames - 1);
    }

    counters.Reset(experiment, data_stream);
}

void AcquisitionDevice::StartAction(const DiffractionExperiment &experiment, uint32_t optional_flags) {
    Cancel();

    if (experiment.GetModulesNum(data_stream) > max_modules)
        throw(JFJochException(JFJochExceptionCategory::InputParameterAboveMax,
                              "Number of modules exceeds max possible for FPGA"));

    counters.Reset(experiment, data_stream);
    expected_frames = experiment.GetFrameNum() / experiment.GetFPGASummation();

    // Ensure internal WR queue is empty
    work_request_queue.Clear();

    Start(experiment, optional_flags);

    for (uint32_t i = 0; i < buffer_device.size(); i++)
        SendWorkRequest(i);

    auto c = work_completion_queue.GetBlocking();
    if (c.type != Completion::Type::Start)
        throw JFJochException(JFJochExceptionCategory::AcquisitionDeviceError, "Mismatch in work completions");

    StartSendingWorkRequests();

    start_time = std::chrono::system_clock::now();

    if (experiment.IsUsingInternalPacketGen())
        RunInternalGenerator(experiment);
}

void AcquisitionDevice::WaitForActionComplete() {
    auto c = work_completion_queue.GetBlocking();

    while (c.type != Completion::Type::End) {
        DeviceOutput* output;

        bool skip_analysis = false;

        try {
            output = GetDeviceOutput(c.handle);
        } catch (const JFJochException &e) {
            if (logger)
                logger->ErrorException(e);
            skip_analysis = true;
        }

        if (!skip_analysis) {
            c.module_number = output->module_statistics.module_number;
            c.packet_count = output->module_statistics.packet_count;
            c.frame_number = output->module_statistics.frame_number;

            if (c.frame_number >= expected_frames) {
                Cancel();
                // this frame is not of any interest, therefore its location can be immediately released
                SendWorkRequest(c.handle);
            } else if (c.module_number >= max_modules) {
                // Module number out of bounds, don't process
                if (logger != nullptr)
                    logger->Error("Completion with wrong module number data stream {} completion frame number {} module {} handle {}",
                                  data_stream, c.frame_number, c.module_number, c.handle);
                SendWorkRequest(c.handle);
            } else if (c.frame_number < counters.GetSlowestFrameNumber()) {
                // Module is falling behind, needs to return the handle then
                SendWorkRequest(c.handle);
            } else {
                try {
                    counters.UpdateCounters(&c);
                } catch (const JFJochException &e) {
                    if (logger)
                        logger->ErrorException(e);
                    SendWorkRequest(c.handle);
                }
            }
            if (logger != nullptr)
                logger->Debug("Data stream {} completion frame number {} module {} handle {}",
                              data_stream, c.frame_number, c.module_number, c.handle);
        }
        c = work_completion_queue.GetBlocking();
    }
    counters.SetAcquisitionFinished();

    end_time = std::chrono::system_clock::now();
    Cancel();
    Finalize();
}

void AcquisitionDevice::SendWorkRequest(uint32_t handle) {
    work_request_queue.Put(WorkRequest{
        .handle = handle
    });
}

uint64_t AcquisitionDevice::GetBytesReceived() const {
    return counters.GetBytesReceived();
}

const DeviceOutput *AcquisitionDevice::GetDeviceOutput(size_t frame_number, uint16_t module_number) const {
    auto handle = counters.GetBufferHandle(frame_number, module_number);
    if (handle != HandleNotValid)
        return GetDeviceOutput(handle);
    else
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Frame not collected");
}

const DeviceOutput *AcquisitionDevice::GetDeviceOutput(size_t handle) const {
    if (handle >= buffer_device.size())
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Handle outside of range");
    else
        return (DeviceOutput *) buffer_device.at(handle);
}

DeviceOutput *AcquisitionDevice::GetDeviceOutput(size_t handle) {
    if (handle >= buffer_device.size())
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Handle outside of range");
    else
        return (DeviceOutput *) buffer_device.at(handle);
}

void AcquisitionDevice::InitializeCalibration(const DiffractionExperiment &experiment, const JFCalibration &calib) {}

void AcquisitionDevice::InitializeIntegrationMap(const DiffractionExperiment &experiment,
                                                 const std::vector<uint16_t> &v,
                                                 const std::vector<float> &weights) {}

void AcquisitionDevice::InitializeIntegrationMap(const uint16_t *map, const float *weights, size_t module_number) {}

void AcquisitionDevice::InitializeSpotFinderResolutionMap(const float *data, size_t module_number) {}

void AcquisitionDevice::InitializeROIMap(const uint16_t *map, size_t module_number) {}

void AcquisitionDevice::InitializePixelMask(const uint32_t *module_mask, size_t module_number) {}

void AcquisitionDevice::InitializeROIMap(const DiffractionExperiment& experiment, const std::vector<uint16_t>& roi_map) {
    if (roi_map.size() != experiment.GetXPixelsNumConv() * experiment.GetYPixelsNumConv())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in array size");

    std::vector<uint16_t> tmp(RAW_MODULE_SIZE);
    auto offset = experiment.GetFirstModuleOfDataStream(data_stream);
    size_t modules = experiment.GetModulesNum(data_stream);
    for (int m = 0; m < modules; m++) {
        ConvertedToRawGeometry(experiment, offset + m, tmp.data(), roi_map.data());
        InitializeROIMap(tmp.data(), m);
    }
}

void AcquisitionDevice::InitializeEmptyPixelMask(const DiffractionExperiment &experiment) {
    std::vector<uint32_t> empty_mask(RAW_MODULE_SIZE);

    size_t modules = experiment.GetModulesNum(data_stream);
    for (int m = 0; m < modules; m++)
        InitializePixelMask(empty_mask.data(), m);
}


void AcquisitionDevice::InitializeDataProcessing(const DiffractionExperiment &experiment,
                                                 const AzimuthalIntegrationMapping &azint) {
    auto offset = experiment.GetFirstModuleOfDataStream(data_stream);
    size_t modules = experiment.GetModulesNum(data_stream);

    // When azimuthal integration is forced onto the CPU, the FPGA must not bin pixels
    // (the integration map can address more bins than the FPGA supports); the CPU path
    // computes the profile from the assembled image instead.
    const bool load_integration_map = !experiment.GetAzimuthalIntegrationSettings().IsForceCPUinFPGAWorkflow();

    if (experiment.IsGeometryTransformed()) {
        std::vector<float> tmp1(RAW_MODULE_SIZE);
        std::vector<uint16_t> tmp2(RAW_MODULE_SIZE);

        for (int m = 0; m < modules; m++) {
            if (load_integration_map) {
                ConvertedToRawGeometry(experiment, offset + m, tmp1.data(), azint.Corrections().data());
                ConvertedToRawGeometry(experiment, offset + m, tmp2.data(), azint.GetPixelToBin().data());
                InitializeIntegrationMap(tmp2.data(), tmp1.data(), m);
            }

            ConvertedToRawGeometry(experiment, offset + m, tmp1.data(), azint.Resolution().data());
            InitializeSpotFinderResolutionMap(tmp1.data(), m);
        }
    } else {
        for (int m = 0; m < modules; m++) {
            if (load_integration_map)
                InitializeIntegrationMap(azint.GetPixelToBin().data() + (offset + m) * RAW_MODULE_SIZE,
                                         azint.Corrections().data() + (offset + m) * RAW_MODULE_SIZE,
                                         m);
            InitializeSpotFinderResolutionMap(azint.Resolution().data() + (m + offset) * RAW_MODULE_SIZE,
                                              m);
        }
    }
}

void AcquisitionDevice::InitializePixelMask(const DiffractionExperiment &experiment, const PixelMask &mask) {
    auto offset = experiment.GetFirstModuleOfDataStream(data_stream);
    size_t modules = experiment.GetModulesNum(data_stream);

    std::vector<uint32_t> tmp(RAW_MODULE_SIZE);
    for (int m = 0; m < modules; m++) {
        ConvertedToRawGeometry(experiment, offset + m, tmp.data(), mask.GetMask().data());
        InitializePixelMask(tmp.data(), m);
    }
}

void AcquisitionDevice::FrameBufferRelease(size_t frame_number, uint16_t module_number) {
    auto handle = counters.GetBufferHandleAndClear(frame_number, module_number);
    if (handle != AcquisitionCounters::HandleNotFound)
        SendWorkRequest(handle);
}

void AcquisitionDevice::EnableLogging(Logger *in_logger) {
    logger = in_logger;
}

int32_t AcquisitionDevice::GetNUMANode() const {
    return -1;
}

uint16_t AcquisitionDevice::GetUDPPort() const {
    return 1234;
}

const AcquisitionCounters &AcquisitionDevice::Counters() const {
    return counters;
}

std::string AcquisitionDevice::GetIPv4Address() const {
    return IPv4AddressToStr(ipv4_addr);
}

std::string AcquisitionDevice::GetMACAddress() const {
    return MacAddressToStr(mac_addr);
}

DataCollectionStatus AcquisitionDevice::GetDataCollectionStatus() const {
    return {};
}

DeviceStatus AcquisitionDevice::GetDeviceStatus() const {
    return {};
}

AcquisitionDeviceStatistics AcquisitionDevice::GetStatistics() const {
    AcquisitionDeviceStatistics ret{};
    ret.bytes_received = GetBytesReceived();
    ret.start_timestamp = std::chrono::system_clock::to_time_t(start_time);
    ret.end_timestamp = std::chrono::system_clock::to_time_t(end_time);
    ret.packets_expected = counters.GetTotalExpectedPackets();
    ret.good_packets = counters.GetTotalPackets();

    for (int i = 0; i < counters.GetModuleNumber(); i++)
        ret.packets_received_per_module.push_back(counters.GetTotalPackets(i));

    if ((ret.packets_expected == 0) || (ret.good_packets == ret.packets_expected))
        ret.efficiency = 1.0;
    else
        ret.efficiency = static_cast<float>(ret.good_packets) / static_cast<float>(ret.packets_expected);

    return ret;
}

void AcquisitionDevice::SetIPv4Address(uint32_t ipv4_addr_network_order) {
    ipv4_addr = ipv4_addr_network_order;
}

AcquisitionDeviceNetConfig AcquisitionDevice::GetNetConfig() const {
    return {
        .mac_addr = GetMACAddress(),
        .ipv4_addr = GetIPv4Address(),
        .udp_port = GetUDPPort()
    };
}

void AcquisitionDevice::RunInternalGenerator(const DiffractionExperiment &experiment) {
    FrameGeneratorConfig config{};
    config.frames = experiment.GetFrameNum() + DELAY_FRAMES_STOP_AND_QUIT + 1;
    config.modules = experiment.GetModulesNum(data_stream);
    config.data_stream = data_stream;
    config.pulse_id = INT_PKT_GEN_BUNCHID;
    config.exptime = INT_PKT_GEN_EXPTTIME;
    config.debug = INT_PKT_GEN_DEBUG;
    config.dest_mac_addr = MacAddressFromStr(GetMACAddress());
    config.dest_ipv4_addr = IPv4AddressFromStr(GetIPv4Address());
    config.images_in_memory = experiment.GetInternalPacketGeneratorImages() - 1;
    switch (experiment.GetDetectorSetup().GetDetectorType()) {
        case DetectorType::JUNGFRAU:
            config.detector_type = SLS_DETECTOR_TYPE_JUNGFRAU;
            break;
        case DetectorType::EIGER:
            config.detector_type = SLS_DETECTOR_TYPE_EIGER;
            config.eiger_bit_depth = experiment.GetBitDepthReadout();
            break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Detector not supported");
    }
    HW_RunInternalGenerator(config);
}

void AcquisitionDevice::SetSpotFinderParameters(const SpotFindingSettings &settings) {
    SpotFinderParameters fpga_parameters{};

    // The FPGA compares d against these as xray_d_t = ap_ufixed<16,9> (hls_jfjoch.h), which tops out just
    // below 512 A and WRAPS above it - so "no low-resolution limit" is that ceiling, not infinity. Sending
    // a larger number would wrap to a small d and mask the whole image. The high end needs no such care:
    // no pixel has d < 0.
    constexpr float FPGA_MAX_D_A = 511.0f;
    fpga_parameters.snr_threshold = settings.signal_to_noise_threshold;
    fpga_parameters.count_threshold = settings.photon_count_threshold;
    fpga_parameters.max_d = settings.low_resolution_limit.value_or(FPGA_MAX_D_A);
    fpga_parameters.min_d = settings.high_resolution_limit.value_or(0.0f);
    fpga_parameters.min_pix_per_spot = settings.min_pix_per_spot.value_or(2);
    HW_SetSpotFinderParameters(fpga_parameters);
}