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

// Copyright (2019-2024) Paul Scherrer Institute

#include <catch2/catch_all.hpp>

#include "../writer/HDF5Objects.h"
#include "../image_analysis/NeuralNetInferenceClient.h"

TEST_CASE("NeuralNetResPredictor_Prepare", "[LinearAlgebra][Coord]") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.DetectorDistance_mm(75).IncidentEnergy_keV(12.4).BeamX_pxl(1000).BeamY_pxl(1000);

    std::vector<int16_t> v(experiment.GetPixelsNum(),0);
    v[1000 * experiment.GetXPixelsNum() + 1000] = 100;
    v[1000 * experiment.GetXPixelsNum() + 1001] = 20;
    v[1001 * experiment.GetXPixelsNum() + 1000] = 30;
    v[1001 * experiment.GetXPixelsNum() + 1001] = INT16_MIN;

    v[1050 * experiment.GetXPixelsNum() + 1050] = 52;
    v[2000 * experiment.GetXPixelsNum() + 1500] = 160;

    v[800 * experiment.GetXPixelsNum() + 600] = 49;
    v[1200 * experiment.GetXPixelsNum() + 600] = 36;
    v[800 * experiment.GetXPixelsNum() + 1400] = 64;

    NeuralNetInferenceClient predictor;

    REQUIRE(predictor.GetMaxPoolFactor(experiment) == 2);

    auto br = predictor.Prepare(experiment, v.data(), Quarter::BottomRight);
    REQUIRE(br.size() == 512 * 512);
    CHECK(br[0] == 10);
    CHECK(br[25 * 512 + 25] == 7);
    CHECK(br[500 * 512 + 250] == 12);

    auto tl = predictor.Prepare(experiment, v.data(), Quarter::TopLeft);
    REQUIRE(tl.size() == 512 * 512);
    CHECK(tl[100 * 512 + 200] == 7);

    auto tr = predictor.Prepare(experiment, v.data(), Quarter::TopRight);
    REQUIRE(tr.size() == 512 * 512);
    CHECK(tr[100 * 512 + 200] == 8);


    auto bl = predictor.Prepare(experiment, v.data(), Quarter::BottomLeft);
    REQUIRE(bl.size() == 512 * 512);
    CHECK(bl[100 * 512 + 200] == 6);
}