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

#include <catch2/catch_all.hpp>

#include "../common/DiffractionExperiment.h"
#include "../reader/JFJochMultiReader.h"

namespace {
    // A stills collection of n frames held in memory: frame i answers with one spot at x = tag + i, so
    // which input and which of its frames was read can be told from the spot.
    class FakeStills : public JFJochReader {
        uint64_t n;
        float tag;
        bool LoadImage_i(std::shared_ptr<JFJochReaderDataset> &, DataMessage &, std::vector<uint8_t> &,
                         int64_t, bool) override { return false; }
    public:
        FakeStills(uint64_t n_, float tag_, const DiffractionExperiment &x, std::vector<uint32_t> mask,
                   std::vector<float> efficiency = {}) : n(n_), tag(tag_) {
            auto d = std::make_shared<JFJochReaderDataset>();
            d->experiment = x;
            d->pixel_mask = std::make_shared<const PixelMask>(mask);
            d->efficiency = std::move(efficiency);
            SetStartMessage(d);
        }
        [[nodiscard]] uint64_t GetNumberOfImages() const override { return n; }
        void Close() override {}
        bool ReadRawImage(int64_t, JFJochReaderRawImage &) override { return true; }
        std::vector<SpotToSave> ReadSpots(int64_t image) const override {
            SpotToSave s;
            s.x = tag + static_cast<float>(image);
            return {s};
        }
    };

    DiffractionExperiment Stills() {
        DiffractionExperiment x(DetJF(1));
        x.BeamX_pxl(500).BeamY_pxl(250).DetectorDistance_mm(100).IncidentEnergy_keV(12.4);
        return x;
    }

    std::vector<uint32_t> EmptyMask(const DiffractionExperiment &x) {
        return std::vector<uint32_t>(x.GetXPixelsNum() * x.GetYPixelsNum(), 0);
    }
}

TEST_CASE("JFJochMultiReader_Concatenates", "[JFJochMultiReader][portable]") {
    const auto x = Stills();
    auto mask_a = EmptyMask(x);
    auto mask_b = EmptyMask(x);
    mask_a[3] = 1;
    mask_b[7] = 4;

    std::vector<std::unique_ptr<JFJochReader>> inputs;
    inputs.push_back(std::make_unique<FakeStills>(3, 100.0f, x, mask_a, std::vector<float>{0.5f, 0.5f, 0.5f}));
    inputs.push_back(std::make_unique<FakeStills>(2, 200.0f, x, mask_b));
    JFJochMultiReader reader(std::move(inputs));

    REQUIRE(reader.GetNumberOfImages() == 5);
    REQUIRE(reader.GetNumberOfInputs() == 2);
    CHECK(reader.FirstImage(1) == 3);
    CHECK(reader.ImagesInInput(1) == 2);

    CHECK(reader.ReadSpots(0)[0].x == Catch::Approx(100.0f));
    CHECK(reader.ReadSpots(2)[0].x == Catch::Approx(102.0f));
    CHECK(reader.ReadSpots(3)[0].x == Catch::Approx(200.0f));
    CHECK(reader.ReadSpots(4)[0].x == Catch::Approx(201.0f));
    CHECK_THROWS(reader.ReadSpots(5));

    const auto dataset = reader.GetDataset();
    const auto &mask = dataset->pixel_mask->GetMask();
    CHECK(mask[3] == 1);
    CHECK(mask[7] == 4);
    CHECK(mask[5] == 0);

    // The second input states no efficiency: its frames count as fully collected.
    REQUIRE(dataset->efficiency.size() == 5);
    CHECK(dataset->efficiency[2] == Catch::Approx(0.5f));
    CHECK(dataset->efficiency[3] == Catch::Approx(1.0f));
}

TEST_CASE("JFJochMultiReader_RefusesOtherGeometry", "[JFJochMultiReader][portable]") {
    const auto x = Stills();
    auto other = Stills();

    SECTION("distance") { other.DetectorDistance_mm(101); }
    SECTION("beam centre") { other.BeamX_pxl(503); }
    SECTION("wavelength") { other.IncidentEnergy_keV(12.6); }

    std::vector<std::unique_ptr<JFJochReader>> inputs;
    inputs.push_back(std::make_unique<FakeStills>(1, 0.0f, x, EmptyMask(x)));
    inputs.push_back(std::make_unique<FakeStills>(1, 0.0f, other, EmptyMask(other)));
    CHECK_THROWS(JFJochMultiReader(std::move(inputs)));
}

TEST_CASE("JFJochMultiReader_RefusesRotation", "[JFJochMultiReader][portable]") {
    const auto x = Stills();
    auto rotation = Stills();
    rotation.Goniometer(GoniometerAxis("omega", 0.0f, 0.1f, Coord(1, 0, 0), {}));

    std::vector<std::unique_ptr<JFJochReader>> inputs;
    inputs.push_back(std::make_unique<FakeStills>(1, 0.0f, x, EmptyMask(x)));
    inputs.push_back(std::make_unique<FakeStills>(1, 0.0f, rotation, EmptyMask(rotation)));
    CHECK_THROWS(JFJochMultiReader(std::move(inputs)));
}