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

#include "ResolutionCutoff.h"

#include <algorithm>
#include <cmath>
#include <limits>

#include "../../common/CorrelationCoefficient.h"

namespace {
    // Fine CC1/2 bins for the fit (finer than the 10 reported shells, per the design). A bin needs
    // this many merged reflections for its CC1/2 to be trusted.
    constexpr int   N_FIT_BINS = 25;
    constexpr int   MIN_BIN_COUNT = 10;
    constexpr int   MIN_FIT_BINS = 5;          // need at least this many usable bins for a fit
    constexpr int   EXTEND_BINS_PAST_FALLOFF = 2;   // bins kept beyond the first sub-target bin
    constexpr double SHELLS_FOR_EXTENSION = 10.0;   // "+1 shell" = one 10-shell width in s (report-independent)

    double Logistic(double s, double k, double s0) {
        return 1.0 / (1.0 + std::exp(k * (s - s0)));
    }

    // Weighted (equal-weight) sum of squared residuals of the logistic against the binned CC1/2.
    double FitSSE(const std::vector<double> &s, const std::vector<double> &cc, double k, double s0) {
        double sse = 0.0;
        for (size_t i = 0; i < s.size(); ++i) {
            const double r = cc[i] - Logistic(s[i], k, s0);
            sse += r * r;
        }
        return sse;
    }
}

ResolutionCutoffResult ComputeCCHalfLogisticCutoff(const std::vector<MergedReflection> &merged,
                                                   double cc_target, Logger &logger) {
    ResolutionCutoffResult result;
    if (!(cc_target > 0.0 && cc_target < 1.0)) {
        result.note = "invalid CC target";
        return result;
    }

    // s = 1/d^2 range over the merged reflections that carry a half-set pair.
    double s_lo = std::numeric_limits<double>::max(), s_hi = 0.0;
    double d_data_min = std::numeric_limits<double>::max();
    for (const auto &m : merged) {
        if (!(m.d > 0.0f) || !std::isfinite(m.I_half[0]) || !std::isfinite(m.I_half[1]))
            continue;
        const double s = 1.0 / (static_cast<double>(m.d) * m.d);
        s_lo = std::min(s_lo, s);
        s_hi = std::max(s_hi, s);
        d_data_min = std::min(d_data_min, static_cast<double>(m.d));
    }
    if (!(s_lo < s_hi)) {
        result.note = "no half-set data for CC1/2 fit";
        return result;
    }

    // Bin CC1/2 against s (equal width in s, matching the reporting shells which are equal in 1/d^2).
    const double bin_w = (s_hi - s_lo) / N_FIT_BINS;
    std::vector<CorrelationCoefficient> bin_cc(N_FIT_BINS);
    std::vector<int> bin_n(N_FIT_BINS, 0);
    for (const auto &m : merged) {
        if (!(m.d > 0.0f) || !std::isfinite(m.I_half[0]) || !std::isfinite(m.I_half[1]))
            continue;
        const double s = 1.0 / (static_cast<double>(m.d) * m.d);
        int b = static_cast<int>((s - s_lo) / bin_w);
        b = std::clamp(b, 0, N_FIT_BINS - 1);
        bin_cc[b].Add(m.I_half[0], m.I_half[1]);
        ++bin_n[b];
    }

    // Usable bins (enough counts), in ascending-s order, with their bin-centre s.
    std::vector<double> s_bin, cc_bin;
    for (int b = 0; b < N_FIT_BINS; ++b) {
        if (bin_n[b] < MIN_BIN_COUNT) continue;
        const double cc = bin_cc[b].GetCC();
        if (!std::isfinite(cc)) continue;
        s_bin.push_back(s_lo + (b + 0.5) * bin_w);
        cc_bin.push_back(cc);
    }
    if (static_cast<int>(s_bin.size()) < MIN_FIT_BINS) {
        result.note = "too few usable CC1/2 bins";
        return result;
    }

    // Restrict to the contiguous fall-off from low res: keep bins up to a couple past the first one
    // that drops below cc_target, so a high-res noise blip cannot pull the fit back up. If the lowest
    // bin is already below cc_target there is no low-res plateau to anchor on - bail out.
    if (cc_bin.front() < cc_target) {
        result.note = "no low-resolution CC1/2 plateau";
        return result;
    }
    size_t keep = cc_bin.size();
    for (size_t i = 0; i < cc_bin.size(); ++i) {
        if (cc_bin[i] < cc_target) {
            keep = std::min(cc_bin.size(), i + 1 + EXTEND_BINS_PAST_FALLOFF);
            break;
        }
    }
    s_bin.resize(keep);
    cc_bin.resize(keep);
    if (static_cast<int>(s_bin.size()) < MIN_FIT_BINS) {
        result.note = "too few CC1/2 bins in the fall-off region";
        return result;
    }

    // Fit the logistic by a grid search over (k>0, s0) then a local coordinate-descent refine
    // (dependency-free; the fall-off is smooth and the grid lands close). s0 spans the s range; k
    // spans transitions from very gradual to very sharp relative to that range.
    const double s_range = s_hi - s_lo;
    double best_k = 0.0, best_s0 = 0.0, best_sse = std::numeric_limits<double>::max();
    constexpr int N_S0 = 60, N_K = 40;
    const double k_min = 2.0 / s_range, k_max = 200.0 / s_range;
    for (int ik = 0; ik < N_K; ++ik) {
        const double k = k_min * std::pow(k_max / k_min, static_cast<double>(ik) / (N_K - 1));
        for (int is = 0; is < N_S0; ++is) {
            const double s0 = s_lo + s_range * static_cast<double>(is) / (N_S0 - 1);
            const double sse = FitSSE(s_bin, cc_bin, k, s0);
            if (sse < best_sse) { best_sse = sse; best_k = k; best_s0 = s0; }
        }
    }

    double k = best_k, s0 = best_s0;
    double step_s0 = s_range / N_S0, step_k = best_k * 0.5;
    for (int iter = 0; iter < 200; ++iter) {
        bool improved = false;
        for (const double ds : {step_s0, -step_s0}) {
            const double sse = FitSSE(s_bin, cc_bin, k, s0 + ds);
            if (sse < best_sse) { best_sse = sse; s0 += ds; improved = true; }
        }
        for (const double dk : {step_k, -step_k}) {
            const double kt = k + dk;
            if (kt <= 0.0) continue;
            const double sse = FitSSE(s_bin, cc_bin, kt, s0);
            if (sse < best_sse) { best_sse = sse; k = kt; improved = true; }
        }
        if (!improved) { step_s0 *= 0.5; step_k *= 0.5; }
        if (step_s0 < 1e-6 * s_range && step_k < 1e-6 * best_k) break;
    }

    // s where the fitted CC1/2 crosses cc_target, then "one shell too far".
    const double s_cross = s0 + std::log(1.0 / cc_target - 1.0) / k;
    const double delta_s = s_range / SHELLS_FOR_EXTENSION;
    const double s_final = s_cross + delta_s;

    // No cut if the fall-off is beyond the measured edge (CC1/2 still healthy at the highest s).
    if (s_final >= s_hi) {
        result.note = "CC1/2 does not fall off within the measured range";
        return result;
    }
    // Low-resolution floor: never cut into good low-res data. A fit that puts the cutoff within two
    // shells of the lowest-res data is not a real fall-off - keep the full range and warn.
    if (s_final <= s_lo + 2.0 * delta_s) {
        logger.Warning("Resolution cutoff fit landed at low resolution (degenerate CC1/2 fall-off); "
                       "keeping the full resolution range");
        result.note = "degenerate low-resolution fit";
        return result;
    }

    double d_cut = 1.0 / std::sqrt(s_final);
    d_cut = std::max(d_cut, d_data_min);   // cannot cut beyond the highest-resolution reflection

    // A cut that is not meaningfully coarser than the data edge is a no-op.
    if (d_cut <= d_data_min * 1.001) {
        result.note = "CC1/2 healthy to the detector edge";
        return result;
    }

    result.d_cut = d_cut;
    result.note = "CC1/2 logistic fall-off, +1 shell";
    return result;
}