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

#include <random>

#include <catch2/catch_all.hpp>

#include "../image_analysis/scale_merge/CellPopulation.h"

TEST_CASE("CellPopulation_OnePopulationKeepsEverything", "[CellPopulation][portable]") {
    std::mt19937 rng(7);
    std::normal_distribution<double> d(0.004, 0.01);
    std::vector<double> r(2000);
    for (auto &x : r) x = d(rng);
    const auto [low, high] = GivenCellPopulation(r);
    for (double x : r)
        CHECK((x >= low && x <= high));
}

TEST_CASE("CellPopulation_SecondFormLeftOut", "[CellPopulation][portable]") {
    // Two crystal forms 9% apart in volume, the given cell one of them; whichever is given, its own
    // population is kept and the other left out.
    std::mt19937 rng(11);
    std::normal_distribution<double> small(0.0, 0.012), large(0.09, 0.008);
    std::vector<double> r;
    for (int i = 0; i < 450; i++) r.push_back(small(rng));
    for (int i = 0; i < 150; i++) r.push_back(large(rng));

    // Kept frames of each form, against either cell.
    const auto kept = [](const std::vector<double> &x, size_t from, size_t to) {
        const auto [low, high] = GivenCellPopulation(x);
        size_t n = 0;
        for (size_t i = from; i < to; i++)
            n += (x[i] >= low && x[i] <= high);
        return n;
    };
    CHECK(kept(r, 0, 450) > 440);
    CHECK(kept(r, 450, 600) < 3);

    std::vector<double> r2;
    for (double x : r) r2.push_back((1.0 + x) / 1.09 - 1.0);
    CHECK(kept(r2, 450, 600) > 145);
    CHECK(kept(r2, 0, 450) < 5);
}

TEST_CASE("CellPopulation_GivenCellOffTheCrystal", "[CellPopulation][portable]") {
    // A given cell 3.5% smaller than the crystal's, and a single stray frame right at it: the crystal's
    // population is kept, not the stray frame.
    std::mt19937 rng(3);
    std::normal_distribution<double> d(0.035, 0.003);
    std::vector<double> r(2000);
    for (auto &x : r) x = d(rng);
    r.push_back(0.0);
    const auto [low, high] = GivenCellPopulation(r);
    size_t kept = 0;
    for (double x : r)
        kept += (x >= low && x <= high);
    CHECK(kept >= 2000);
}

TEST_CASE("CellPopulation_TooFewFrames", "[CellPopulation][portable]") {
    const auto [low, high] = GivenCellPopulation({0.0, 0.1, 0.2});
    CHECK(std::isinf(low));
    CHECK(std::isinf(high));
}