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

#include <catch2/catch_all.hpp>
#include "../image_analysis/scale_merge/HKLKey.h"
#include "../image_analysis/scale_merge/Merge.h"
#include "gemmi/reciproc.hpp"

TEST_CASE("HKLKey_NoSG_noMergeFriedel") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3) != hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(-1,-2,-3) != hkl_key_gen(1,-2,-3));
}

TEST_CASE("HKLKey_NoSG_MergeFriedel") {
    HKLKeyGenerator hkl_key_gen(true, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3) == hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(-1,-2,-3) != hkl_key_gen(1,-2,-3));
}

TEST_CASE("HKLKey_SG1_MergeFriedel") {
    HKLKeyGenerator hkl_key_gen(true, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3) == hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(-1,-2,-3) != hkl_key_gen(1,-2,-3));
}

TEST_CASE("HKLKey_SG1_NoMergeFriedel") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3) != hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(-1,-2,-3) != hkl_key_gen(1,-2,-3));
}

TEST_CASE("HKLKey_SG96_MergeFriedel") {
    HKLKeyGenerator hkl_key_gen(true, *gemmi::find_spacegroup_by_number(96));
    CHECK(hkl_key_gen(-1, -2, -3) == hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2,1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-1,-2,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2,-1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(1,-2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-1,2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2,1,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2, -1, -3));
    
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2,-1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2, 1, 3));
}

TEST_CASE("HKLKey_SG96_NoMergeFriedel") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(96));
    CHECK(hkl_key_gen(-1, -2, -3) != hkl_key_gen(1,2,3));
    CHECK(hkl_key_gen(-1,-2,-3) == hkl_key_gen(-1,-2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2,1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-1,-2,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2,-1,3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(1,-2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-1,2,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(2,1,-3));
    CHECK(hkl_key_gen(1,2,3) == hkl_key_gen(-2, -1, -3));
    
    CHECK(hkl_key_gen(1,2,3) != hkl_key_gen(-2,-1,3));
    CHECK(hkl_key_gen(1,2,3) != hkl_key_gen(2, 1, 3));
}

TEST_CASE("HKLKey_pack_friedel") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3).pack() != hkl_key_gen(1,2,3).pack());
    CHECK(hkl_key_gen(-1,-2,-3).pack() == hkl_key_gen(-1,-2,-3).pack());
    CHECK(hkl_key_gen(-1,-2,-3).pack() != hkl_key_gen(1,-2,-3).pack());
}

TEST_CASE("HKLKey_pack_no_friedel") {
    HKLKeyGenerator hkl_key_gen(true, *gemmi::find_spacegroup_by_number(1));
    CHECK(hkl_key_gen(-1, -2, -3).pack() == hkl_key_gen(1,2,3).pack());
    CHECK(hkl_key_gen(-1,-2,-3).pack() == hkl_key_gen(-1,-2,-3).pack());
    CHECK(hkl_key_gen(-1,-2,-3).pack() != hkl_key_gen(1,-2,-3).pack());
}


TEST_CASE("HKLKey_sys_absence_P212121") {
    HKLKeyGenerator hkl_key_gen(false, *gemmi::find_spacegroup_by_number(19));
    CHECK(hkl_key_gen.IsSystematicallyAbsent(5,0,0));
    CHECK(!hkl_key_gen.IsSystematicallyAbsent(6,0,0));
    CHECK(hkl_key_gen.IsSystematicallyAbsent(0,5,0));
    CHECK(hkl_key_gen.IsSystematicallyAbsent(0,0,5));
    CHECK(!hkl_key_gen.IsSystematicallyAbsent(0,4,0));
    CHECK(!hkl_key_gen.IsSystematicallyAbsent(5,5,5));
}
TEST_CASE("AcceptReflection_ResolutionLimits") {
    Reflection r{};
    r.I = 100.0f;
    r.sigma = 5.0f;
    r.rlp = 1.0f;
    r.d = 20.0f;

    // No limits: only the finiteness checks apply.
    CHECK(AcceptReflection(r, std::nullopt, std::nullopt));

    // Low-resolution limit rejects anything coarser than the limit, and is exclusive at it.
    CHECK_FALSE(AcceptReflection(r, std::nullopt, std::optional<double>(15.0)));
    CHECK(AcceptReflection(r, std::nullopt, std::optional<double>(20.0)));
    CHECK(AcceptReflection(r, std::nullopt, std::optional<double>(50.0)));

    // High-resolution limit still rejects anything finer, in the same direction as before.
    CHECK_FALSE(AcceptReflection(r, std::optional<double>(25.0), std::nullopt));
    CHECK(AcceptReflection(r, std::optional<double>(2.0), std::optional<double>(50.0)));

    // The plain-double overload treats 0 as "no limit" at both ends.
    CHECK(AcceptReflection(r, 0.0, 0.0));
    CHECK_FALSE(AcceptReflection(r, 0.0, 15.0));
    CHECK(AcceptReflection(r, 2.0, 50.0));
}

// --- Completeness denominator --------------------------------------------------------------------

namespace {
    // A merged set built straight out of the reflections the cell and space group can give, so the
    // test says exactly which of them were measured: every unique reflection between d_min and
    // d_max_measured and none outside. Without Friedel merging an acentric contributes both hands.
    std::vector<MergedReflection> MeasuredBetween(const gemmi::SpaceGroup &sg, const UnitCell &cell,
                                                  double d_min, double d_max_measured,
                                                  bool merge_friedel) {
        const gemmi::UnitCell gemmi_cell = cell;
        const gemmi::GroupOps gops = sg.operations();
        std::vector<MergedReflection> out;
        for (const auto &hkl: gemmi::make_miller_vector(gemmi_cell, &sg, d_min, d_max_measured, true)) {
            MergedReflection r;
            r.h = hkl[0];
            r.k = hkl[1];
            r.l = hkl[2];
            r.d = static_cast<float>(gemmi_cell.calculate_d(hkl));
            r.I = 100.0f;
            r.sigma = 10.0f;
            r.I_half[0] = 100.0f;
            r.I_half[1] = 100.0f;
            out.push_back(r);
            if (!merge_friedel && !gops.is_reflection_centric(hkl)) {
                r.h = -hkl[0];
                r.k = -hkl[1];
                r.l = -hkl[2];
                out.push_back(r);
            }
        }
        return out;
    }

    MergeStatistics StatsWithLowLimit(const gemmi::SpaceGroup &sg, const std::optional<UnitCell> &cell,
                                      const std::vector<MergedReflection> &merged,
                                      std::optional<double> low_limit, bool merge_friedel) {
        DiffractionExperiment x;
        x.SetSpaceGroup(sg);
        ScalingSettings s = x.GetScalingSettings();
        s.LowResolutionLimit_A(low_limit);
        s.MergeFriedel(merge_friedel);
        x.ImportScalingSettings(s);

        MergeOnTheFly merge(x);
        merge.ReferenceCell(cell);
        return merge.MergeStats(merged, {});
    }

    double Completeness(const MergeStatisticsShell &s) {
        return s.possible_unique_reflections > 0
            ? 100.0 * s.unique_reflections / s.possible_unique_reflections : 0.0;
    }

    const gemmi::SpaceGroup &TestSpaceGroup() { return gemmi::get_spacegroup_by_name("P 1 2 1"); }
    constexpr UnitCell TEST_CELL{40, 45, 50, 90, 100, 90};   // synthetic; coarsest reflection ~49 A
}

// The low-resolution terms a beam stop ate must count as missing: the denominator is the declared
// range, so widening the declared range lowers completeness rather than leaving it alone.
TEST_CASE("MergeStats_CompletenessFallsWhenTheLowBoundCrossesAMaskedRegion") {
    const auto &sg = TestSpaceGroup();
    // Nothing coarser than 20 A was measured - it is all behind the stop.
    const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 20.0, true);
    REQUIRE(!merged.empty());

    const auto at_20 = StatsWithLowLimit(sg, TEST_CELL, merged, 20.0, true);
    const auto at_50 = StatsWithLowLimit(sg, TEST_CELL, merged, 50.0, true);

    // Declared exactly where the data stop: everything possible was measured.
    CHECK(Completeness(at_20.overall) > 99.0);

    // Declared out to 50 A: the 20-50 A shell is in the denominator and in nothing else.
    CHECK(at_50.overall.possible_unique_reflections > at_20.overall.possible_unique_reflections);
    CHECK(at_50.overall.unique_reflections == at_20.overall.unique_reflections);
    CHECK(Completeness(at_50.overall) < Completeness(at_20.overall));
    // The innermost shell is where it bites.
    CHECK(Completeness(at_50.shells.front()) < Completeness(at_20.shells.front()));
}

// No low-resolution limit means the whole sphere. The cell has no reflection coarser than 50 A, so
// freeing the 50 A limit must count the same set - the fix does not presuppose either default.
TEST_CASE("MergeStats_CompletenessWithNoDeclaredLowLimit") {
    const auto &sg = TestSpaceGroup();
    const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 20.0, true);

    const auto at_50 = StatsWithLowLimit(sg, TEST_CELL, merged, 50.0, true);
    const auto unlimited = StatsWithLowLimit(sg, TEST_CELL, merged, std::nullopt, true);

    CHECK(unlimited.overall.possible_unique_reflections == at_50.overall.possible_unique_reflections);
    CHECK(Completeness(unlimited.overall) == Catch::Approx(Completeness(at_50.overall)));
    // The shell table stays finite even though the bound is not.
    CHECK(std::isfinite(unlimited.shells.front().d_max));
}

// Counting the two Bijvoet mates of an acentric separately doubles the denominator too, so a fully
// measured anomalous set is 100% complete and not 200%.
TEST_CASE("MergeStats_CompletenessNeverExceeds100") {
    const auto &sg = TestSpaceGroup();
    for (const bool merge_friedel: {true, false}) {
        const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 50.0, merge_friedel);
        const auto stats = StatsWithLowLimit(sg, TEST_CELL, merged, 50.0, merge_friedel);
        INFO("merge_friedel = " << merge_friedel);
        CHECK(Completeness(stats.overall) <= 100.0);
        CHECK(Completeness(stats.overall) > 99.0);
        for (const auto &sh: stats.shells)
            CHECK(Completeness(sh) <= 100.0);
    }
}

// The shell grid and the denominator share their bounds, so every possible reflection lands in a
// shell: the shells sum to the overall, and the overall is the sphere the run declared.
TEST_CASE("MergeStats_PossibleSumsOverTheShellsToTheDeclaredSphere") {
    const auto &sg = TestSpaceGroup();
    const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 20.0, true);
    const auto stats = StatsWithLowLimit(sg, TEST_CELL, merged, 50.0, true);

    int sum = 0;
    for (const auto &sh: stats.shells)
        sum += sh.possible_unique_reflections;
    CHECK(sum == stats.overall.possible_unique_reflections);

    // Counted independently over the same declared range - nothing is lost between the two.
    const gemmi::UnitCell gemmi_cell = TEST_CELL;
    const int expected = gemmi::count_reflections(gemmi_cell, &sg, stats.overall.d_min * 0.999, 50.0, true);
    CHECK(stats.overall.possible_unique_reflections == expected);
}

// Without a reference cell there is no set to count against; completeness stays unmeasured rather
// than becoming a number, with or without a declared low limit.
TEST_CASE("MergeStats_NoReferenceCellLeavesCompletenessUnmeasured") {
    const auto &sg = TestSpaceGroup();
    const auto merged = MeasuredBetween(sg, TEST_CELL, 2.0, 20.0, true);

    for (const std::optional<double> low_limit: {std::optional<double>(50.0), std::optional<double>()}) {
        const auto stats = StatsWithLowLimit(sg, std::nullopt, merged, low_limit, true);
        CHECK(stats.overall.possible_unique_reflections == 0);
        CHECK(stats.overall.unique_reflections > 0);
        CHECK(Completeness(stats.overall) == 0.0);
    }
}