// 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_StrayFramesBeyondTheTail", "[CellPopulation][portable]") {
    // Two forms, the given cell at the smaller one, and a few misindexed frames far out beyond its tail.
    // Each stray frame is a local maximum of the histogram with no valley of its own; it is not a
    // population, and must not stand for "everything", both forms included.
    const auto block = [](double mu, double w, int n, std::vector<double> &r) {
        for (int i = 0; i < n; i++) r.push_back(mu + w * ((i + 0.5) / n - 0.5));
    };
    std::vector<double> r;
    block(-0.055, 0.0582, 400, r);
    block(0.02, 0.0194, 220, r);
    for (double x : {0.0748, 0.0442, 0.0423, 0.8394, 0.7456}) r.push_back(x);
    const auto [low, high] = GivenCellPopulation(r);
    size_t kept_other = 0, kept_given = 0;
    for (int i = 0; i < 400; i++) kept_other += (r[i] >= low && r[i] <= high);
    for (int i = 400; i < 620; i++) kept_given += (r[i] >= low && r[i] <= high);
    CHECK(kept_given == 220);
    CHECK(kept_other < 10);
}

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));
}