Jungfraujoch/reader/JFJochReader.cpp

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

#include "JFJochReader.h"
#include <future>

JFJochReader &JFJochReader::Experiment(const DiffractionExperiment &experiment) {
std::unique_lock ul(m);
    default_experiment = experiment;
    return *this;
}

void JFJochReader::SummationThread(int64_t image0, int64_t n_image, int64_t image_jump, JFJochReaderImage &image) {
    std::vector<uint8_t> buffer;
    DataMessage msg;
    for (int64_t i = image0; i < n_image; i += image_jump) {
        bool ret = LoadImage_i(dataset, msg, buffer, i, false);
        if (ret) {
            auto image_sum = std::make_shared<JFJochReaderImage>(msg, dataset);
            {
                std::unique_lock ul(summation_mutex);
                image.AddImage(*image_sum);
            }
        }
    }
}

std::shared_ptr<JFJochReaderImage> JFJochReader::LoadImage(int64_t image_number, int64_t summation_factor) {
    // It would be a mess to load two images at the same time
    // so loading is protected via mutex
    // yet copying share_ptr pointer is atomic and needs no mutex protection
    std::unique_lock ul(m);

    std::vector<uint8_t> buffer;
    DataMessage msg;

    if (LoadImage_i(dataset, msg, buffer, image_number, true)) {
        auto image = std::make_shared<JFJochReaderImage>(msg, dataset);
        if (summation_factor > 4) {
            int64_t nthread = std::min<int64_t>(summation_factor - 1, 8);

            std::vector<std::future<void>> futures;
            for (int i = 0; i < nthread; i++)
                futures.emplace_back(std::async(std::launch::async,
                                                &JFJochReader::SummationThread, this,
                                                image_number + 1 + i,
                                                image_number + summation_factor,
                                                nthread,
                                                std::ref(*image)));
            for (auto &f: futures)
                f.get();
        } else if (summation_factor > 1) {
            SummationThread(image_number + 1, image_number + summation_factor, 1, *image);
        }
        return image;
    }
    return {};
}

std::shared_ptr<JFJochReaderRawImage> JFJochReader::GetRawImage(int64_t image_number) {
    auto ret = std::make_shared<JFJochReaderRawImage>();
    if (!ReadRawImage(image_number, *ret))
        return {};
    return ret;
}

void JFJochReader::SetStartMessage(const std::shared_ptr<JFJochReaderDataset> &val) {
    std::unique_lock ul(m);
    dataset = val;
}

std::shared_ptr<const JFJochReaderDataset> JFJochReader::GetDataset() const {
    std::unique_lock ul(m);
    return dataset;
}

void JFJochReader::UpdateGeomMetadata(const DiffractionExperiment &experiment) {
    std::unique_lock ul(m);
    if (!dataset)
        return;

    auto new_dataset = std::make_shared<JFJochReaderDataset>(*dataset);
    // At the moment subset of options is limited to safe ones...need to change it in the future

    new_dataset->experiment.BeamX_pxl(experiment.GetBeamX_pxl());
    new_dataset->experiment.BeamY_pxl(experiment.GetBeamY_pxl());
    new_dataset->experiment.DetectorDistance_mm(experiment.GetDetectorDistance_mm());
    new_dataset->experiment.IncidentEnergy_keV(experiment.GetIncidentEnergy_keV());
    new_dataset->experiment.PoniRot1_rad(experiment.GetDatasetSettings().GetPoniRot1_rad());
    new_dataset->experiment.PoniRot2_rad(experiment.GetDatasetSettings().GetPoniRot2_rad());
    new_dataset->experiment.PoniRot3_rad(experiment.GetDatasetSettings().GetPoniRot3_rad());
    new_dataset->experiment.SetUnitCell(experiment.GetUnitCell());
    new_dataset->experiment.SetSpaceGroup(experiment.GetGemmiSpaceGroup());
    new_dataset->experiment.PolarizationFactor(experiment.GetPolarizationFactor());
    new_dataset->experiment.Goniometer(experiment.GetGoniometer());
    new_dataset->experiment.GridScan(experiment.GetGridScan());
    new_dataset->experiment.ImportIndexingSettings(experiment.GetIndexingSettings());
    new_dataset->experiment.ImportBraggIntegrationSettings(experiment.GetBraggIntegrationSettings());
    new_dataset->experiment.DetectIceRings(experiment.IsDetectIceRings());
    dataset = new_dataset;
}

void JFJochReader::UpdateUserMask(const std::vector<uint32_t> &mask) {
    std::unique_lock ul(m);
    if (!dataset)
        return;
    auto new_dataset = std::make_shared<JFJochReaderDataset>(*dataset);
    // Copy-on-write: the mask is shared with the old snapshot, so edit a fresh copy, not in place.
    auto new_mask = std::make_shared<PixelMask>(*dataset->pixel_mask);
    new_mask->LoadUserMask(dataset->experiment, mask);
    new_dataset->pixel_mask = new_mask;
    dataset = new_dataset;
}

void JFJochReader::AddDeadPixelMask(const std::vector<uint32_t> &mask) {
    std::unique_lock ul(m);
    if (!dataset)
        return;
    auto new_dataset = std::make_shared<JFJochReaderDataset>(*dataset);
    auto new_mask = std::make_shared<PixelMask>(*dataset->pixel_mask);
    new_mask->AddDeadPixelMask(dataset->experiment, mask);
    new_dataset->pixel_mask = new_mask;
    dataset = new_dataset;
}

std::shared_ptr<JFJochReaderSpots> JFJochReader::ReadAllSpots(int64_t start_image, int64_t end_image,
    int64_t stride) const {

    if (start_image < 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
            "Start image must be non-negative");

    if (start_image > end_image)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
            "Start image number is greater than end image number");

    if (stride == 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
            "Stride cannot be zero");

    size_t nelems = (end_image - start_image) / stride + 1;

    auto ret = std::make_shared<JFJochReaderSpots>();
    ret->start_image = static_cast<int64_t>(start_image);
    ret->stride = static_cast<int64_t>(stride);
    ret->spots.reserve(nelems);

    for (int i = 0; i < nelems; i++)
        ret->spots.emplace_back(ReadSpots(start_image + i * stride));
    return ret;
}