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

#include <catch2/catch_test_macros.hpp>

#include <cmath>
#include <fstream>
#include <sstream>

#include "../rugnux/DiagnosticOutput.h"
#include "../image_analysis/scale_merge/Merge.h"

namespace {
    bool same(const rgb &a, const rgb &b) {
        return a.r == b.r && a.g == b.g && a.b == b.b;
    }

    // The columns of one row of <prefix>_plot.txt, whitespace-separated.
    std::vector<std::string> Columns(const std::string &line) {
        std::istringstream is(line);
        std::vector<std::string> ret;
        for (std::string token; is >> token;)
            ret.push_back(token);
        return ret;
    }

    std::vector<std::string> ReadLines(const std::string &path) {
        std::ifstream f(path);
        std::vector<std::string> ret;
        for (std::string line; std::getline(f, line);)
            ret.push_back(line);
        return ret;
    }
}

TEST_CASE("DiagnosticImage_Colors") {
    ColorScale scale;

    // One pixel of each kind, in the order the renderer tests them. Two measured pixels read zero,
    // so the median is zero and the third one saturates.
    std::vector<uint32_t> mask_bits{
            0,                                          // plain, no counts
            0,                                          // plain, no counts
            0,                                          // plain, the brightest pixel
            1u << PixelMask::ModuleGapPixelBit,
            1u << PixelMask::BeamStopPixelBit,
            1u << PixelMask::UserMaskedPixelBit,
            (1u << PixelMask::ModuleGapPixelBit) | (1u << PixelMask::BeamStopPixelBit),
            (1u << PixelMask::BeamStopPixelBit) | (1u << PixelMask::NoisyPixelBit)};
    const PixelMask mask(mask_bits);
    const std::vector<float> projection{0.0f, 0.0f, 1000.0f, NAN, NAN, NAN, NAN, NAN};

    const auto image = DetectorDiagnosticRGB(mask, projection);
    REQUIRE(image.size() == mask_bits.size());

    CHECK(same(image[0], scale.Apply(0.0f)));                          // white: no counts
    CHECK(same(image[2], scale.Apply(1.0f)));                          // indigo: the saturation value
    CHECK(same(image[3], scale.Apply(ColorScaleSpecial::Gap)));        // grey
    CHECK(same(image[4], scale.Apply(ColorScaleSpecial::BeamStop)));   // coral
    CHECK(same(image[5], scale.Apply(ColorScaleSpecial::BadPixel)));   // magenta
    // A gap cannot be shadowed - it carries no background for the detection to miss - so the gap
    // wins; every other mask bit loses to the shadow, which is what the picture is there to show.
    CHECK(same(image[6], scale.Apply(ColorScaleSpecial::Gap)));
    CHECK(same(image[7], scale.Apply(ColorScaleSpecial::BeamStop)));
}

TEST_CASE("DiagnosticImage_Contrast") {
    ColorScale scale;

    // The ramp saturates at eight times the MEDIAN of the measured pixels - 128 here, so 1024 - so
    // the background spans the ramp and the peak clips. Set from the brightest pixels instead, as
    // the preview and the viewer set their contrast, the peak would take the whole range and this
    // background would sit at the very bottom of the scale, a flat white: a missing background is
    // what a shadow is read off, so that picture cannot be judged.
    std::vector<uint32_t> mask_bits{0, 0, 0, 0, 0};
    const PixelMask mask(mask_bits);
    const std::vector<float> projection{0.0f, 128.0f, 128.0f, 1024.0f, 5000.0f};

    const auto image = DetectorDiagnosticRGB(mask, projection);
    const auto &lut = scale.LUTData();
    REQUIRE(image.size() == projection.size());

    CHECK(same(image[0], lut.front()));
    CHECK(same(image[1], lut[lut.size() / 8]));
    CHECK(same(image[2], lut[lut.size() / 8]));
    CHECK(same(image[3], lut.back()));
    // Clipped, and that costs nothing here - the spots are not what the picture is judged on.
    CHECK(same(image[4], lut.back()));
}

TEST_CASE("DiagnosticPlot_Format") {
    const std::string prefix = "diagnostic_plot_format_test";

    PerImagePlot plot;
    plot.bkg_estimate = {1.0f, 2.0f, 3.0f, 4.0f};
    plot.resolution_A = {2.5f, 2.4f, NAN, 2.6f};
    plot.spot_count = {100.0f, 120.0f, 0.0f, 90.0f};
    plot.scale_g = {1.0f, 1.1f, NAN, 0.9f};
    plot.mosaicity_deg = {0.12f, 0.12f, NAN, 0.13f};
    plot.cc_to_merge = {0.98f, 0.97f, NAN, 0.5f};
    plot.cc_n = {500, 480, 0, 120};
    plot.frame_disposition = {static_cast<uint8_t>(FrameDisposition::Merged),
                              static_cast<uint8_t>(FrameDisposition::Downgraded),
                              static_cast<uint8_t>(FrameDisposition::Rejected),
                              static_cast<uint8_t>(FrameDisposition::Rejected)};

    const GoniometerAxis goniometer("omega", 10.0f, 0.5f, {0, 1, 0}, {});
    WritePerImagePlot(prefix, plot, goniometer, 4);

    const auto lines = ReadLines(prefix + "_plot.txt");
    REQUIRE(lines.size() == 5);

    // Exactly one legend row, and it is a comment.
    REQUIRE(lines[0][0] == '#');
    for (size_t i = 1; i < lines.size(); i++)
        CHECK(lines[i].find('#') == std::string::npos);

    const auto legend = Columns(lines[0]);
    REQUIRE(legend.size() == 10);
    CHECK(legend[0] == "#image");
    CHECK(legend[1] == "angle_deg");
    CHECK(legend[6] == "sigma_M_deg");
    CHECK(legend[8] == "cc_n");
    CHECK(legend[9] == "merged");

    for (size_t i = 1; i < lines.size(); i++) {
        const auto row = Columns(lines[i]);
        REQUIRE(row.size() == legend.size());
        CHECK(row[0] == std::to_string(i - 1));
    }

    // Mid-exposure angle: start + ordinal * increment + increment / 2.
    CHECK(std::stof(Columns(lines[1])[1]) == 10.25f);
    CHECK(std::stof(Columns(lines[3])[1]) == 11.25f);

    // Merged and downgraded frames are in the merged data; a rejected one is not.
    CHECK(Columns(lines[1])[9] == "1");
    CHECK(Columns(lines[2])[9] == "1");
    CHECK(Columns(lines[3])[9] == "0");
    CHECK(Columns(lines[4])[9] == "0");

    // The count the correlation beside it was taken over.
    CHECK(Columns(lines[1])[8] == "500");

    // A frame nothing measured leaves its own columns unplottable, not the row missing.
    CHECK(Columns(lines[3])[3] == "nan");
}

TEST_CASE("DiagnosticPlot_NoDisposition") {
    const std::string prefix = "diagnostic_plot_nodisp_test";

    // Without the sweep-quality diagnostic, a frame counts as merged when it was scaled at all.
    PerImagePlot plot;
    plot.scale_g = {1.0f, NAN, 0.8f};
    WritePerImagePlot(prefix, plot, {}, 3);

    const auto lines = ReadLines(prefix + "_plot.txt");
    REQUIRE(lines.size() == 4);
    CHECK(Columns(lines[1])[9] == "1");
    CHECK(Columns(lines[2])[9] == "0");
    CHECK(Columns(lines[3])[9] == "1");

    // No goniometer: the angle column is present and unplottable, so the columns never shift.
    CHECK(Columns(lines[1]).size() == 10);
    CHECK(Columns(lines[1])[1] == "nan");
}

TEST_CASE("DiagnosticPlot_NoMerge") {
    const std::string prefix = "diagnostic_plot_nomerge_test";

    // Nothing was merged - --no-merge, or a sweep nothing indexed on - so "merged" has no answer to
    // give and must not read as a sweep every frame of which was thrown out.
    PerImagePlot plot;
    plot.bkg_estimate = {12.0f, 13.0f};
    plot.scale_g = {NAN, NAN};
    WritePerImagePlot(prefix, plot, {}, 2);

    const auto lines = ReadLines(prefix + "_plot.txt");
    REQUIRE(lines.size() == 3);
    CHECK(Columns(lines[1])[9] == "nan");
    CHECK(Columns(lines[2])[9] == "nan");
    // The columns a run without a merge still fills are there.
    CHECK(std::stof(Columns(lines[1])[2]) == 12.0f);
}