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

#include "TwinningAnalysis.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <iomanip>
#include <limits>
#include <sstream>
#include <unordered_map>
#include <vector>

namespace {
    int64_t PackHKL(int h, int k, int l) {
        constexpr int64_t bias = 1 << 20;          // indices assumed within +/- 2^20
        return ((h + bias) << 42) | ((k + bias) << 21) | (l + bias);
    }

    bool UsableIntensity(const MergedReflection& r) {
        return std::isfinite(r.I) && std::isfinite(r.d) && r.d > 0.0;
    }

    // Merohedral twinning needs a twin law - a lattice symmetry operation that is not a symmetry of the
    // crystal - which exists only when the Laue class is a proper subgroup of the lattice holohedry.
    // The holohedral high-symmetry Laue classes (4/mmm, 6/mmm, m-3m, and -3m on a rhombohedral lattice)
    // admit no such operation, so twinning is geometrically impossible and the intensity statistics
    // cannot be indicating it. Low-symmetry classes stay eligible because pseudo-merohedral twinning
    // through an accidental metric specialisation cannot be excluded from the symmetry alone.
    bool MerohedralTwinningPossible(const gemmi::SpaceGroup* sg) {
        if (!sg)
            return true;                    // P1 / unknown symmetry: cannot rule twinning out
        switch (sg->laue_class()) {
            case gemmi::Laue::L4mmm:        // 4/mmm - tetragonal holohedry
            case gemmi::Laue::L6mmm:        // 6/mmm - hexagonal holohedry
            case gemmi::Laue::Lm3m:         // m-3m  - cubic holohedry
                return false;
            case gemmi::Laue::L3m:          // -3m is holohedral on a rhombohedral (R) lattice, but a
                return sg->hm[0] != 'R';    //   hexagonal-P 32/3m crystal can still twin towards 6/mmm
            default:
                return true;
        }
    }
}

TwinningAnalysisResult AnalyzeTwinning(const std::vector<MergedReflection>& merged,
                                       const gemmi::SpaceGroup* space_group, int resolution_shells,
                                       const std::array<double, 3>* tncs_vector) {
    TwinningAnalysisResult result;
    if (merged.empty())
        return result;

    // Centric reflections follow different statistics and must be excluded. In P1 (no space group)
    // none are centric.
    const gemmi::GroupOps gops = space_group ? space_group->operations() : gemmi::GroupOps{};
    auto acentric = [&](const MergedReflection& r) {
        return !space_group || !gops.is_reflection_centric(gemmi::Op::Miller{{r.h, r.k, r.l}});
    };

    // --- L-test ---
    // Following Padilla & Yeates (2003) Acta Cryst. D59, 1124-1130
    // Pair each reflection with a symmetry-independent neighbour two steps away along an axis (the
    // step of 2 keeps the partner local in resolution while avoiding the reflection itself). The
    // merged reflections are unique in the asymmetric unit, so any other merged reflection is
    // genuinely non-equivalent - exactly the pairing the L-test wants. Only acentric reflections
    // with positive intensity enter, which also keeps L = (I1-I2)/(I1+I2) bounded in [-1, 1].
    std::unordered_map<int64_t, double> intensity;
    intensity.reserve(merged.size() * 2);
    for (const auto& r : merged)
        if (UsableIntensity(r) && r.I > 0.0 && acentric(r))
            intensity.emplace(PackHKL(r.h, r.k, r.l), r.I);

    // Step to a nearby, symmetry-independent partner. The axis step of 2 is load-bearing for TWO
    // reasons, and only the first was ever written down.
    //
    // (1) It preserves the reflection condition of P/I/C/F/A/B lattices (parity-based). See below.
    //
    // (2) It preserves the class of a HALF-INTEGER translational pseudo-symmetry. A pseudo-
    //     translation u multiplies the intensity by |1 + exp(2 pi i h.u)|^2; a partner at h + s
    //     carries the same factor exactly when s.u is an integer, which for an axis step of 2 means
    //     every component of u is 0 or 1/2. The commonest tNCS - a second copy at (1/2,1/2,1/2) or
    //     similar - therefore leaves <|L|> untouched, and the L-test works on it by construction
    //     rather than by luck. Measured on a real half-integer case: <|L|> = 0.4755 with these steps,
    //     0.4880 with parity-breaking steps of 1. Do not "shorten" or "lengthen" these steps without
    //     re-measuring that: a step of 6, which would also survive a one-third pseudo-translation,
    //     puts <|L|> above 0.500 on two thirds of a 137-dataset corpus purely because the partner is
    //     too far away in resolution, which is worse than the problem it solves.
    //
    //     A pseudo-translation that is NOT half-integer is a different matter, and which STEP is used
    //     decides it: only the steps with s.u not an integer break the class. A vector along c alone,
    //     u = (0,0,1/3), is paired by (2,0,0) on almost every reflection and is harmless; a vector
    //     with a one-third component along a breaks (2,0,0) and biases <|L|> hard. Measured on
    //     synthetic data, a PERFECT TWIN carrying a one-third pseudo-translation along a reads
    //     <|L|> = 0.454 instead of 0.377 while its second moment rises from 1.50 to 1.97 - both
    //     indicators destroyed and the twin call lost, with nothing in the report to say why.
    //
    //     So when a pseudo-translation is known, the partner steps are restricted to those that
    //     preserve its class. That is a repair, not a compromise: on the same synthetic twin it
    //     returns <|L|> = 0.376 against 0.377 for the tNCS-free control, at the cost of about 0.4%
    //     of the pairs. When no step qualifies - u = (1/3,1/3,1/3), say, where every one of the five
    //     breaks the class - the L-test simply cannot be measured on this crystal, and the verdict
    //     stops reading it in either direction.
    //
    // The axis step violates R-centring (-h+k+l = 0 mod 3,
    // as 2 != 0 mod 3) -> the partner is systematically absent and rhombohedral crystals yield zero
    // pairs. The diagonal (1,1,0)/(1,1,3) steps preserve the mod-3 condition in BOTH obverse and
    // reverse settings; they are tried only when the axis steps find no present partner, so P/I/C/F
    // behaviour is unchanged (first present partner wins).
    const std::vector<std::array<int, 3>> all_steps{
        {2, 0, 0}, {0, 2, 0}, {0, 0, 2}, {1, 1, 0}, {1, 1, 3}};

    // A partner at h + s carries the same pseudo-translation factor as h exactly when s.u is an
    // integer. Keep the steps for which it is, in the original order so the "first present partner
    // wins" behaviour is unchanged among those that remain.
    std::vector<std::array<int, 3>> steps = all_steps;
    if (tncs_vector != nullptr) {
        std::vector<std::array<int, 3>> preserving;
        for (const auto& s : all_steps) {
            const double t = s[0] * (*tncs_vector)[0] + s[1] * (*tncs_vector)[1]
                             + s[2] * (*tncs_vector)[2];
            if (std::fabs(t - std::round(t)) <= 0.06)
                preserving.push_back(s);
        }
        if (preserving.empty())
            result.l_test_contaminated_by_tncs = true;
        else {
            steps = preserving;
            result.l_test_tncs_step_restricted = true;
        }
    }

    auto run_l_test = [&](const std::vector<std::array<int, 3>>& use) {
        double sum_abs_l = 0.0, sum_l2 = 0.0;
        int n_pairs = 0;
        for (const auto& [key, i1] : intensity) {
            const int h = static_cast<int>((key >> 42) & 0x1FFFFF) - (1 << 20);
            const int k = static_cast<int>((key >> 21) & 0x1FFFFF) - (1 << 20);
            const int l = static_cast<int>(key & 0x1FFFFF) - (1 << 20);
            for (const auto& s : use) {
                const auto it = intensity.find(PackHKL(h + s[0], k + s[1], l + s[2]));
                if (it == intensity.end())
                    continue;
                const double lstat = (i1 - it->second) / (i1 + it->second);
                sum_abs_l += std::fabs(lstat);
                sum_l2 += lstat * lstat;
                ++n_pairs;
                break;                                 // one neighbour per reflection
            }
        }
        result.l_test_pairs = n_pairs;
        result.mean_abs_l = n_pairs > 0 ? sum_abs_l / n_pairs : 0.0;
        result.mean_l_squared = n_pairs > 0 ? sum_l2 / n_pairs : 0.0;
    };

    run_l_test(steps);
    // A restricted step set can leave nothing to pair with - an R-centred lattice whose only present
    // partners are the diagonals, and those are the steps the pseudo-translation broke. Fall back to
    // the full set and mark the statistic unreadable rather than reporting <|L|> from a handful of
    // pairs.
    if (result.l_test_tncs_step_restricted && result.l_test_pairs < 100) {
        result.l_test_tncs_step_restricted = false;
        result.l_test_contaminated_by_tncs = true;
        run_l_test(all_steps);
    }

    // --- Second moment <I^2>/<I>^2 of acentric intensities, normalised per resolution shell ---
    // Binning by 1/d^2 removes the resolution fall-off, so the moment is 2.0 (untwinned) or 1.5
    // (perfect twin) regardless of the overall B-factor. The moment divides by the *square* of the
    // shell-mean intensity, so it is not robust: on weak or mis-integrated data a shell mean can
    // collapse to the noise floor and one outlier reflection then dominates (I/mean)^2 (a single
    // I=158 in a mean~1 shell contributed 78% of a whole dataset's value). To keep this a twinning
    // indicator rather than a data-quality artefact - as phenix.xtriage does - we skip noise-only
    // shells (<I/sigma> below 1) and reject Wilson outliers (E^2 above 8, ~exp(-8) upper tail) with
    // one shell-mean re-iteration so the outlier does not corrupt the normalising mean either.
    constexpr double min_shell_isig = 1.0;     // shells below this are noise, not signal
    constexpr double wilson_outlier_e2 = 8.0;  // reject improbably strong reflections (P ~ e^-8)
    int n_shells = std::max(1, resolution_shells);
    double min_s = std::numeric_limits<double>::infinity();
    double max_s = -std::numeric_limits<double>::infinity();
    for (const auto& r : merged) {
        if (!UsableIntensity(r) || !acentric(r))
            continue;
        const double s = 1.0 / (r.d * r.d);
        min_s = std::min(min_s, s);
        max_s = std::max(max_s, s);
    }
    if (std::isfinite(min_s) && max_s > min_s) {
        auto shell_of = [&](double d) {
            const double t = (1.0 / (d * d) - min_s) / (max_s - min_s);
            return std::min(n_shells - 1, std::max(0, static_cast<int>(t * n_shells)));
        };
        // Group acentric intensities by shell, and accumulate <I/sigma> to gauge each shell's signal.
        std::vector<std::vector<double>> shell_I(n_shells);
        std::vector<double> shell_isig_sum(n_shells, 0.0);
        std::vector<int> shell_isig_n(n_shells, 0);
        for (const auto& r : merged) {
            if (!UsableIntensity(r) || !acentric(r))
                continue;
            const int b = shell_of(r.d);
            shell_I[b].push_back(r.I);
            if (std::isfinite(r.sigma) && r.sigma > 0.0) {
                shell_isig_sum[b] += r.I / r.sigma;
                shell_isig_n[b] += 1;
            }
        }

        double sum_e4 = 0.0;
        int n_moment = 0;
        for (int b = 0; b < n_shells; ++b) {
            const auto& intensities = shell_I[b];
            if (intensities.empty() || shell_isig_n[b] == 0
                || shell_isig_sum[b] / shell_isig_n[b] < min_shell_isig)
                continue;

            double sum = 0.0;
            for (double I : intensities)
                sum += I;
            double mean = sum / intensities.size();
            if (mean <= 0.0)
                continue;
            // Re-fit the mean over the reflections that pass the outlier cut, so the outlier does not
            // inflate the very mean it is measured against.
            sum = 0.0;
            int n_kept = 0;
            for (double I : intensities)
                if (I / mean <= wilson_outlier_e2) {
                    sum += I;
                    ++n_kept;
                }
            if (n_kept == 0)
                continue;
            mean = sum / n_kept;
            if (mean <= 0.0)
                continue;

            for (double I : intensities) {
                const double e2 = I / mean;
                if (e2 > wilson_outlier_e2)
                    continue;
                sum_e4 += e2 * e2;
                ++n_moment;
            }
        }
        result.moment_reflections = n_moment;
        if (n_moment > 0)
            result.second_moment = sum_e4 / n_moment;
    }

    // Twin fraction from the second moment M = 2(1 - a + a^2): a = (1 - sqrt(2M-3))/2.
    if (result.second_moment > 0.0) {
        const double m = std::clamp(result.second_moment, 1.5, 2.0);
        result.estimated_twin_fraction = (1.0 - std::sqrt(std::max(0.0, 2.0 * m - 3.0))) / 2.0;
    }
    // Either indicator dropping clearly below its untwinned value is suspicious - but only where a twin
    // law can actually exist. In a holohedral Laue class (e.g. 422) no merohedral twinning is
    // possible, so a low <|L|> is a statistical artefact (correlated near-neighbours) rather than a
    // twin, and must not be flagged.
    result.merohedral_twinning_possible = MerohedralTwinningPossible(space_group);
    // The two indicators can only move one way under twinning: <|L|> down from 0.500 towards 0.375,
    // the second moment down from 2.0 towards 1.5. A narrow intensity distribution sitting next to an
    // <|L|> at or ABOVE its untwinned value therefore has some other cause, and calling it a twin is
    // the one reading the data rule out. It happens on small-cell rotation data, where a twin fraction
    // of 0.50 was reported for a crystal whose <|L|> was 0.63 - a value a twin cannot produce.
    // A high <|L|> here is NOT evidence of a centrosymmetric structure. Measured across the corpus,
    // the highest <|L|> of all - 0.712, second moment 3.089 - belongs to a small (~10 A) cell in an
    // orthorhombic group whose three 2_1 axes are each proven by absences, no violations against a
    // control class of the same size, so it is chiral. Both numbers sit ABOVE the centric
    // expectations of 0.637 and 3.0, which is the point: no Wilson distribution, centric or acentric,
    // reaches them, so they are not a statement about the structure's symmetry at all. What does put
    // them there is not established here, but it is neither of the obvious guesses - a structure of
    // few atoms gives a second moment of 2 - 1/N, BELOW 2, and a large term common to every structure
    // factor drives it towards 1. Read this statistic as "not a twin", and nothing further.
    //
    // Unless the L-test could not be measured free of a pseudo-translation. Three real crystals
    // carrying a one-third pseudo-translation read <|L|> = 0.527, 0.529 and 0.587, against 0.468-0.491
    // for clean untwinned controls - every one of them trips the >= 0.50 branch and has its twin test
    // vetoed for a reason that has nothing to do with twinning. Where the step restriction above could
    // not repair it, the statistic is dropped from the verdict in BOTH directions: it can no longer
    // say "not a twin", and it can no longer say "a twin" either. The second moment then decides
    // alone, and the text says the L-test could not be read. This REMOVES a veto; it never adds one.
    const bool l_test_usable = result.l_test_pairs > 0 && !result.l_test_contaminated_by_tncs;
    const bool l_test_contradicts_twin = l_test_usable && result.mean_abs_l >= 0.50;
    result.twinning_suspected = result.merohedral_twinning_possible && !l_test_contradicts_twin &&
                                ((l_test_usable && result.mean_abs_l < 0.44) ||
                                 (result.moment_reflections > 0 && result.second_moment < 1.85));
    return result;
}

std::string TwinningAnalysisToText(const TwinningAnalysisResult& result) {
    std::ostringstream os;
    os << std::fixed << std::setprecision(3);
    os << "Twinning analysis\n";
    if (result.l_test_pairs > 0) {
        os << "  L-test (Padilla-Yeates): <|L|> = " << result.mean_abs_l
           << ", <L^2> = " << result.mean_l_squared
           << "  [untwinned 0.500 / 0.333, perfect twin 0.375 / 0.200; " << result.l_test_pairs
           << " pairs]\n";
        if (result.l_test_contaminated_by_tncs)
            os << "     NOT READ: this crystal has a translational pseudo-symmetry, and no partner\n"
               << "     reflection available to this test shares the class that pseudo-translation\n"
               << "     defines, so <|L|> is biased upwards by it. The number above is not a statement\n"
               << "     about twinning in either direction; the verdict below rests on the second\n"
               << "     moment alone.\n";
        else if (result.l_test_tncs_step_restricted)
            os << "     Measured against partner reflections chosen to share the class of this\n"
               << "     crystal's pseudo-translation, so the pseudo-symmetry does not bias it.\n";
    }
    if (result.moment_reflections > 0)
        os << "  Second moment <I^2>/<I>^2 = " << result.second_moment
           << "  [untwinned 2.00, perfect twin 1.50]\n";
    if (result.twinning_suspected && result.estimated_twin_fraction > 0.01)
        os << "  => Twinning suspected (estimated twin fraction ~"
           << result.estimated_twin_fraction << "). Statistics flag the presence of twinning, not\n"
           << "     the twin law; the law itself is a question for a dedicated twin-law analysis.\n";
    else if (result.twinning_suspected)
        // The fraction is derived from the second moment alone, so a twin called by the L-test on a
        // crystal whose second moment sits at or above its untwinned value has no fraction to quote -
        // printing the 0.000 the formula returns would contradict the sentence it sits in. The
        // disagreement is worth stating, and where a pseudo-symmetry is present it has a named cause.
        os << "  => Twinning suspected on the L-test. The second moment does not agree ("
           << result.second_moment << ", at or\n"
           << "     above its untwinned 2.00), so no twin fraction is quoted"
           << (result.l_test_tncs_step_restricted
                       ? " - and this\n"
                         "     crystal's translational pseudo-symmetry is a known reason for the second\n"
                         "     moment to run high, so the L-test carries the verdict here.\n"
                       : ".\n")
           << "     Statistics flag the presence of twinning, not the twin law; the law itself is\n"
           << "     a question for a dedicated twin-law analysis.\n";
    else if (!result.merohedral_twinning_possible && result.laue_class_was_chosen_by_promotion)
        os << "  => Cannot rule out twinning from these numbers: the Laue class is holohedral, so no\n"
           << "     merohedral twin law exists WITHIN it - but this Laue class was chosen by the\n"
           << "     space-group search itself, and promoting into a twin's holohedry is precisely what\n"
           << "     a merohedral twin looks like. The subgroup statistics reported by the search,\n"
           << "     not these numbers, are where the twinning is decided.\n";
    else if (!result.merohedral_twinning_possible)
        os << "  => No twinning: the Laue class is holohedral, so no merohedral twin law exists\n"
           << "     (any <|L|> below 0.5 here is a statistical artefact, not twinning).\n";
    else if (result.l_test_contaminated_by_tncs)
        os << "  => No twinning indicated by the second moment. The L-test, which is normally the\n"
           << "     stronger of the two, could not be read on this crystal (see above), so this is a\n"
           << "     weaker statement than usual - where the space-group search reported subgroup\n"
           << "     statistics, those are the stronger evidence.\n";
    else
        os << "  => No twinning indicated.\n";
    return os.str();
}