// 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/Definitions.h"
#include "../image_analysis/IceScore.h"
#include "../image_analysis/spot_finding/SpotUtils.h"

namespace {
    constexpr float TWO_PI = 6.283185307f;
    // The pipeline's own ice band half-width (SpotFindingSettings::ice_ring_width_Q_recipA).
    constexpr float ICE_W = 0.03f;

    SpotToSave spot(float d_A, float intensity) {
        return SpotToSave{.intensity = intensity, .d_A = d_A};
    }

    AzimuthalIntegrationSettings ice_settings() {
        AzimuthalIntegrationSettings settings;
        settings.QSpacing_recipA(0.006f).QRange_recipA(0.1f, 4.5f);
        return settings;
    }

    int bin_of(const AzimuthalIntegrationSettings &settings, float d_A) {
        return static_cast<int>(std::lround((TWO_PI / d_A - settings.GetLowQ_recipA())
                                            / settings.GetQSpacing_recipA() - 0.5f));
    }
}

TEST_CASE("ProteinScore_EmptyAndHighResolutionOnly") {
    CHECK(ProteinScore({}) == 0.0f);

    // Nothing beyond 5 A: a salt or ice powder pattern, however strong, is not protein.
    std::vector<SpotToSave> spots;
    for (int i = 0; i < 200; i++)
        spots.push_back(spot(1.0f + 0.015f * static_cast<float>(i), 5000.0f));
    CHECK(ProteinScore(spots) == 0.0f);
}

TEST_CASE("ProteinScore_SaturatesAndCountsShellsNotSpots") {
    // Six well-separated shells above 5 A, one strong spot each: enough evidence to be sure.
    std::vector<SpotToSave> shells;
    for (const float d: {5.5f, 6.5f, 8.0f, 10.0f, 14.0f, 20.0f})
        shells.push_back(spot(d, 1000.0f));
    const float six = ProteinScore(shells);
    CHECK(six > 0.5f);

    // Ten times as many spots, all in the SAME shell: a parasitic ring, not a lattice. One shell is
    // worth at most 1, so it must score far below the six shells above.
    std::vector<SpotToSave> one_ring;
    for (int i = 0; i < 60; i++)
        one_ring.push_back(spot(5.5f, 1000.0f));
    CHECK(ProteinScore(one_ring) < 0.3f);
    CHECK(ProteinScore(one_ring) < six);

    // Saturation: making every spot a hundred times stronger does not raise the score, because the
    // weight is measured against the frame's own median spot.
    std::vector<SpotToSave> strong;
    for (const float d: {5.5f, 6.5f, 8.0f, 10.0f, 14.0f, 20.0f})
        strong.push_back(spot(d, 100000.0f));
    CHECK(ProteinScore(strong) == Catch::Approx(six));

    // And it stays inside [0, 1] however much evidence there is.
    std::vector<SpotToSave> many;
    for (int i = 0; i < 400; i++)
        many.push_back(spot(5.1f + 0.08f * static_cast<float>(i % 300), 1000.0f));
    CHECK(ProteinScore(many) > 0.9f);
    CHECK(ProteinScore(many) < 1.0f);
}

TEST_CASE("IceScoreRadial_FlatProfileIsNotIce") {
    const auto settings = ice_settings();
    const int q_bins = settings.GetQBinCount();
    const std::vector<float> flat(q_bins, 100.0f);
    // Per-pixel standard deviation and pixel count: the score uses std / sqrt(count) = 2 photons.
    const std::vector<float> sigma(q_bins, 20.0f);
    const std::vector<uint64_t> count(q_bins, 100);
    CHECK(IceScoreRadial(flat, sigma, count, q_bins, settings) == 0.0f);

    // No standard deviation, no radial channel: the FPGA azimuthal integration does not produce one.
    CHECK(IceScoreRadial(flat, {}, count, q_bins, settings) == 0.0f);
}

TEST_CASE("IceScoreRadial_HexagonalAndCubicPatterns") {
    const auto settings = ice_settings();
    const int q_bins = settings.GetQBinCount();
    const std::vector<float> sigma(q_bins, 20.0f);
    const std::vector<uint64_t> count(q_bins, 100);

    // A single strong band is not ice - real ice shows a whole pattern, and the band-count
    // concordance test is what refuses one bin.
    std::vector<float> one_band(q_bins, 100.0f);
    one_band[bin_of(settings, ICE_RING_RES_A[0])] = 400.0f;
    CHECK(IceScoreRadial(one_band, sigma, count, q_bins, settings) < 0.5f);

    // The whole hexagonal pattern is.
    std::vector<float> hexagonal(q_bins, 100.0f);
    for (const float d: ICE_RING_RES_A) {
        const int b = bin_of(settings, d);
        if (b >= 0 && b < q_bins)
            hexagonal[b] = 130.0f;
    }
    CHECK(IceScoreRadial(hexagonal, sigma, count, q_bins, settings) > 0.5f);

    // So is the cubic one, which shares only three lines with it - the phase that a hexagonal-only
    // detector misses entirely.
    std::vector<float> cubic(q_bins, 100.0f);
    for (const float d: ICE_RING_CUBIC_RES_A) {
        const int b = bin_of(settings, d);
        if (b >= 0 && b < q_bins)
            cubic[b] = 130.0f;
    }
    CHECK(IceScoreRadial(cubic, sigma, count, q_bins, settings) > 0.5f);

    // The same excess spread over bins belonging to no phase is not ice.
    std::vector<float> off_band(q_bins, 100.0f);
    for (int i = 100; i < q_bins - 100; i += 37)
        off_band[i] = 130.0f;
    CHECK(IceScoreRadial(off_band, sigma, count, q_bins, settings) < 0.5f);
}

TEST_CASE("IceScoreRadial_AzimuthalProfileFoldsToTheSameAnswer") {
    AzimuthalIntegrationSettings settings;
    settings.QSpacing_recipA(0.006f).QRange_recipA(0.1f, 4.5f).AzimuthalBinCount(4);
    const int q_bins = settings.GetQBinCount();

    std::vector<float> flat(q_bins, 100.0f);
    std::vector<float> sigma(q_bins, 20.0f);
    std::vector<uint64_t> count(q_bins, 400);
    for (const float d: ICE_RING_RES_A) {
        const int b = bin_of(settings, d);
        if (b >= 0 && b < q_bins)
            flat[b] = 130.0f;
    }
    std::vector<float> sectors(static_cast<size_t>(q_bins) * 4);
    std::vector<float> sectors_sigma(static_cast<size_t>(q_bins) * 4);
    std::vector<uint64_t> sectors_count(static_cast<size_t>(q_bins) * 4);
    for (int az = 0; az < 4; az++)
        for (int q = 0; q < q_bins; q++) {
            sectors[static_cast<size_t>(az) * q_bins + q] = flat[q];
            sectors_sigma[static_cast<size_t>(az) * q_bins + q] = sigma[q];
            sectors_count[static_cast<size_t>(az) * q_bins + q] = count[q] / 4;
        }

    CHECK(IceScoreRadial(sectors, sectors_sigma, sectors_count, q_bins, settings)
          == Catch::Approx(IceScoreRadial(flat, sigma, count, q_bins, settings)));
}

namespace {
    // IceScoreSpots reads d-spacings, not spots: it is fed the list from BEFORE the max-spot-count
    // cap, which orders ice-band spots last and would otherwise discard the very spots it counts.
    std::vector<float> d_of(const std::vector<SpotToSave> &spots) {
        std::vector<float> d;
        d.reserve(spots.size());
        for (const auto &s: spots)
            d.push_back(s.d_A);
        return d;
    }
}

TEST_CASE("IceScoreSpots_ExcessOnTheIceRadii") {
    const auto settings = ice_settings();
    const int q_bins = settings.GetQBinCount();
    // A detector that covers every radius equally, so the control offsets are directly comparable.
    const std::vector<uint64_t> count(q_bins, 10000);

    // 400 spots spread evenly in q: whatever lands on an ice radius is what the control predicts.
    std::vector<SpotToSave> even;
    const float q_lo = 1.3f, q_hi = 4.2f;
    for (int i = 0; i < 400; i++) {
        const float q = q_lo + (q_hi - q_lo) * static_cast<float>(i) / 399.0f;
        even.push_back(spot(TWO_PI / q, 1000.0f));
    }
    CHECK(IceScoreSpots(d_of(even), count, q_bins, settings, ICE_W) < 0.5f);

    // The same frame with 10 extra spots planted on each hexagonal radius.
    std::vector<SpotToSave> with_ice = even;
    for (const float d: ICE_RING_RES_A)
        for (int i = 0; i < 10; i++)
            with_ice.push_back(spot(d, 1000.0f));
    CHECK(IceScoreSpots(d_of(with_ice), count, q_bins, settings, ICE_W) > 0.5f);

    // Two spots that both happen to sit on a ring are not ice: the ratio term refuses them even
    // though the Poisson tail alone would not.
    std::vector<SpotToSave> two;
    two.push_back(spot(ICE_RING_RES_A[0], 1000.0f));
    two.push_back(spot(ICE_RING_RES_A[1], 1000.0f));
    CHECK(IceScoreSpots(d_of(two), count, q_bins, settings, ICE_W) < 0.5f);

    CHECK(IceScoreSpots({}, count, q_bins, settings, ICE_W) == 0.0f);
}

TEST_CASE("IceScore_TakesTheStrongerChannel") {
    const auto settings = ice_settings();
    const int q_bins = settings.GetQBinCount();
    const std::vector<uint64_t> count(q_bins, 10000);
    const std::vector<float> sigma(q_bins, 200.0f);

    // Powder ice, no spots at all: the radial channel carries it on its own.
    std::vector<float> powder(q_bins, 100.0f);
    for (const float d: ICE_RING_RES_A) {
        const int b = bin_of(settings, d);
        if (b >= 0 && b < q_bins)
            powder[b] = 130.0f;
    }
    CHECK(IceScore(powder, sigma, count, q_bins, settings, {}, ICE_W) > 0.5f);

    // Ice as discrete crystallites: the profile is flat and only the spot channel sees it.
    const std::vector<float> flat(q_bins, 100.0f);
    std::vector<SpotToSave> textured;
    for (int i = 0; i < 400; i++)
        textured.push_back(spot(TWO_PI / (1.3f + 2.9f * static_cast<float>(i) / 399.0f), 1000.0f));
    for (const float d: ICE_RING_RES_A)
        for (int i = 0; i < 10; i++)
            textured.push_back(spot(d, 1000.0f));
    CHECK(IceScoreRadial(flat, sigma, count, q_bins, settings) == 0.0f);
    CHECK(IceScore(flat, sigma, count, q_bins, settings, d_of(textured), ICE_W) > 0.5f);
}

// The spot channel must be fed the list from BEFORE the max-spot-count cap. FilterSpotsByCount
// orders ice-band spots LAST when indexing is not to use them, so on a frame with more spots than
// the budget it discards the ice first - and the channel that exists for ice arriving as discrete
// spots then reads zero on exactly the frames it was written for.
TEST_CASE("IceScoreSpots_ReadsThePreCapList") {
    const auto settings = ice_settings();
    const int q_bins = settings.GetQBinCount();
    const std::vector<uint64_t> count(q_bins, 10000);

    std::vector<SpotToSave> spots;
    const float q_lo = 1.3f, q_hi = 4.2f;
    for (int i = 0; i < 1800; i++) {
        const float q = q_lo + (q_hi - q_lo) * static_cast<float>(i) / 1799.0f;
        spots.push_back(spot(TWO_PI / q, 1000.0f));
    }
    for (const float d: ICE_RING_RES_A)
        for (int i = 0; i < 60; i++)
            spots.push_back(spot(d, 1000.0f));
    MarkIceRings(spots, ICE_W);

    const std::vector<float> before = d_of(spots);
    FilterSpotsByCount(spots, 1000, true);
    const std::vector<float> after = d_of(spots);

    REQUIRE(before.size() > after.size());
    CHECK(IceScoreSpots(before, count, q_bins, settings, ICE_W) > 0.5f);
    CHECK(IceScoreSpots(after, count, q_bins, settings, ICE_W) < 0.5f);
}