// 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"

namespace {
    // A monoclinic cell, and a second form of the same protein whose edges differ from it by 1.5-4%.
    const UnitCell given_form{96.0f, 94.5f, 50.5f, 90.0f, 91.2f, 90.0f};
    const UnitCell other_form{92.0f, 93.0f, 49.0f, 90.0f, 92.7f, 90.0f};

    // The cell grown on every edge alike to (1 + r) times its own volume: a frame of it refined at a detector
    // distance off by about a third of r.
    UnitCell Scaled(const UnitCell &u, double r) {
        const auto s = static_cast<float>(std::cbrt(1.0 + r));
        return {u.a * s, u.b * s, u.c * s, u.alpha, u.beta, u.gamma};
    }

    // How many of cells[from, to) are kept.
    size_t Kept(const std::vector<UnitCell> &cells, const UnitCell &given, size_t from, size_t to) {
        const auto [low, high] = GivenCellPopulation(cells, given);
        const double v = gemmi::UnitCell(given).volume;
        size_t n = 0;
        for (size_t i = from; i < to; i++) {
            const double r = gemmi::UnitCell(cells[i]).volume / v - 1.0;
            n += (r >= low && r <= high);
        }
        return n;
    }
}

TEST_CASE("CellPopulation_OnePopulationKeepsEverything", "[CellPopulation][portable]") {
    std::mt19937 rng(7);
    std::normal_distribution<double> d(0.004, 0.01);
    std::vector<UnitCell> cells;
    for (int i = 0; i < 2000; i++) cells.push_back(Scaled(given_form, d(rng)));
    CHECK(Kept(cells, given_form, 0, cells.size()) == cells.size());
}

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> spread(0.0, 0.01);
    std::vector<UnitCell> cells;
    for (int i = 0; i < 450; i++) cells.push_back(Scaled(given_form, spread(rng)));
    for (int i = 0; i < 150; i++) cells.push_back(Scaled(other_form, spread(rng)));
    CHECK(Kept(cells, given_form, 0, 450) > 440);
    CHECK(Kept(cells, given_form, 450, 600) < 3);
    CHECK(Kept(cells, other_form, 450, 600) > 145);
    CHECK(Kept(cells, other_form, 0, 450) < 5);
}

TEST_CASE("CellPopulation_SecondFormNearerInVolume", "[CellPopulation][portable]") {
    // The same two forms refined at a detector distance 3.5% long: every cell 11% larger in volume, so that the
    // other form, 9% smaller than the given one, comes out 1% from the given cell and the given form 11% from
    // it. The given form's proportions still say which population is its own.
    std::mt19937 rng(17);
    std::normal_distribution<double> spread(0.0, 0.01);
    std::vector<UnitCell> cells;
    for (int i = 0; i < 300; i++) cells.push_back(Scaled(given_form, 0.11 + spread(rng)));
    for (int i = 0; i < 500; i++) cells.push_back(Scaled(other_form, 0.11 + spread(rng)));
    CHECK(Kept(cells, given_form, 0, 300) > 290);
    CHECK(Kept(cells, given_form, 300, 800) < 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<UnitCell> cells;
    for (int i = 0; i < 2000; i++) cells.push_back(Scaled(given_form, d(rng)));
    cells.push_back(given_form);
    CHECK(Kept(cells, given_form, 0, 2000) == 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 = [](const UnitCell &u, double mu, double w, int n, std::vector<UnitCell> &cells) {
        for (int i = 0; i < n; i++) cells.push_back(Scaled(u, mu + w * ((i + 0.5) / n - 0.5)));
    };
    std::vector<UnitCell> cells;
    block(other_form, -0.055, 0.0582, 400, cells);
    block(given_form, 0.02, 0.0194, 220, cells);
    for (double x : {0.0748, 0.0442, 0.0423, 0.8394, 0.7456}) cells.push_back(Scaled(given_form, x));
    CHECK(Kept(cells, given_form, 400, 620) == 220);
    CHECK(Kept(cells, given_form, 0, 400) < 10);
}

TEST_CASE("CellPopulation_TooFewFrames", "[CellPopulation][portable]") {
    const auto [low, high] = GivenCellPopulation({given_form, Scaled(given_form, 0.1), Scaled(given_form, 0.2)},
                                                 given_form);
    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);
}