// 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;
    }
}

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);

    const std::array<std::array<int, 3>, 3> steps{{{2, 0, 0}, {0, 2, 0}, {0, 0, 2}}};

    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.
    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)));
        };
        std::vector<double> shell_sum(n_shells, 0.0);
        std::vector<int> shell_n(n_shells, 0);
        for (const auto& r : merged) {
            if (!UsableIntensity(r) || !acentric(r))
                continue;
            const int b = shell_of(r.d);
            shell_sum[b] += r.I;
            shell_n[b] += 1;
        }
        std::vector<double> shell_mean(n_shells, 0.0);
        for (int b = 0; b < n_shells; ++b)
            if (shell_n[b] > 0)
                shell_mean[b] = shell_sum[b] / shell_n[b];

        double sum_e4 = 0.0;
        int n_moment = 0;
        for (const auto& r : merged) {
            if (!UsableIntensity(r) || !acentric(r))
                continue;
            const double mean = shell_mean[shell_of(r.d)];
            if (mean <= 0.0)
                continue;
            const double e2 = r.I / mean;
            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.
    result.twinning_suspected = (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
        os << "  => No twinning indicated.\n";
    return os.str();
}