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

#include <catch2/catch_all.hpp>

#include <array>
#include <cmath>
#include <numbers>
#include <vector>

#include "../image_analysis/scale_merge/TranslationalNCS.h"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"

namespace {
    // For a half-integer pseudo-translation the classes carry (1 +- f) times the mean intensity, and
    // the Patterson peak at that vector comes out at exactly f of the origin. f = 0.6 therefore models
    // a 60% peak - as strong as the strongest genuine pseudo-symmetry in a 137-dataset corpus. f = 1
    // would mean the weak class is EXTINCT, which is a lattice translation, not a pseudo-symmetry.
    constexpr double kRealisticDepth = 0.6;

    gemmi::UnitCell Cell(double a, double b, double c) {
        gemmi::UnitCell out;
        out.set(a, b, c, 90.0, 90.0, 90.0);
        return out;
    }

    // A synthetic merge over the reciprocal asymmetric unit: Wilson-distributed intensities from a
    // hash of the index (so the crystal is reproduced bit for bit), optionally multiplied by the
    // factor a translational pseudo-symmetry u imposes, |1 + exp(2 pi i h.u)|^2 / 2 written with a
    // depth f so the weak class is suppressed rather than extinguished.
    std::vector<MergedReflection> Synthetic(const gemmi::UnitCell &cell, const char *space_group,
                                            const std::array<double, 3> *u, double f,
                                            double d_min = 3.0) {
        const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(space_group);
        REQUIRE(sg != nullptr);
        const gemmi::GroupOps ops = sg->operations();
        const gemmi::ReciprocalAsu asu(sg);
        std::vector<MergedReflection> out;
        const int hmax = static_cast<int>(cell.a / d_min) + 1;
        const int kmax = static_cast<int>(cell.b / d_min) + 1;
        const int lmax = static_cast<int>(cell.c / d_min) + 1;
        for (int h = -hmax; h <= hmax; ++h)
            for (int k = -kmax; k <= kmax; ++k)
                for (int l = -lmax; l <= lmax; ++l) {
                    const gemmi::Op::Miller hkl{{h, k, l}};
                    if ((h == 0 && k == 0 && l == 0) || !asu.is_in(hkl)
                        || ops.is_systematically_absent(hkl))
                        continue;
                    const double d2 = cell.calculate_1_d2(hkl);
                    if (d2 <= 0 || 1.0 / std::sqrt(d2) < d_min)
                        continue;
                    const uint32_t seed = static_cast<uint32_t>(h * 73856093 ^ k * 19349663
                                                                ^ l * 83492791);
                    const double uni = ((seed * 2654435761u) >> 8) / static_cast<double>(1 << 24);
                    double I = 1000.0 * -std::log(std::max(1e-6, uni)) * std::exp(-0.5 * 20.0 * d2);
                    if (u != nullptr) {
                        const double phase = 2.0 * std::numbers::pi
                                             * (h * (*u)[0] + k * (*u)[1] + l * (*u)[2]);
                        I *= 1.0 + f * std::cos(phase);
                    }
                    MergedReflection r;
                    r.h = h; r.k = k; r.l = l;
                    r.d = static_cast<float>(1.0 / std::sqrt(d2));
                    r.I = static_cast<float>(I);
                    r.sigma = static_cast<float>(0.02 * std::fabs(r.I) + 1.0);
                    out.push_back(r);
                }
        return out;
    }
}

// The detector must find a pseudo-translation that is there, and name it.
TEST_CASE("A pseudo-translation is detected and its vector recovered", "[tncs]") {
    const gemmi::UnitCell cell = Cell(70.0, 85.0, 110.0);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    const std::array<double, 3> u{0.5, 0.5, 0.5};
    const auto merged = Synthetic(cell, "P 21 21 21", &u, kRealisticDepth);
    const auto r = AnalyzeTranslationalNCS(merged, cell, sg);

    REQUIRE(r.measurable);
    CHECK(r.detected);
    CHECK(r.peak_percent > 20.0);
    CHECK(r.undeclared_lattice_translations.empty());
    CHECK(r.peak_z > 5.0);
    CHECK(r.modulation > 2.5 * r.modulation_null);
    // The vector is refined to maximise the modulation, not to minimise its distance from the truth,
    // and the greedy search will trade a little accuracy for a little contrast. It is a starting
    // point for a molecular-replacement program (which refines it), not a measured constant, so the
    // tolerance here is the honest one - about 0.05 fractional, roughly 5 A on a 100 A axis.
    for (int i = 0; i < 3; ++i)
        CHECK(std::fabs(r.vector_frac[i] - 0.5) < 0.05);
    CHECK(r.commensurate);
    CHECK(r.commensurate_denominator == 2);
    // A half-integer pseudo-translation is pseudo-centring, not a wrong cell: the suppressed class is
    // weak, not absent, so the claim "your cell may be a supercell" must NOT be made here.
    CHECK_FALSE(r.near_extinct_class);
    CHECK(TranslationalNCSToText(r).find("supercell candidate") == std::string::npos);
}

// The false-positive rate is the whole product: the same crystal without a pseudo-translation must
// say nothing.
TEST_CASE("A crystal without a pseudo-translation stays silent", "[tncs]") {
    const gemmi::UnitCell cell = Cell(70.0, 85.0, 110.0);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    const auto merged = Synthetic(cell, "P 21 21 21", nullptr, 0.0);
    const auto r = AnalyzeTranslationalNCS(merged, cell, sg);

    REQUIRE(r.measurable);
    CHECK_FALSE(r.detected);
}

// The case that would make the detector useless if it were got wrong. A centred lattice has a real,
// full-height off-origin Patterson peak at its centring translation BY CONSTRUCTION - every C-, I-
// and F-centred crystal has one, and none of them has a pseudo-symmetry because of it. The centring
// vectors are origin-equivalent points and must be excluded along with the origin.
TEST_CASE("A genuinely centred lattice is not called", "[tncs]") {
    struct Case { const char *sg; gemmi::UnitCell cell; };
    const std::vector<Case> cases{
        {"C 1 2 1", Cell(120.0, 75.0, 90.0)},
        {"I 2 2 2", Cell(80.0, 95.0, 115.0)},
        {"F 2 2 2", Cell(110.0, 125.0, 140.0)},
    };
    for (const auto &c : cases) {
        const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(c.sg);
        REQUIRE(sg != nullptr);
        const auto merged = Synthetic(c.cell, c.sg, nullptr, 0.0);
        const auto r = AnalyzeTranslationalNCS(merged, c.cell, sg);
        INFO(c.sg << ": peak " << r.peak_percent << "% z " << r.peak_z << " modulation "
                  << r.modulation << " vs null " << r.modulation_null);
        REQUIRE(r.measurable);
        CHECK_FALSE(r.detected);
    }
}

// A centred lattice that ALSO has a pseudo-translation must still be called - excluding the centring
// vectors must not blind the test to a real one somewhere else.
TEST_CASE("A centred lattice with a real pseudo-translation is still called", "[tncs]") {
    const gemmi::UnitCell cell = Cell(120.0, 75.0, 90.0);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("C 1 2 1");
    const std::array<double, 3> u{0.0, 0.0, 0.5};
    const auto merged = Synthetic(cell, "C 1 2 1", &u, kRealisticDepth);
    const auto r = AnalyzeTranslationalNCS(merged, cell, sg);
    INFO("peak " << r.peak_percent << "% z " << r.peak_z << " u " << r.vector_frac[0] << " "
                 << r.vector_frac[1] << " " << r.vector_frac[2]);
    REQUIRE(r.measurable);
    CHECK(r.detected);
    CHECK(std::fabs(r.vector_frac[2] - 0.5) < 0.03);
}

// A cell too small to have anything in the 20-5 A band must say so, and must NOT say "no tNCS".
TEST_CASE("A small-molecule cell reports that it could not be measured", "[tncs]") {
    const gemmi::UnitCell cell = Cell(9.0, 11.0, 13.0);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    const auto merged = Synthetic(cell, "P 21 21 21", nullptr, 0.0, 0.8);
    const auto r = AnalyzeTranslationalNCS(merged, cell, sg);
    CHECK_FALSE(r.measurable);
    CHECK_FALSE(r.detected);
    CHECK_FALSE(r.refusal.empty());
    CHECK(TranslationalNCSToText(r).find("not a statement") != std::string::npos);
}

// P1 - no symmetry to expand with, and no centric reflections. The commonest case for a merge the
// space-group search declined to promote, so it must work rather than divide by zero.
TEST_CASE("P1 data are handled with no space group at all", "[tncs]") {
    const gemmi::UnitCell cell = Cell(70.0, 85.0, 110.0);
    const std::array<double, 3> u{0.5, 0.5, 0.5};
    const auto merged = Synthetic(cell, "P 1", &u, kRealisticDepth);
    const auto r = AnalyzeTranslationalNCS(merged, cell, nullptr);
    REQUIRE(r.measurable);
    CHECK(r.detected);
}

// The failure this detector would be useless without. A C-centred crystal MERGED IN P1 - which is
// what --mode scale does on a file with no space group, and what the pre-promotion merge is - still
// carries its centring vector, and the Patterson there is as high as at the origin. With no space
// group to name the centring it cannot be excluded from a list, so it has to be recognised for what
// it is: a translation the data are exactly invariant under is a LATTICE vector, not a pseudo-
// symmetry. Measured before this was handled, 11 of 95 P1 merges in the corpus were falsely called.
TEST_CASE("A centred lattice merged in P1 is reported as a lattice translation, not tNCS", "[tncs]") {
    const gemmi::UnitCell cell = Cell(120.0, 75.0, 90.0);
    // Generated with C-centring (so the absences are real), then handed over as if it were P1.
    const auto merged = Synthetic(cell, "C 1 2 1", nullptr, 0.0);
    const auto r = AnalyzeTranslationalNCS(merged, cell, nullptr);
    REQUIRE(r.measurable);
    REQUIRE_FALSE(r.undeclared_lattice_translations.empty());
    const auto &t = r.undeclared_lattice_translations.front();
    CHECK(std::fabs(t[0] - 0.5) < 0.03);
    CHECK(std::fabs(t[1] - 0.5) < 0.03);
    CHECK_FALSE(r.detected);
    CHECK(TranslationalNCSToText(r).find("LATTICE translation") != std::string::npos);
}

// ...and a real pseudo-translation hiding underneath that centring vector must still be found, or
// the guard above would have traded one blind spot for another.
TEST_CASE("A pseudo-translation under an undeclared centring is still found", "[tncs]") {
    const gemmi::UnitCell cell = Cell(120.0, 75.0, 90.0);
    const std::array<double, 3> u{0.0, 0.0, 0.5};
    const auto merged = Synthetic(cell, "C 1 2 1", &u, kRealisticDepth);
    const auto r = AnalyzeTranslationalNCS(merged, cell, nullptr);
    REQUIRE(r.measurable);
    REQUIRE_FALSE(r.undeclared_lattice_translations.empty());
    INFO("peak " << r.peak_percent << "% u " << r.vector_frac[0] << " " << r.vector_frac[1] << " "
                 << r.vector_frac[2]);
    CHECK(r.detected);
    CHECK(std::fabs(r.vector_frac[2] - 0.5) < 0.05);
}

// The other end of the same scale: a translation the data are EXACTLY invariant under - the weak
// class extinct rather than weak - is a lattice vector however it arose, and must be reported as one
// even when the space group declares no centring at all.
TEST_CASE("An exact translation is a lattice vector, not a strong pseudo-symmetry", "[tncs]") {
    const gemmi::UnitCell cell = Cell(70.0, 85.0, 110.0);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21");
    const std::array<double, 3> u{0.5, 0.5, 0.5};
    const auto merged = Synthetic(cell, "P 21 21 21", &u, 1.0);   // weak class exactly zero
    const auto r = AnalyzeTranslationalNCS(merged, cell, sg);
    REQUIRE(r.measurable);
    REQUIRE_FALSE(r.undeclared_lattice_translations.empty());
    CHECK_FALSE(r.detected);
}

// The most valuable thing this detector can say is not "there is a pseudo-symmetry" but "your cell
// is twice the size it should be". That case does NOT land on a simple fraction of the reported
// axes when the true sub-cell is triclinic, so it is recognised from the depth of the modulation
// instead: a class that is nearly extinct rather than merely weak is a lattice translation in
// disguise.
TEST_CASE("A near-extinct class is reported as a sub-lattice", "[tncs]") {
    const gemmi::UnitCell cell = Cell(70.0, 85.0, 110.0);
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 1 2 1");
    const std::array<double, 3> u{0.0, 0.425, 0.5};   // no p/q with q <= 6 fits the whole vector
    const auto merged = Synthetic(cell, "P 1 2 1", &u, 0.96);   // weak class at 2% of the strong
    const auto r = AnalyzeTranslationalNCS(merged, cell, sg);
    REQUIRE(r.measurable);
    REQUIRE(r.detected);
    CHECK_FALSE(r.commensurate);
    CHECK(r.near_extinct_class);
    CHECK(TranslationalNCSToText(r).find("supercell candidate") != std::string::npos);
}