// 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) {
    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 ---
    // 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 preserves the reflection
    // condition of P/I/C/F/A/B lattices (parity-based), but 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::array<std::array<int, 3>, 5> steps{
        {{2, 0, 0}, {0, 2, 0}, {0, 0, 2}, {1, 1, 0}, {1, 1, 3}}};

    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 : steps) {
            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;
    if (n_pairs > 0) {
        result.mean_abs_l = sum_abs_l / n_pairs;
        result.mean_l_squared = sum_l2 / n_pairs;
    }

    // --- 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. lysozyme's 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);
    result.twinning_suspected = result.merohedral_twinning_possible &&
                                ((result.l_test_pairs > 0 && 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.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)
        os << "  => Twinning suspected (estimated twin fraction ~"
           << result.estimated_twin_fraction << "). Statistics flag the presence of twinning, not\n"
           << "     the twin law; confirm with a dedicated twin-law analysis.\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
        os << "  => No twinning indicated.\n";
    return os.str();
}