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

namespace {
    // A P1 merge whose centric reflections are centric: SyntheticMergedReflections draws every
    // intensity from the acentric distribution, which is all the searches it was written for need, but
    // the zone test reads exactly the difference. E^2 is chi^2 with one degree of freedom for a
    // reflection centric in the true group (a squared Box-Muller normal), exponential otherwise, keyed
    // on the asymmetric-unit index so the true group's symmetry is exact. `anisotropy_b` is an extra
    // Debye-Waller B along c* alone (A^2) - the same in both twin domains, a twin law being a lattice
    // operation - and `scale_jitter` a log-normal factor of that width on every observed intensity,
    // a nuisance with no direction that no normalisation removes.
    std::vector<MergedReflection> WilsonTetragonal(const char *true_group, double twin_fraction,
                                                   double anisotropy_b = 0.0, double scale_jitter = 0.0) {
        const gemmi::SpaceGroup &sub = gemmi::get_spacegroup_by_name(true_group);
        const gemmi::SpaceGroup &super = gemmi::get_spacegroup_by_name("P 4 2 2");
        const gemmi::Op twin = jfjoch_test::TwinLaw(sub, super);
        const gemmi::GroupOps gops = sub.operations();
        const gemmi::ReciprocalAsu rasu(&sub);
        const gemmi::UnitCell cell(47, 47, 63, 90, 90, 90);
        auto true_intensity = [&](const gemmi::Op::Miller &hkl) {
            const auto asu = rasu.to_asu(hkl, gops).first;
            const double u1 = jfjoch_test::detail::UniformFromHkl(asu);
            const double u2 = jfjoch_test::detail::UniformFromHkl({{asu[0] + 300, asu[1] + 300, asu[2] + 300}});
            const double c = std::cos(2.0 * std::numbers::pi * u2);
            const double e2 = gops.is_reflection_centric(hkl) ? -2.0 * std::log(u1) * c * c : -std::log(u1);
            const double s_c = hkl[2] / cell.c;
            return 1000.0 * e2 * std::exp(-20.0 / (2.0 * std::pow(cell.calculate_d(hkl), 2)))
                   * std::exp(-anisotropy_b * s_c * s_c / 2.0);
        };
        std::vector<MergedReflection> out;
        for (int h = -19; h <= 19; ++h)
            for (int k = -19; k <= 19; ++k)
                for (int l = -26; l <= 26; ++l) {
                    if (std::make_tuple(h, k, l) <= std::make_tuple(-h, -k, -l))
                        continue;
                    const gemmi::Op::Miller hkl{{h, k, l}};
                    const double d = cell.calculate_d(hkl);
                    if (d < 2.5)
                        continue;
                    MergedReflection r;
                    r.h = h;
                    r.k = k;
                    r.l = l;
                    r.d = static_cast<float>(d);
                    const double u3 = jfjoch_test::detail::UniformFromHkl({{h + 700, k + 700, l + 700}});
                    const double u4 = jfjoch_test::detail::UniformFromHkl({{h + 900, k + 900, l + 900}});
                    const double jitter = std::exp(scale_jitter * std::sqrt(-2.0 * std::log(u3))
                                                   * std::cos(2.0 * std::numbers::pi * u4));
                    r.I = static_cast<float>(((1.0 - twin_fraction) * true_intensity(hkl)
                                              + twin_fraction * true_intensity(twin.apply_to_hkl(hkl))) * jitter);
                    r.sigma = static_cast<float>(0.01 * r.I + 1.0);
                    out.push_back(r);
                }
        return out;
    }
}

// Reflections centric in the adopted group but acentric in a subgroup are their own twin mates under
// the operators the group adds. A perfect twin makes every operator statistic read "real", but these
// stay acentric under it and are centric only where the added operators are real.
TEST_CASE("Twin-immune zones tell a perfect twin from real symmetry", "[twinning]") {
    const gemmi::SpaceGroup &p422 = gemmi::get_spacegroup_by_name("P 4 2 2");
    const gemmi::UnitCell tetragonal(47, 47, 63, 90, 90, 90);

    const auto real = AnalyzeTwinImmuneZones(WilsonTetragonal("P 4 2 2", 0.0), tetragonal, p422);
    REQUIRE(real.zones.size() == 3);                       // over 4, over 222, over the diagonal 222
    for (const auto &z : real.zones) {
        CAPTURE(z.operators);
        CHECK(z.mean_abs_e2_minus_1 == Catch::Approx(0.968).margin(0.08));
        CHECK(z.evidence_nats > 0.0);
    }
    CHECK(real.control.mean_abs_e2_minus_1 == Catch::Approx(0.736).margin(0.05));

    // A perfect P4 twin: whichever subgroup is taken as the parent, the operators added over it are
    // not all real, and the zone reads acentric.
    const auto twin = AnalyzeTwinImmuneZones(WilsonTetragonal("P 4", 0.5), tetragonal, p422);
    REQUIRE(twin.zones.size() == 3);
    for (const auto &z : twin.zones) {
        CAPTURE(z.operators);
        CHECK(z.mean_abs_e2_minus_1 == Catch::Approx(0.736).margin(0.08));
        CHECK(z.evidence_nats < 0.0);
        CHECK(z.calibrated_evidence_nats < 0.0);
    }
    // Wilson data normalise cleanly: the control reads the acentric expectation and the
    // calibration takes nothing off.
    CHECK(real.control_excess_per_reflection < 0.01);
    CHECK(twin.control_excess_per_reflection < 0.01);
}

// The zone is read absolutely, so it is only as good as the normalisation. A strong anisotropy
// leaves every class of an isotropically normalised merge reading centric, the acentric control
// included, and a zone - a plane - more so than the control's sphere; it is fitted and taken out
// before anything is read. What no normalisation removes, the control certifies: its excess over the
// acentric expectation is taken off every zone, so a partial twin's zone reads acentric and a real
// operator's centric, as they do on clean data.
TEST_CASE("Twin-immune zones are normalised for anisotropy and calibrated by the acentric control",
          "[twinning]") {
    const gemmi::SpaceGroup &p422 = gemmi::get_spacegroup_by_name("P 4 2 2");
    const gemmi::SpaceGroup &p4 = gemmi::get_spacegroup_by_name("P 4");
    const gemmi::UnitCell tetragonal(47, 47, 63, 90, 90, 90);

    SECTION("a 60 A^2 anisotropy along c* is taken out") {
        const auto twin = AnalyzeTwinImmuneZone(WilsonTetragonal("P 4", 0.2, 60.0), tetragonal, p422, p4);
        REQUIRE(twin.zones.size() == 1);
        CHECK(twin.anisotropy_delta_b_A2 == Catch::Approx(60.0).margin(15.0));
        CHECK(twin.control.mean_abs_e2_minus_1 < 0.76);    // a partial twin's control: at or below acentric
        CHECK(twin.zones[0].calibrated_evidence_nats < -20.0);

        const auto real = AnalyzeTwinImmuneZone(WilsonTetragonal("P 4 2 2", 0.0, 60.0), tetragonal, p422, p4);
        REQUIRE(real.zones.size() == 1);
        CHECK(real.control.mean_abs_e2_minus_1 == Catch::Approx(0.736).margin(0.05));
        CHECK(real.zones[0].calibrated_evidence_nats > 20.0);
    }

    SECTION("a scale jitter no normalisation removes is calibrated off") {
        const auto real = AnalyzeTwinImmuneZone(WilsonTetragonal("P 4 2 2", 0.0, 0.0, 0.6), tetragonal, p422, p4);
        REQUIRE(real.zones.size() == 1);
        CHECK(real.control.mean_abs_e2_minus_1 > 0.80);    // inflated: the normalisation is wrong
        CHECK(real.control_excess_per_reflection > 0.02);
        CHECK(real.zones[0].calibrated_evidence_nats > 20.0);

        // The twin's control is deflated by the twinning as much as the jitter inflates it, so
        // nothing is taken off - and its zone, acentric, still reads so.
        const auto twin = AnalyzeTwinImmuneZone(WilsonTetragonal("P 4", 0.2, 0.0, 0.6), tetragonal, p422, p4);
        REQUIRE(twin.zones.size() == 1);
        CHECK(twin.zones[0].calibrated_evidence_nats < -20.0);
    }
}