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

    if (!dataset)
        return {};

    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 {};
}

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);
    if (!dataset)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "File not loaded");
    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.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);
    new_dataset->pixel_mask.LoadUserMask(dataset->experiment, mask);
    dataset = new_dataset;
}