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

namespace {
    const UnitCell reference{70.0f, 71.0f, 40.0f, 90.0f, 93.0f, 90.0f};

    // Frames of a monoclinic cell (a, b, c, beta) with a percent's noise in each edge and half a degree in
    // the angles, its two near-equal edges a and b indexed the other way round in every second frame.
    void AddFrames(std::vector<UnitCell> &cells, int n, double a, double b, double c, double beta, std::mt19937 &rng) {
        std::normal_distribution<double> edge(1.0, 0.01), angle(0.0, 0.5);
        for (int i = 0; i < n; i++) {
            UnitCell u{static_cast<float>(a * edge(rng)), static_cast<float>(b * edge(rng)),
                       static_cast<float>(c * edge(rng)), static_cast<float>(90 + angle(rng)),
                       static_cast<float>(beta + angle(rng)), static_cast<float>(90 + angle(rng))};
            if (i % 2)
                u = {u.b, u.a, u.c, u.beta, u.alpha, u.gamma};
            cells.push_back(u);
        }
    }
}

TEST_CASE("CellForms_TwoFormsMinorFound", "[CellPopulation][portable]") {
    // Two forms 1.5-4 % apart in each edge, the second a quarter of the frames.
    std::mt19937 rng(3);
    std::vector<UnitCell> cells;
    AddFrames(cells, 1500, 70.0, 71.0, 40.0, 93.0, rng);
    AddFrames(cells, 500, 73.0, 72.0, 41.5, 91.5, rng);
    const auto forms = SeparateCellForms(cells, reference);
    REQUIRE(forms.minor_cell);
    size_t major_flagged = 0, minor_flagged = 0;
    for (size_t i = 0; i < cells.size(); i++)
        (i < 1500 ? major_flagged : minor_flagged) += forms.form[i];
    CHECK(major_flagged < 15);
    CHECK(minor_flagged > 450);
    CHECK(forms.minor_frames == major_flagged + minor_flagged);
    // Its cell with the axes in the reference's order of length - its own a > b the other way round, the
    // angle opposite each edge moving with it.
    CHECK(forms.minor_cell->a == Catch::Approx(72.0).margin(0.6));
    CHECK(forms.minor_cell->b == Catch::Approx(73.0).margin(0.6));
    CHECK(forms.minor_cell->c == Catch::Approx(41.5).margin(0.3));
    CHECK(forms.minor_cell->alpha == Catch::Approx(91.5).margin(0.6));
}

TEST_CASE("CellForms_OneFormNoMinor", "[CellPopulation][portable]") {
    std::mt19937 rng(5);
    std::vector<UnitCell> cells;
    AddFrames(cells, 2000, 70.0, 71.0, 40.0, 93.0, rng);
    const auto forms = SeparateCellForms(cells, reference);
    CHECK(forms.minor_frames == 0);
    CHECK(!forms.minor_cell);
}

TEST_CASE("CellForms_DriftIsOneForm", "[CellPopulation][portable]") {
    // One form whose cell grows by 0.1 % per edge halfway through, with a tight spread so that the two
    // halves could separate: closer than MIN_FORM_SEPARATION, it is one form.
    std::mt19937 rng(9);
    std::normal_distribution<double> edge(1.0, 0.0005);
    std::vector<UnitCell> cells;
    for (int i = 0; i < 2000; i++) {
        const double s = i < 1000 ? 1.0 : 1.001;
        cells.push_back({static_cast<float>(70.0 * s * edge(rng)), static_cast<float>(71.0 * s * edge(rng)),
                         static_cast<float>(40.0 * s * edge(rng)), 90.0f, 93.0f, 90.0f});
    }
    CHECK(SeparateCellForms(cells, reference).minor_frames == 0);
}

TEST_CASE("CellForms_StrayGroupIsNotAForm", "[CellPopulation][portable]") {
    // A second cell on 5 % of the frames is a stray group, not a crystal form.
    std::mt19937 rng(13);
    std::vector<UnitCell> cells;
    AddFrames(cells, 1900, 70.0, 71.0, 40.0, 93.0, rng);
    AddFrames(cells, 100, 73.0, 72.0, 41.5, 91.5, rng);
    CHECK(SeparateCellForms(cells, reference).minor_frames == 0);
}