// 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 <map>
#include <numbers>
#include <vector>

#include "../image_analysis/scale_merge/TwinningAnalysis.h"
#include "SyntheticMergedReflections.h"

namespace {
    // A translational pseudo-symmetry u puts the factor |1 + exp(2 pi i h.u)|^2 on every intensity.
    // Modelled here with a depth f so the weak class is suppressed rather than extinguished, which is
    // what a real pair of copies with different orientations gives.
    std::vector<MergedReflection> WithPseudoTranslation(std::vector<MergedReflection> merged,
                                                        const std::array<double, 3> &u, double f) {
        for (auto &r : merged) {
            const double phase = 2.0 * std::numbers::pi * (r.h * u[0] + r.k * u[1] + r.l * u[2]);
            r.I = static_cast<float>(r.I * (1.0 + f * std::cos(phase)));
            r.sigma = static_cast<float>(0.02 * std::fabs(r.I) + 1.0);
        }
        return merged;
    }

    std::vector<MergedReflection> Crystal(double twin_fraction) {
        jfjoch_test::SyntheticMergeParams p;
        p.true_space_group = "P 1 2 1";
        p.twin_supergroup = "P 2 2 2";
        p.twin_fraction = twin_fraction;
        p.d_min_A = 3.0;
        return jfjoch_test::GenerateSyntheticMerged(p);
    }
}

// The L-test pairs a reflection with a partner two steps along an axis, and that choice is
// load-bearing for a reason the code did not state until now: an even step preserves the class of a
// HALF-INTEGER pseudo-translation, so the commonest tNCS leaves <|L|> alone by construction. A
// one-third pseudo-translation along the same axis does not, and it moves <|L|> by enough to change
// the verdict. This is the regression test for both halves of that.
TEST_CASE("L-test partner steps and a pseudo-translation", "[twinning][tncs]") {
    const auto clean = Crystal(0.0);
    const double l_clean = AnalyzeTwinning(clean, nullptr).mean_abs_l;
    REQUIRE(l_clean == Catch::Approx(0.5).margin(0.02));

    SECTION("a half-integer pseudo-translation does not move it") {
        const auto modulated = WithPseudoTranslation(clean, {0.5, 0.0, 0.0}, 0.8);
        const auto r = AnalyzeTwinning(modulated, nullptr);
        CHECK(r.mean_abs_l == Catch::Approx(l_clean).margin(0.005));
    }

    SECTION("a one-third pseudo-translation moves it a long way up") {
        const auto modulated = WithPseudoTranslation(clean, {1.0 / 3.0, 0.0, 0.0}, 0.8);
        const auto r = AnalyzeTwinning(modulated, nullptr);
        CHECK(r.mean_abs_l > l_clean + 0.03);
        CHECK(r.mean_abs_l > 0.50);                        // into the "contradicts a twin" branch
    }

    SECTION("declaring the vector repairs it") {
        const std::array<double, 3> u{1.0 / 3.0, 0.0, 0.0};
        const auto modulated = WithPseudoTranslation(clean, u, 0.8);
        const auto r = AnalyzeTwinning(modulated, nullptr, 20, &u);
        CHECK(r.l_test_tncs_step_restricted);
        CHECK_FALSE(r.l_test_contaminated_by_tncs);
        // The control is the SAME crystal without the pseudo-translation measured with the SAME
        // restricted steps: <|L|> differs slightly between step directions (they are different
        // distances in reciprocal space), so comparing against the unrestricted number would be
        // measuring that instead of the repair.
        const auto control = AnalyzeTwinning(clean, nullptr, 20, &u);
        REQUIRE(control.l_test_tncs_step_restricted);
        CHECK(r.mean_abs_l == Catch::Approx(control.mean_abs_l).margin(0.005));
        CHECK(r.l_test_pairs > 0.9 * AnalyzeTwinning(modulated, nullptr).l_test_pairs);
    }

    SECTION("a vector no step can preserve is declared unreadable, not repaired") {
        const std::array<double, 3> u{1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0};
        const auto modulated = WithPseudoTranslation(clean, u, 0.8);
        const auto r = AnalyzeTwinning(modulated, nullptr, 20, &u);
        CHECK(r.l_test_contaminated_by_tncs);
        CHECK_FALSE(r.l_test_tncs_step_restricted);
        CHECK(r.l_test_pairs > 0);                         // still reported, just not read
    }

    SECTION("declaring a half-integer vector changes nothing") {
        const std::array<double, 3> u{0.5, 0.0, 0.0};
        const auto modulated = WithPseudoTranslation(clean, u, 0.8);
        const auto with = AnalyzeTwinning(modulated, nullptr, 20, &u);
        const auto without = AnalyzeTwinning(modulated, nullptr);
        CHECK(with.mean_abs_l == Catch::Approx(without.mean_abs_l).margin(0.002));
        CHECK_FALSE(with.l_test_contaminated_by_tncs);
    }
}

// The defect the repair exists for: a twinned crystal that also carries a one-third
// pseudo-translation loses its twin call entirely, because the pseudo-symmetry pushes <|L|> up out
// of the twinned range AND pushes the second moment up out of it at the same time.
TEST_CASE("A pseudo-translation can hide a twin, and declaring it restores the call",
          "[twinning][tncs]") {
    const auto twin = Crystal(0.5);
    const auto baseline = AnalyzeTwinning(twin, nullptr);
    REQUIRE(baseline.mean_abs_l < 0.44);                  // an unambiguous twin when nothing masks it
    REQUIRE(baseline.twinning_suspected);

    const std::array<double, 3> u{1.0 / 3.0, 0.0, 0.0};
    const auto masked = WithPseudoTranslation(twin, u, 0.8);

    const auto undeclared = AnalyzeTwinning(masked, nullptr);
    CHECK(undeclared.mean_abs_l > baseline.mean_abs_l + 0.05);
    CHECK_FALSE(undeclared.twinning_suspected);           // the twin call is lost

    const auto declared = AnalyzeTwinning(masked, nullptr, 20, &u);
    // Against the same twin without the pseudo-translation, measured with the same restricted steps.
    const auto control = AnalyzeTwinning(twin, nullptr, 20, &u);
    CHECK(declared.mean_abs_l == Catch::Approx(control.mean_abs_l).margin(0.005));
    CHECK(declared.mean_abs_l < 0.44);
    CHECK(declared.twinning_suspected);                   // and restored
}

namespace {
    // The merge of a synthetic crystal in `group`: one row per reflection of that group's asymmetric
    // unit, the average of the P1 intensities of its orbit - which is what merging under the group
    // does, whether or not its operators are real.
    std::vector<MergedReflection> MergedIn(const std::vector<MergedReflection> &p1, const char *group) {
        const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(group);
        const gemmi::GroupOps gops = sg->operations();
        const gemmi::ReciprocalAsu asu(sg);
        std::map<std::array<int, 3>, std::pair<MergedReflection, int>> sum;
        for (const auto &r : p1) {
            const auto key = asu.to_asu(gemmi::Op::Miller{{r.h, r.k, r.l}}, gops).first;
            auto [it, fresh] = sum.try_emplace({key[0], key[1], key[2]}, r, 0);
            if (!fresh)
                it->second.first.I += r.I;
            it->second.second += 1;
        }
        std::vector<MergedReflection> out;
        for (auto &[key, v] : sum) {
            MergedReflection r = v.first;
            r.h = key[0];
            r.k = key[1];
            r.l = key[2];
            r.I /= static_cast<float>(v.second);
            out.push_back(r);
        }
        return out;
    }

    std::vector<MergedReflection> Tetragonal(const char *true_group, double twin_fraction) {
        jfjoch_test::SyntheticMergeParams p;
        p.true_space_group = true_group;
        p.twin_supergroup = "P 4 2 2";
        p.twin_fraction = twin_fraction;
        p.d_min_A = 2.5;
        return jfjoch_test::GenerateSyntheticMerged(p);
    }
}

// The L-test compares two intensities as if they had the same expected value, and two index steps are
// not the same resolution: on a small cell with a steep fall-off the raw pair differs systematically,
// which reads as an untwinned crystal being "more untwinned" than 0.5. The shell normalisation removes
// that, and the fraction quoted is the L-test's own.
TEST_CASE("L-test on shell-normalised intensities", "[twinning]") {
    jfjoch_test::SyntheticMergeParams p;
    p.true_space_group = "P 1";
    p.twin_supergroup = "P 1";
    p.wilson_b_A2 = 60.0;
    p.d_min_A = 2.5;
    const auto r = AnalyzeTwinning(jfjoch_test::GenerateSyntheticMerged(p), nullptr);
    INFO("<|L|> " << r.mean_abs_l);
    CHECK(r.mean_abs_l == Catch::Approx(0.5).margin(0.02));
    CHECK_FALSE(r.twinning_suspected);

    const auto twin = AnalyzeTwinning(Crystal(0.2), nullptr);
    CHECK(twin.twin_fraction_source == TwinFractionSource::LTest);
    CHECK(twin.estimated_twin_fraction == Catch::Approx(0.2).margin(0.05));
}

// A P1 merge of a centred lattice holds the reflections the centring extinguishes. Told the centring,
// the test leaves them out and reads the crystal; not told, it pairs present with absent ones.
TEST_CASE("L-test on a P1 merge of a centred lattice", "[twinning]") {
    jfjoch_test::SyntheticMergeParams p;
    p.true_space_group = "R 3 :H";
    p.twin_supergroup = "R 3 :H";
    p.d_min_A = 3.0;
    const auto merged = jfjoch_test::GenerateSyntheticMerged(p);
    const auto told = AnalyzeTwinning(merged, nullptr, 20, nullptr, 'R');
    CHECK(told.mean_abs_l == Catch::Approx(0.5).margin(0.02));
    const auto not_told = AnalyzeTwinning(merged, nullptr);
    CHECK(std::fabs(not_told.mean_abs_l - 0.5) > 0.05);
}

// A genuine holohedral crystal: no twin law exists, and its operators read as real.
TEST_CASE("No twin is called in a genuine holohedral Laue class", "[twinning]") {
    const auto merged = MergedIn(Tetragonal("P 4 2 2", 0.0), "P 4 2 2");
    const auto r = AnalyzeTwinning(merged, gemmi::find_spacegroup_by_name("P 4 2 2"));
    CHECK_FALSE(r.merohedral_twinning_possible);
    CHECK(r.mean_abs_l == Catch::Approx(0.5).margin(0.03));
    CHECK_FALSE(r.twinning_suspected);
    CHECK_FALSE(r.adopted_operators_suspect);
}

// Merging under a false operator averages I(h) with I(Th), and (I(h) + I(Th))/2 has the perfect-twin
// distribution whatever the twin fraction - so a subgroup crystal merged in its lattice's holohedry
// reads <|L|> ~0.375 at every fraction, and the holohedral class must say so rather than "no twin
// law exists".
TEST_CASE("A holohedral merge under a false operator is called suspect", "[twinning]") {
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 4 2 2");
    for (double alpha : {0.0, 0.2, 0.5}) {
        const auto r = AnalyzeTwinning(MergedIn(Tetragonal("P 4", alpha), "P 4 2 2"), sg);
        CAPTURE(alpha);
        CHECK(r.mean_abs_l == Catch::Approx(0.375).margin(0.02));
        CHECK(r.adopted_operators_suspect);
        CHECK(r.twin_fraction_source == TwinFractionSource::None);   // lost in the merge
        CHECK(TwinningVerdictLine(r).rfind("SYMMETRY SUSPECT", 0) == 0);
    }
}