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

#include <catch2/catch_all.hpp>

#include "../image_analysis/spot_finding/SpotUtils.h"
#include "../image_analysis/indexing/AnalyzeIndexing.h"

TEST_CASE("FilterSpuriousHighResolutionSpots") {
    std::vector<SpotToSave> spots;
    spots.push_back(SpotToSave{.x = 1, .y = 2, .intensity = 3, .d_A = 18.0, .indexed = false});
    spots.push_back(SpotToSave{.x = 1, .y = 2, .intensity = 3, .d_A = 20.0, .indexed = false});
    spots.push_back(SpotToSave{.x = 1, .y = 2, .intensity = 3, .d_A = 30.0, .indexed = false});

    spots.push_back(SpotToSave{.x = 1, .y = 2, .intensity = 3, .d_A = 6.0, .indexed = false});

    spots.push_back(SpotToSave{.x = 1, .y = 2, .intensity = 3, .d_A = 2.0, .indexed = false});
    spots.push_back(SpotToSave{.x = 1, .y = 2, .intensity = 3, .d_A = 1.9, .indexed = false});
    spots.push_back(SpotToSave{.x = 1, .y = 2, .intensity = 3, .d_A = 1.3, .indexed = false});

    FilterSpuriousHighResolutionSpots(spots, 1.57); // roughly 0.25 in 1/d

    REQUIRE(spots.size() == 4);
    // Spots are sorted by resolution
    CHECK(spots[0].d_A == Catch::Approx(30.0));
    CHECK(spots[1].d_A == Catch::Approx(20.0));
    CHECK(spots[2].d_A == Catch::Approx(18.0));
    CHECK(spots[3].d_A == Catch::Approx(6.0));
}


TEST_CASE("GetResolution") {
    // Eleven equally strong spots at 1/d^2 = 0.1, 0.2, ... 1.1. Walking in from the highest-resolution
    // one, four of the eleven are the first to carry 30% of the weight, so the quantile is the fourth
    // spot in, 1/d^2 = 0.8. The estimate is that resolution taken 2.25x further in 1/d.
    std::vector<SpotToSave> spots;
    for (int i = 1; i <= 11; i++)
        spots.push_back(SpotToSave{.intensity = 100.0f, .d_A = 1.0f / std::sqrt(0.1f * static_cast<float>(i))});
    const auto d = GetResolution(spots);
    REQUIRE(d.has_value());
    CHECK(*d == Catch::Approx(1.0 / (2.25 * std::sqrt(0.8))).epsilon(1e-4));

    // The answer is not limited to what a detector records. Keeping only the five spots a detector
    // reaching 1/d^2 = 0.5 would have recorded leaves the quantile at 0.4, and the estimate still
    // extrapolates 2.25x past it instead of stopping at the cut.
    const std::vector<SpotToSave> cut(spots.begin(), spots.begin() + 5);
    CHECK(*GetResolution(cut) == Catch::Approx(1.0 / (2.25 * std::sqrt(0.4))).epsilon(1e-4));

    // Ice-flagged spots take no part, however strong they are.
    std::vector<SpotToSave> with_ice = spots;
    with_ice.push_back(SpotToSave{.intensity = 1e6f, .d_A = 0.5f, .ice_ring = true});
    CHECK(*GetResolution(with_ice) == Catch::Approx(*d));

    // A weak high-resolution spot moves the answer far less than a strong one, which is the point of
    // weighting by sqrt(I) rather than counting: the old order statistic would follow it entirely.
    std::vector<SpotToSave> with_spur = spots;
    with_spur.push_back(SpotToSave{.intensity = 1.0f, .d_A = 0.5f});
    CHECK(*GetResolution(with_spur) == Catch::Approx(*d).epsilon(0.02));

    // Too few spots to have a fall-off at all.
    CHECK_FALSE(GetResolution(std::vector<SpotToSave>(3)).has_value());
}

TEST_CASE("SpotBudgetFromEvidence") {
    // One image's worth of spots, repeated over 60 frames as the first pass does: the first 100 index
    // and the next 100 do not. Every indexed spot adds 1 - 0.2 and every unindexed one takes 0.2 away,
    // so the running tally rises to rank 100 and falls after it.
    std::vector<SpotToSave> spots(200);
    for (size_t i = 0; i < spots.size(); i++)
        spots[i].indexed = i < 100;

    constexpr int frames = 60;
    std::vector<int64_t> indexed(spots.size(), 0), counted(spots.size(), 0);
    for (int f = 0; f < frames; f++)
        AddSpotBudgetEvidence(spots, false, indexed, counted);
    CHECK(SpotBudgetFromEvidence(indexed, counted) == 100);

    // Spots that go on indexing all the way down: the tally never falls, so there is nothing to cut.
    for (auto &s: spots)
        s.indexed = true;
    std::vector<int64_t> all_hit(spots.size(), 0), all_seen(spots.size(), 0);
    for (int f = 0; f < frames; f++)
        AddSpotBudgetEvidence(spots, false, all_hit, all_seen);
    CHECK(SpotBudgetFromEvidence(all_hit, all_seen) == 0);

    // A budget already cut to its peak has no fall left in it, so a second measurement takes nothing
    // further off: the rule does not ratchet down on repetition.
    CHECK(SpotBudgetFromEvidence({indexed.begin(), indexed.begin() + 100},
                                 {counted.begin(), counted.begin() + 100}) == 0);

    // Ice-flagged spots take no part, so a run of them neither ends the budget nor moves it: the peak
    // stays at the last indexed non-ice rank before them.
    std::vector<SpotToSave> with_ice(300);
    for (size_t i = 0; i < with_ice.size(); i++) {
        with_ice[i].ice_ring = (i >= 100 && i < 160);
        with_ice[i].indexed = i < 100;
    }
    std::vector<int64_t> ice_indexed(with_ice.size(), 0), ice_counted(with_ice.size(), 0);
    for (int f = 0; f < frames; f++)
        AddSpotBudgetEvidence(with_ice, false, ice_indexed, ice_counted);
    CHECK(SpotBudgetFromEvidence(ice_indexed, ice_counted) == 100);

    // Nothing indexes: no rank carries evidence and there is no budget to report.
    for (auto &s: spots)
        s.indexed = false;
    std::vector<int64_t> none_indexed(spots.size(), 0), none_counted(spots.size(), 0);
    for (int f = 0; f < frames; f++)
        AddSpotBudgetEvidence(spots, false, none_indexed, none_counted);
    CHECK(SpotBudgetFromEvidence(none_indexed, none_counted) == 0);

    // The case a bare argmax gets wrong: the spots index at exactly the gate's own fraction at every
    // depth, so there is no depth at which the list stops being reflections. The tally still has a
    // maximum - it always does - but the fall from it is inside the counting noise, and nothing is cut.
    std::vector<SpotToSave> flat(1000);
    for (size_t i = 0; i < flat.size(); i++)
        flat[i].indexed = (i % 5 == 0);
    std::vector<int64_t> flat_indexed(flat.size(), 0), flat_counted(flat.size(), 0);
    for (int f = 0; f < frames; f++)
        AddSpotBudgetEvidence(flat, false, flat_indexed, flat_counted);
    CHECK(SpotBudgetFromEvidence(flat_indexed, flat_counted) == 0);
}