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

#include "SigmaA.h"

#include <algorithm>
#include <cmath>
#include <vector>

#include "gemmi/bessel.hpp"   // bessel_i1_over_i0, log_bessel_i0
#include "gemmi/binner.hpp"   // Binner
#include "gemmi/math.hpp"     // log_cosh

namespace {

// How many free reflections a shell is given. sigma_A is one number per shell, so a few tens of
// reflections already pin it; sizing the SHELLS from the free count (rather than fixing the number of
// shells and then patching the thin ones) is what keeps every shell usable on a small dataset, where
// the free set is 5 % of a few thousand reflections.
constexpr int FREE_PER_SHELL = 50;
constexpr int MAX_SHELLS = 20;

// Below this a shell's own maximum likelihood is noise, and the estimate over the whole free set is
// used for it instead. Equal-count shells are cut on all the reflections, so a shell's free count
// still fluctuates about FREE_PER_SHELL and the odd one comes out short.
constexpr int MIN_FREE_FOR_SHELL = 10;

constexpr double SIGMA_A_MIN = 0.01;
constexpr double SIGMA_A_MAX = 0.99;

struct Normalized {
    double eo = 0, ec = 0;
    bool centric = false;
};

// Rice (acentric) and Woolfson (centric) log likelihoods of |Eo| given |Ec| and sigma_A, with the
// terms that do not depend on sigma_A dropped.
double LogLikelihood(const std::vector<Normalized> &e, const std::vector<int> &idx, double sigma_a) {
    const double u = 1 - sigma_a * sigma_a;
    double ll = 0;
    for (int i : idx) {
        const double eo = e[i].eo, ec = e[i].ec;
        const double quad = eo * eo + sigma_a * sigma_a * ec * ec;
        if (e[i].centric)
            ll += -0.5 * std::log(u) - quad / (2 * u) + gemmi::log_cosh(sigma_a * eo * ec / u);
        else
            ll += -std::log(u) - quad / u + gemmi::log_bessel_i0(2 * sigma_a * eo * ec / u);
    }
    return ll;
}

// The likelihood is unimodal in sigma_A, so a golden-section search finds its maximum without
// derivatives and without a starting guess.
double MaximizeSigmaA(const std::vector<Normalized> &e, const std::vector<int> &idx) {
    constexpr double GOLDEN = 0.6180339887498949;
    double lo = SIGMA_A_MIN, hi = SIGMA_A_MAX;
    double x1 = hi - GOLDEN * (hi - lo), x2 = lo + GOLDEN * (hi - lo);
    double f1 = LogLikelihood(e, idx, x1), f2 = LogLikelihood(e, idx, x2);
    for (int it = 0; it < 60; it++) {
        if (f1 > f2) {
            hi = x2; x2 = x1; f2 = f1;
            x1 = hi - GOLDEN * (hi - lo); f1 = LogLikelihood(e, idx, x1);
        } else {
            lo = x1; x1 = x2; f1 = f2;
            x2 = lo + GOLDEN * (hi - lo); f2 = LogLikelihood(e, idx, x2);
        }
    }
    return 0.5 * (lo + hi);
}

}  // namespace

SigmaAResult EstimateSigmaA(const std::vector<SigmaAReflection> &refl, const gemmi::UnitCell &cell) {
    SigmaAResult result;
    result.weight.assign(refl.size(), SigmaAWeight{});
    if (refl.empty())
        return result;

    for (const SigmaAReflection &r : refl)
        if (r.free)
            ++result.free_reflections;
    if (result.free_reflections == 0)
        return result;   // nothing to estimate on: leave the coefficients unweighted

    const int nbins = std::clamp(result.free_reflections / FREE_PER_SHELL, 1, MAX_SHELLS);
    std::vector<double> inv_d2;
    inv_d2.reserve(refl.size());
    for (const SigmaAReflection &r : refl)
        inv_d2.push_back(r.inv_d2);
    gemmi::Binner binner;
    binner.setup_from_1_d2(nbins, gemmi::Binner::Method::EqualCount, std::move(inv_d2), &cell);
    std::vector<int> bin(refl.size());
    for (size_t i = 0; i < refl.size(); i++)
        bin[i] = binner.get_bin_from_1_d2(refl[i].inv_d2);
    result.shells = nbins;

    // Normalize both amplitudes to <|E|^2> = 1 within the shell, which is the scale the likelihood
    // above is written in. The epsilon factor is the symmetry enhancement of a reflection's expected
    // intensity, and dividing it out is what makes the reflections of one shell comparable.
    std::vector<double> sum_o(nbins, 0), sum_c(nbins, 0);
    std::vector<int> count(nbins, 0);
    for (size_t i = 0; i < refl.size(); i++) {
        const double eps = std::max(1, refl[i].epsilon);
        sum_o[bin[i]] += refl[i].f_obs * refl[i].f_obs / eps;
        sum_c[bin[i]] += refl[i].f_calc * refl[i].f_calc / eps;
        ++count[bin[i]];
    }
    std::vector<double> sigma_o(nbins), sigma_c(nbins);
    for (int b = 0; b < nbins; b++) {
        sigma_o[b] = count[b] > 0 ? sum_o[b] / count[b] : 0;
        sigma_c[b] = count[b] > 0 ? sum_c[b] / count[b] : 0;
    }

    std::vector<Normalized> e(refl.size());
    for (size_t i = 0; i < refl.size(); i++) {
        const int b = bin[i];
        const double eps = std::max(1, refl[i].epsilon);
        e[i].eo = sigma_o[b] > 0 ? refl[i].f_obs / std::sqrt(eps * sigma_o[b]) : 0;
        e[i].ec = sigma_c[b] > 0 ? refl[i].f_calc / std::sqrt(eps * sigma_c[b]) : 0;
        e[i].centric = refl[i].centric;
    }

    std::vector<std::vector<int>> free_in_bin(nbins);
    std::vector<int> free_all;
    for (size_t i = 0; i < refl.size(); i++)
        if (refl[i].free) {
            free_in_bin[bin[i]].push_back(static_cast<int>(i));
            free_all.push_back(static_cast<int>(i));
        }
    const double sigma_a_overall = MaximizeSigmaA(e, free_all);
    std::vector<double> sigma_a(nbins, sigma_a_overall);
    for (int b = 0; b < nbins; b++)
        if (static_cast<int>(free_in_bin[b].size()) >= MIN_FREE_FOR_SHELL)
            sigma_a[b] = MaximizeSigmaA(e, free_in_bin[b]);
    result.sigma_a_lowest_shell = sigma_a.front();
    result.sigma_a_highest_shell = sigma_a.back();

    double fom_sum = 0;
    for (size_t i = 0; i < refl.size(); i++) {
        const int b = bin[i];
        const double s = sigma_a[b];
        const double x = s * e[i].eo * e[i].ec / (1 - s * s);
        result.weight[i].m = refl[i].centric ? std::tanh(x) : gemmi::bessel_i1_over_i0(2 * x);
        // D takes Fc from its own shell scale to the observed one; the two are already close, because
        // Fcalc reaches here scaled to Fobs overall and with a bulk-solvent and anisotropic-B model.
        result.weight[i].d = sigma_c[b] > 0 ? s * std::sqrt(sigma_o[b] / sigma_c[b]) : s;
        fom_sum += result.weight[i].m;
    }
    result.mean_fom = fom_sum / static_cast<double>(refl.size());
    return result;
}