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

#include "JFJochMultiReader.h"

#include <algorithm>
#include <cmath>

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

namespace {
    // How far two inputs may state the geometry apart and still be one run: what two collections of one
    // setup differ by in their headers, well under what a stills geometry refinement moves.
    constexpr double WAVELENGTH_TOLERANCE = 1e-3;     // relative
    constexpr double DISTANCE_TOLERANCE_MM = 0.1;
    constexpr double BEAM_TOLERANCE_PXL = 1.0;

    void CheckSameGeometry(const DiffractionExperiment &a, const DiffractionExperiment &b, size_t input) {
        auto refuse = [input](const std::string &what) {
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Input " + std::to_string(input + 1) + " is not at the geometry of input 1: "
                                  + what + " differs. Inputs read as one run must share one detector, "
                                  "wavelength, distance and beam centre; process them separately");
        };
        if (a.GetXPixelsNum() != b.GetXPixelsNum() || a.GetYPixelsNum() != b.GetYPixelsNum()
            || std::fabs(a.GetPixelSize_mm() - b.GetPixelSize_mm()) > 1e-6)
            refuse("the detector");
        if (std::fabs(a.GetWavelength_A() - b.GetWavelength_A()) > WAVELENGTH_TOLERANCE * a.GetWavelength_A())
            refuse("the wavelength");
        if (std::fabs(a.GetDetectorDistance_mm() - b.GetDetectorDistance_mm()) > DISTANCE_TOLERANCE_MM)
            refuse("the detector distance");
        if (std::fabs(a.GetBeamX_pxl() - b.GetBeamX_pxl()) > BEAM_TOLERANCE_PXL
            || std::fabs(a.GetBeamY_pxl() - b.GetBeamY_pxl()) > BEAM_TOLERANCE_PXL)
            refuse("the beam centre");
    }
}

JFJochMultiReader::JFJochMultiReader(std::vector<std::unique_ptr<JFJochReader>> in)
        : inputs(std::move(in)) {
    if (inputs.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No input to read");

    const auto first = inputs[0]->GetDataset();
    if (!first)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Input 1 holds no dataset");

    auto dataset = std::make_shared<JFJochReaderDataset>(*first);
    std::vector<uint32_t> mask = first->pixel_mask->GetMask();
    dataset->efficiency.clear();

    for (size_t k = 0; k < inputs.size(); k++) {
        const auto d = inputs[k]->GetDataset();
        if (!d)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Input " + std::to_string(k + 1) + " holds no dataset");
        // Rotation sweeps are not concatenated: their frames are one crystal turned, numbered by angle.
        if (const auto g = d->experiment.GetGoniometer(); g && g->GetIncrement_deg() != 0.0f)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Input " + std::to_string(k + 1) + " is a rotation sweep; only stills "
                                  "collections are read together");
        CheckSameGeometry(first->experiment, d->experiment, k);

        const auto &m = d->pixel_mask->GetMask();
        if (m.size() != mask.size())
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Input " + std::to_string(k + 1) + " has a pixel mask of another size");
        // A pixel is out where any input has it out.
        for (size_t i = 0; i < mask.size(); i++)
            mask[i] |= m[i];

        const auto n = inputs[k]->GetNumberOfImages();
        first_image.push_back(static_cast<int64_t>(number_of_images));
        // An input stating no efficiency has every frame at 1, as the image pass reads a missing one.
        auto efficiency = d->efficiency;
        efficiency.resize(n, 1.0f);
        dataset->efficiency.insert(dataset->efficiency.end(), efficiency.begin(), efficiency.end());
        number_of_images += n;
    }

    // Per-image results a previous run stored in the first input's file are not the whole run's.
    dataset->spot_count.clear();
    dataset->spot_count_indexed.clear();
    dataset->spot_count_low_res.clear();
    dataset->spot_count_ice_rings.clear();
    dataset->spot_count_ice_control.clear();
    dataset->indexing_result.clear();
    dataset->indexing_lattice_count.clear();
    dataset->bkg_estimate.clear();
    dataset->ice_ring_score.clear();
    dataset->resolution_estimate.clear();
    dataset->image_scale_factor.clear();
    dataset->image_scale_cc.clear();

    dataset->pixel_mask = std::make_shared<const PixelMask>(mask);
    SetStartMessage(dataset);
}

std::pair<size_t, int64_t> JFJochMultiReader::Locate(int64_t image_number) const {
    if (image_number < 0 || image_number >= static_cast<int64_t>(number_of_images))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Image number out of range");
    const auto it = std::upper_bound(first_image.begin(), first_image.end(), image_number);
    const auto k = static_cast<size_t>(it - first_image.begin()) - 1;
    return {k, image_number - first_image[k]};
}

bool JFJochMultiReader::LoadImage_i(std::shared_ptr<JFJochReaderDataset> &dataset, DataMessage &message,
                                    std::vector<uint8_t> &buffer, int64_t image_number, bool update_dataset) {
    const auto [k, local] = Locate(image_number);
    // The input decodes the frame against its own dataset; the whole run's stays as it is.
    auto input_dataset = std::const_pointer_cast<JFJochReaderDataset>(inputs[k]->GetDataset());
    if (!inputs[k]->LoadImage_i(input_dataset, message, buffer, local, false))
        return false;
    message.number = image_number;
    return true;
}

uint64_t JFJochMultiReader::GetNumberOfImages() const {
    return number_of_images;
}

void JFJochMultiReader::Close() {
    for (auto &i: inputs)
        i->Close();
}

bool JFJochMultiReader::ReadRawImage(int64_t image_number, JFJochReaderRawImage &image) {
    const auto [k, local] = Locate(image_number);
    return inputs[k]->ReadRawImage(local, image);
}

std::vector<SpotToSave> JFJochMultiReader::ReadSpots(int64_t image) const {
    const auto [k, local] = Locate(image);
    return inputs[k]->ReadSpots(local);
}

size_t JFJochMultiReader::GetNumberOfInputs() const {
    return inputs.size();
}

int64_t JFJochMultiReader::FirstImage(size_t input) const {
    return first_image.at(input);
}

uint64_t JFJochMultiReader::ImagesInInput(size_t input) const {
    return inputs.at(input)->GetNumberOfImages();
}