// 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 "../image_analysis/scale_merge/SearchSpaceGroup.h"
#include "SyntheticMergedReflections.h"
#include "gemmi/symmetry.hpp"

#include <cmath>
#include <iomanip>
#include <iostream>
#include <limits>
#include <optional>
#include <set>
#include <sstream>
#include <string>
#include <vector>

// Point-group decision on MEROHEDRALLY TWINNED data.
//
// SearchSpaceGroupTest.cpp exercises Stage B (systematic absences) on noise-free, exactly-symmetric
// intensities. The decision that actually goes wrong on real crystals is Stage A: whether the extra
// operator of the metric holohedry is a real symmetry or a twin law. That decision runs through the
// merge chi^2 gate, its systematic-b rescue and the systematic-b veto - none of which the noise-free
// set can reach, because it has no errors for a chi^2 to be reduced by.
//
// Three crystals are modelled per lattice, all with the SAME metric symmetry (a merohedral twin has
// the supergroup's metric, so the lattice cannot arbitrate - only the intensities can):
//   * genuine supergroup  - structure factors invariant under the supergroup; must be promoted;
//   * untwinned subgroup  - structure factors invariant under the subgroup only; must NOT be;
//   * twinned subgroup    - the same crystal at twin fraction alpha; must NOT be promoted for any
//                           0 <= alpha < 0.5, because promoting averages the two twin domains into
//                           one intensity and the twin is then unrecoverable downstream.
// At alpha = 0.5 the twin is physically indistinguishable from real symmetry, so only "terminates
// and returns one of the two" is asserted.
//
// Each is measured through two merge-quality regimes and five merge multiplicities. Neither changes
// any physics - they change only how well the same crystal was measured and how honest its sigmas
// are - so no decision above may move with them. That is the property the harness exists to pin
// down; every case below asserts it outright.

namespace {
    using jfjoch_test::SyntheticMergeParams;

    struct TwinCrystal {
        std::string name;
        std::string sub;                // the crystal's true space group when twinned
        std::string super;              // supergroup of index 2; its extra operator is the twin law
        gemmi::CrystalSystem system;    // metric (lattice) symmetry, as rugnux passes it from indexing
        // The OTHER maximal subgroup of the supergroup of the same order as `sub`, when one exists.
        // 422 has two - 4 and 222 - and only one of them is the crystal. Which one a parent-normalised
        // statistic divides by decides the promotion, so the harness reports both; empty when the
        // supergroup has only one maximal subgroup of that order (32 over 3) and the choice cannot arise.
        std::string rival_parent;
    };

    const std::vector<TwinCrystal> crystals = {
        {"trigonal 3 -> 32 (R3 / R32, twin law k,h,-l)", "R 3 :H", "R 32 :H",
         gemmi::CrystalSystem::Trigonal, ""},
        {"tetragonal 4 -> 422 (P4 / P422, twin law h,-k,-l)", "P 4", "P 4 2 2",
         gemmi::CrystalSystem::Tetragonal, "P 2 2 2"},
    };

    // How honest the merged sigmas are. Both regimes are ways a fitted error model misses in
    // practice, and neither is a property of the crystal's symmetry.
    struct MergeQuality {
        std::string name;
        double sigma_miscalibration;
        double error_model_b;
        std::optional<double> true_systematic_b;
    };

    const std::vector<MergeQuality> merge_quality = {
        // The usual case: merged sigmas come out ~1.7x too small across the board, with the error
        // model's b matching the systematic scatter that is actually there.
        {"sigmas 1.7x too small", 1.7, 0.05, std::nullopt},
        // The other way a fitted error model misses: the statistical sigmas come out somewhat too
        // LARGE while b - the asymptotic I/sigma ceiling, ISa = 1/b - is fitted 3x too optimistic, so
        // the systematic scatter present is 3x what the merged sigmas' floor admits.
        {"ISa 3x too optimistic", 0.6, 0.02, 0.06},
    };

    const std::vector<int> multiplicities = {2, 3, 6, 9, 18};
    const std::vector<double> twin_fractions = {0.0, 0.05, 0.10, 0.20, 0.35, 0.50};

    struct Decision {
        std::string point_group;
        std::string space_group;
        size_t n_merged = 0;
        std::vector<SpaceGroupOperatorScore> operators;
        std::string report;
    };

    // true_group is the symmetry the structure factors have: the subgroup for the twin series, the
    // supergroup for the genuine-high-symmetry control (where the twin law is a real symmetry
    // operator, so the twin fraction has no effect).
    Decision Decide(const TwinCrystal& c, const MergeQuality& q,
                    const std::string& true_group, double alpha, int multiplicity) {
        SyntheticMergeParams p;
        p.true_space_group = true_group;
        p.twin_supergroup = c.super;
        p.twin_fraction = alpha;
        p.multiplicity = multiplicity;
        p.sigma_miscalibration = q.sigma_miscalibration;
        p.error_model_b = q.error_model_b;
        p.true_systematic_b = q.true_systematic_b;
        // One fixed seed for the whole harness: every case draws the same unit-normal stream, so two
        // cases differ only in the knob under test and every test is reproducible.
        p.seed = 20260727;

        const auto merged = jfjoch_test::GenerateSyntheticMerged(p);
        REQUIRE(merged.size() > 5000);            // a realistic dataset, not a handful of reflections

        SearchSpaceGroupOptions opt;              // as rugnux/Rugnux.cpp sets it
        opt.merge_friedel = true;
        opt.lattice_system = c.system;

        const auto result = SearchSpaceGroup(merged, opt);

        Decision d;
        d.point_group = result.point_group_hm;
        d.space_group = result.best_space_group.has_value()
                        ? result.best_space_group->short_name() : "none";
        d.n_merged = merged.size();
        d.operators = result.operator_scores;
        d.report = SearchSpaceGroupResultToText(result);
        return d;
    }

    std::string PointGroupOf(const std::string& space_group_name) {
        return gemmi::get_spacegroup_by_name(space_group_name).point_group_hm();
    }

    std::string ShortNameOf(const std::string& space_group_name) {
        return gemmi::get_spacegroup_by_name(space_group_name).short_name();
    }

    // The hkl triplets SearchSpaceGroup labels a space group's own rotations with. Lets a test tell the
    // crystal's real symmetry operators from the twin law among result.operator_scores.
    std::set<std::string> OperatorTripletsOf(const std::string& space_group_name) {
        std::set<std::string> out;
        const auto& sg = gemmi::get_spacegroup_by_name(space_group_name);
        for (const auto& op : sg.operations().derive_symmorphic().sym_ops) {
            if (op.rot == gemmi::Op::identity().rot)
                continue;
            out.insert(gemmi::Op{op.rot, {0, 0, 0}, op.notation}.as_hkl().triplet('h'));
        }
        return out;
    }

    // The statistic the promotion is actually decided on: the mean operator disagreement
    // H = <|I1-I2|/(I1+I2)> over the operators the promotion ADDS, divided by the mean over the parent
    // group's own operators, measured on the same reflections. Mirrors what SearchSpaceGroup computes
    // for the sub -> super step, so a test can report the margin the max_operator_h_ratio bound has.
    //
    // The parent normalisation is the whole design, not a detail. An ABSOLUTE per-operator bound cannot
    // work: two reflections related by a real symmetry operator still disagree, because they carry
    // DIFFERENT systematic error - absorption, illumination, partiality - and how much of that a
    // crystal has is a property of the measurement, not of its symmetry. So a genuine operator's own
    // disagreement ranges over whatever the data quality happens to be, and any fixed bound placed on
    // it rejects good crystals at one end or waves twins through at the other. Dividing by the parent
    // operators - already confirmed, measured on the same reflections, carrying the same systematic
    // floor - cancels the data quality and leaves only the question being asked: does the ADDED
    // operator relate intensities as equal as the parent's do (real symmetry), or systematically less
    // equal (a twin law mixing non-equivalent reflections)?
    double HRatioOfPromotion(const std::vector<SpaceGroupOperatorScore>& operators,
                             const std::string& parent_group) {
        const auto parent_ops = OperatorTripletsOf(parent_group);
        double h_added = 0.0, h_parent = 0.0;
        int n_added = 0, n_parent = 0;
        for (const auto& s : operators) {
            if (s.n_pairs < 200)                    // SearchSpaceGroupOptions::min_pairs_for_h
                continue;
            if (parent_ops.count(s.op_triplet_hkl) > 0) { h_parent += s.h_stat; ++n_parent; }
            else                                        { h_added  += s.h_stat; ++n_added; }
        }
        if (n_added == 0 || n_parent == 0 || h_parent <= 0.0)
            return std::numeric_limits<double>::quiet_NaN();
        return (h_added / n_added) / (h_parent / n_parent);
    }

    std::string Describe(const TwinCrystal& c, const MergeQuality& q, double alpha, int multiplicity) {
        std::ostringstream os;
        os << c.name << ", " << q.name << ", twin fraction " << std::fixed << std::setprecision(2)
           << alpha << ", multiplicity " << multiplicity;
        return os.str();
    }
}

// Positive control: a crystal whose structure factors really do have the higher symmetry must be
// promoted to it. Guards the twin tests below against a criterion that simply never promotes.
TEST_CASE("SearchSpaceGroup promotes a genuinely high-symmetry crystal",
          "[SearchSpaceGroup][twin]") {
    for (const auto& c : crystals)
        for (const auto& q : merge_quality)
            for (int mult : multiplicities) {
                DYNAMIC_SECTION(c.name << ", " << q.name << ", genuine supergroup, multiplicity " << mult) {
                    const auto d = Decide(c, q, c.super, 0.0, mult);
                    INFO(d.report);
                    CHECK(d.point_group == PointGroupOf(c.super));
                    CHECK(d.space_group == ShortNameOf(c.super));
                }
            }
}

// Negative control: an UNTWINNED crystal of the true subgroup (alpha = 0) must not be promoted - its
// extra metric operator relates reflections that are simply not equivalent.
TEST_CASE("SearchSpaceGroup keeps an untwinned low-symmetry crystal in its subgroup",
          "[SearchSpaceGroup][twin]") {
    for (const auto& c : crystals)
        for (const auto& q : merge_quality)
            for (int mult : multiplicities) {
                DYNAMIC_SECTION(Describe(c, q, 0.0, mult)) {
                    const auto d = Decide(c, q, c.sub, 0.0, mult);
                    INFO(d.report);
                    CHECK(d.point_group == PointGroupOf(c.sub));
                    CHECK(d.space_group == ShortNameOf(c.sub));
                }
            }
}

// A partial merohedral twin must stay in its true subgroup. Promoting it averages the two twin
// domains into one intensity, which no later step can undo: the twin fraction is not recoverable and
// the merged data are simply wrong.
TEST_CASE("SearchSpaceGroup keeps a partially twinned crystal in its true subgroup",
          "[SearchSpaceGroup][twin]") {
    for (const auto& c : crystals)
        for (const auto& q : merge_quality)
            for (double alpha : {0.05, 0.10, 0.20, 0.35})
                for (int mult : multiplicities) {
                    DYNAMIC_SECTION(Describe(c, q, alpha, mult)) {
                        const auto d = Decide(c, q, c.sub, alpha, mult);
                        INFO(d.report);
                        CHECK(d.point_group == PointGroupOf(c.sub));
                        CHECK(d.space_group == ShortNameOf(c.sub));
                    }
                }
}

// A PERFECT (alpha = 0.5) merohedral twin produces intensities that are exactly invariant under the
// twin law: I_obs(h) = I_obs(twin h) for every reflection. No intensity statistic can tell it from a
// crystal that genuinely has the higher symmetry - the information is not in the data (it takes a
// different measurement, e.g. the |E| distribution's second moment, to even suspect it). So the only
// thing asserted here is that the search terminates and returns one of the two.
TEST_CASE("SearchSpaceGroup on a perfect merohedral twin returns one of the two symmetries",
          "[SearchSpaceGroup][twin]") {
    for (const auto& c : crystals)
        for (const auto& q : merge_quality)
            for (int mult : multiplicities) {
                DYNAMIC_SECTION(Describe(c, q, 0.5, mult)) {
                    const auto d = Decide(c, q, c.sub, 0.5, mult);
                    INFO(d.report);
                    REQUIRE(d.space_group != "none");
                    CHECK((d.point_group == PointGroupOf(c.sub) ||
                           d.point_group == PointGroupOf(c.super)));
                }
            }
}

// THE property this harness exists for. Multiplicity changes only the sigmas - the random part of a
// merged sigma averages down as 1/sqrt(n) while the systematic floor b*|I| does not - so it changes
// how well the SAME crystal is measured, never what its symmetry is. A symmetry decision that moves
// when the same crystal is merged 2x instead of 18x is a defect of the criterion, not a property of
// the data.
//
// This is what a criterion thresholded on merge chi^2 or on a systematic-b RATIO cannot deliver: both
// are ratios to an error model that multiplicity and the sigma calibration move, so the tetragonal
// 4 -> 422 twin at alpha 0.20 / 0.35 used to flip - promoted at multiplicity 2 and kept at 18 with
// under-calibrated sigmas, and the other way round with an over-optimistic ISa. The operator
// disagreement ratio H_added/H_parent holds instead because there is no sigma in it at all: it
// compares intensities with intensities, and normalising against the parent group's own operators on
// the same reflections divides out both the data quality and the systematic floor that multiplicity
// and the error model move. An absolute bound on a single operator's H would not survive this - see
// HRatioOfPromotion.
TEST_CASE("SearchSpaceGroup point-group decision does not depend on merge multiplicity",
          "[SearchSpaceGroup][twin]") {
    for (const auto& c : crystals)
        for (const auto& q : merge_quality)
            for (double alpha : twin_fractions) {
                DYNAMIC_SECTION(Describe(c, q, alpha, 2) + " vs multiplicity 18") {
                    const auto low = Decide(c, q, c.sub, alpha, 2);
                    const auto high = Decide(c, q, c.sub, alpha, 18);
                    INFO("multiplicity 2:\n" << low.report << "\nmultiplicity 18:\n" << high.report);
                    CHECK(low.point_group == high.point_group);
                    CHECK(low.space_group == high.space_group);
                }
            }
}

// Diagnostic, not run by default:  ./jfjoch_test "[twin-h]"
// Prints the operator-disagreement ratio H_added/H_parent that the promotion is decided on, for the
// genuine high-symmetry crystal and for each twin fraction, across both merge-quality regimes and
// every multiplicity - i.e. how much margin the max_operator_h_ratio bound actually has, and whether
// either side of it drifts with data quality or data amount.
TEST_CASE("SearchSpaceGroup operator H ratio margins", "[.][twin-h]") {
    SearchSpaceGroupOptions defaults;
    std::cout << "H_added / H_parent for the sub -> super promotion; bound "
              << defaults.max_operator_h_ratio << " (above = refused as a twin)\n";
    for (const auto& c : crystals)
        for (const auto& q : merge_quality)
            // Both normalisations where the supergroup has two maximal subgroups of the same order:
            // against the crystal's true parent, and against its rival.
            for (const auto& parent : c.rival_parent.empty()
                                      ? std::vector<std::string>{c.sub}
                                      : std::vector<std::string>{c.sub, c.rival_parent}) {
                std::cout << "\n" << c.name << "\n  merge quality: " << q.name
                          << "   normalised against " << ShortNameOf(parent)
                          << (parent == c.sub ? " (the crystal's own parent)" : " (the RIVAL parent)")
                          << "\n";
                std::cout << "  " << std::setw(22) << std::left << "true symmetry / alpha" << std::right;
                for (int mult : multiplicities)
                    std::cout << std::setw(12) << ("mult " + std::to_string(mult));
                std::cout << "\n  " << std::setw(22) << std::left << "genuine supergroup" << std::right;
                for (int mult : multiplicities)
                    std::cout << std::setw(12) << std::fixed << std::setprecision(3)
                              << HRatioOfPromotion(Decide(c, q, c.super, 0.0, mult).operators, parent);
                std::cout << "\n";
                for (double alpha : twin_fractions) {
                    std::ostringstream label;
                    label << "subgroup, alpha " << std::fixed << std::setprecision(2) << alpha;
                    std::cout << "  " << std::setw(22) << std::left << label.str() << std::right;
                    for (int mult : multiplicities)
                        std::cout << std::setw(12) << std::fixed << std::setprecision(3)
                                  << HRatioOfPromotion(Decide(c, q, c.sub, alpha, mult).operators, parent);
                    std::cout << "\n";
                }
            }
    SUCCEED();
}

// Diagnostic, not run by default (hidden by the [.] tag):
//   ./jfjoch_test "[twin-table]"
// prints the decision for every (true symmetry, twin fraction, multiplicity) combination in both
// merge-quality regimes - the table a redesign of the point-group criterion should be judged against.
TEST_CASE("SearchSpaceGroup twin decision table", "[.][twin-table]") {
    for (const auto& c : crystals)
        for (const auto& q : merge_quality) {
            const auto reference = Decide(c, q, c.super, 0.0, 6);
            std::cout << "\n" << c.name << "\n  merge quality: " << q.name
                      << "   (true subgroup " << ShortNameOf(c.sub)
                      << ", supergroup " << ShortNameOf(c.super) << ", "
                      << reference.n_merged << " merged reflections)\n";
            std::cout << "  " << std::setw(22) << std::left << "true symmetry / alpha" << std::right;
            for (int mult : multiplicities)
                std::cout << std::setw(12) << ("mult " + std::to_string(mult));
            std::cout << "\n  " << std::setw(22) << std::left << "genuine supergroup" << std::right;
            for (int mult : multiplicities)
                std::cout << std::setw(12) << Decide(c, q, c.super, 0.0, mult).space_group;
            std::cout << "\n";
            for (double alpha : twin_fractions) {
                std::ostringstream label;
                label << "subgroup, alpha " << std::fixed << std::setprecision(2) << alpha;
                std::cout << "  " << std::setw(22) << std::left << label.str() << std::right;
                for (int mult : multiplicities)
                    std::cout << std::setw(12) << Decide(c, q, c.sub, alpha, mult).space_group;
                std::cout << "\n";
            }
        }
    SUCCEED();
}