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

#include <algorithm>
#include <cmath>

#include "IceScore.h"
#include "../common/Definitions.h"

namespace {
    constexpr float TWO_PI = 6.283185307f;

    // Below this q the profile is direct beam and beam-stop shadow, not diffraction. The lowest ice
    // line of either phase sits at 1.61 A^-1, so nothing is lost.
    constexpr float Q_MIN_EVAL = 1.20f;
    // Running-median half-window for the background under the bands. A median over 13 bins rejects a
    // 3-5 bin powder ring but follows the ~40-bin-wide vitreous-water halo, so the halo never reaches
    // the residual and cannot score.
    constexpr int BACKGROUND_HALF_WINDOW = 6;
    // Half-window over which the bin mean's own error is smoothed, so that a ring cannot inflate its
    // own denominator.
    constexpr int SIGMA_HALF_WINDOW = 20;
    // A band must be a real excess to contribute at all; a uniform low-level positive bias spread over
    // many bands then contributes nothing. Capped so that one saturated band cannot carry a hypothesis.
    constexpr float Z_FLOOR = 2.0f;
    constexpr float Z_CAP = 6.0f;
    // The conventional 3-sigma detection maps to a score of 0.5.
    constexpr float S_HALF = 3.0f;
    // How far either side of a predicted band the radial channel looks for the bin that carries it.
    // This is NOT a second band half-width competing with the spot channel's half_width_q: it is a
    // tolerance for a mis-set beam centre, which smears and splits a ring, and it is set a priori from
    // the measured centre error (+-0.6 % in d) rather than fitted. There is one band width in this
    // file, and it is the pipeline's ice_ring_width_Q_recipA, passed in.
    constexpr float CENTRE_SMEAR_Q = 0.012f;
    // Bins nearer than this to a band of EITHER phase are not part of the null.
    constexpr float NULL_EXCLUSION_Q = 0.025f;
    // Fewer null bins than this and the null is not measured.
    constexpr int NULL_MIN_BINS = 20;

    // How far a band may be slid for its control, and in what steps.
    constexpr float SPOT_OFFSET_MAX_Q = 0.45f;
    constexpr float SPOT_OFFSET_STEP_Q = 0.0025f;
    // A control radius the detector barely covers is not a control.
    constexpr float SPOT_MIN_AREA_FRACTION = 0.05f;
    constexpr int SPOT_MIN_OFFSETS = 5;
    // A control count below this is not measurable; do not divide by it.
    constexpr float SPOT_CONTROL_FLOOR = 0.5f;
    // -log10 p at which the significance term saturates, and the band/control ratio at which the size
    // term does.
    constexpr float SPOT_SIGNIFICANCE_SAT = 4.0f;
    constexpr float SPOT_RATIO_SAT = 3.0f;

    // Average a q_bins x azimuthal-bins array over azimuth. A profile that is already 1-D passes
    // through. Bins with nothing finite in them come out NaN.
    std::vector<float> FoldToQ(const std::vector<float> &in, int nq) {
        std::vector<float> out(nq, 0.0f);
        std::vector<int> n(nq, 0);
        for (size_t i = 0; i < in.size() && nq > 0; i++) {
            const int q = static_cast<int>(i % nq);
            if (std::isfinite(in[i])) {
                out[q] += in[i];
                n[q]++;
            }
        }
        for (int q = 0; q < nq; q++)
            out[q] = n[q] ? out[q] / static_cast<float>(n[q]) : NAN;
        return out;
    }

    // The same fold for the live pixel count, which adds rather than averages.
    std::vector<double> FoldCountToQ(const std::vector<uint64_t> &in, int nq) {
        std::vector<double> out(nq, 0.0);
        for (size_t i = 0; i < in.size() && nq > 0; i++)
            out[i % nq] += static_cast<double>(in[i]);
        return out;
    }

    // Running median, NaN where the window holds nothing finite.
    std::vector<float> RunningMedian(const std::vector<float> &v, int half) {
        std::vector<float> out(v.size(), NAN);
        std::vector<float> window;
        for (int i = 0; i < static_cast<int>(v.size()); i++) {
            window.clear();
            const int lo = std::max(0, i - half);
            const int hi = std::min(static_cast<int>(v.size()), i + half + 1);
            for (int j = lo; j < hi; j++)
                if (std::isfinite(v[j]))
                    window.push_back(v[j]);
            if (window.empty())
                continue;
            std::ranges::nth_element(window, window.begin() + window.size() / 2);
            out[i] = window[window.size() / 2];
        }
        return out;
    }

    // One ice phase read as a hypothesis over the per-bin evidence c: the standardised amplitude of its
    // bands and the standardised count of bands that show anything, against the frame's own null.
    float HypothesisScore(const std::vector<float> &c, const std::vector<float> &band_q,
                          const std::vector<float> &band_weight, int min_bands,
                          float low_q, float dq, int bh,
                          float mu0, float v0, float p0) {
        const int nq = static_cast<int>(c.size());
        float w1 = 0.0f, w2 = 0.0f, evidence = 0.0f;
        int nbands = 0, nhit = 0;
        for (size_t i = 0; i < band_q.size(); i++) {
            if (band_q[i] < Q_MIN_EVAL)
                continue;
            const int b = static_cast<int>(std::lround((band_q[i] - low_q) / dq - 0.5f));
            if (b < bh || b >= nq - bh || !std::isfinite(c[b]))
                continue;
            nbands++;
            w1 += band_weight[i];
            w2 += band_weight[i] * band_weight[i];
            evidence += c[b] * band_weight[i];
            if (c[b] > 0.0f)
                nhit++;
        }
        if (nbands < min_bands)
            return 0.0f;
        const float s_amp = (evidence - mu0 * w1) / std::sqrt(std::max(v0 * w2, 1e-6f));
        const float n = static_cast<float>(nbands);
        const float s_cnt = (static_cast<float>(nhit) - n * p0)
                            / std::sqrt(std::max(n * p0 * (1.0f - p0), 0.0f) + 0.25f);
        return std::max(std::min(s_amp, s_cnt), 0.0f);
    }

    // P(X >= k) for X ~ Poisson(mu), summed upward from k. The lower tail would lose a small
    // probability to cancellation, and it is the small probabilities this is wanted for.
    double PoissonUpperTail(int64_t k, double mu) {
        if (k <= 0)
            return 1.0;
        if (!(mu > 0.0))
            return 0.0;
        const double log_mu = std::log(mu);
        double sum = 0.0;
        for (int64_t j = k; j < k + 100000; j++) {
            const double term = std::exp(-mu + static_cast<double>(j) * log_mu - std::lgamma(static_cast<double>(j) + 1.0));
            sum += term;
            if (static_cast<double>(j) > mu && term < 1e-18 * sum)
                break;
        }
        return std::min(sum, 1.0);
    }

    // Live pixels within +-w of q.
    double AreaAt(const std::vector<double> &count_q, float low_q, float dq, float q, float w) {
        double area = 0.0;
        const int lo = static_cast<int>(std::ceil((q - w - low_q) / dq - 0.5f));
        const int hi = static_cast<int>(std::floor((q + w - low_q) / dq - 0.5f));
        for (int i = std::max(lo, 0); i <= std::min(hi, static_cast<int>(count_q.size()) - 1); i++)
            area += count_q[i];
        return area;
    }

    // Spots within +-w of q, from the sorted spot q list.
    double CountAt(const std::vector<float> &sorted_q, float q, float w) {
        const auto lo = std::lower_bound(sorted_q.begin(), sorted_q.end(), q - w);
        const auto hi = std::upper_bound(sorted_q.begin(), sorted_q.end(), q + w);
        return static_cast<double>(std::distance(lo, hi));
    }
}

float IceScoreRadial(const std::vector<float> &profile, const std::vector<float> &profile_std,
                     const std::vector<uint64_t> &profile_count, int32_t q_bins,
                     const AzimuthalIntegrationSettings &settings) {
    const int nq = std::max<int>(q_bins, 0);
    const float low_q = settings.GetLowQ_recipA();
    const float dq = settings.GetQSpacing_recipA();
    if (nq < 2 * SIGMA_HALF_WINDOW || !(dq > 0.0f) || profile_std.empty())
        return 0.0f;

    std::vector<float> prof = FoldToQ(profile, nq);
    std::vector<float> sigma = FoldToQ(profile_std, nq);
    const std::vector<double> count = FoldCountToQ(profile_count, nq);
    for (int i = 0; i < nq; i++) {
        if (!(prof[i] > 0.0f))
            prof[i] = NAN;
        // The error of the bin mean, not of one pixel.
        if (sigma[i] > 0.0f && count[i] > 0.0)
            sigma[i] /= static_cast<float>(std::sqrt(count[i]));
        else
            sigma[i] = NAN;
    }

    const std::vector<float> background = RunningMedian(prof, BACKGROUND_HALF_WINDOW);
    const std::vector<float> sigma_smooth = RunningMedian(sigma, SIGMA_HALF_WINDOW);
    const int bh = std::max(1, static_cast<int>(std::lround(CENTRE_SMEAR_Q / dq)));
    auto q_of = [&](int i) { return low_q + (static_cast<float>(i) + 0.5f) * dq; };

    // Per-bin excess in units of the bin mean's own error, floored so that only a real excess counts.
    // The sigma floor at 1 % of the background keeps a bin whose error is reported as tiny - or not at
    // all - from turning noise into a detection.
    std::vector<float> z(nq, NAN);
    for (int i = 0; i < nq; i++) {
        if (q_of(i) < Q_MIN_EVAL || !std::isfinite(prof[i]) || !std::isfinite(background[i]))
            continue;
        const float floor_sigma = 0.01f * background[i];
        const float sg = std::max(std::isfinite(sigma_smooth[i]) ? sigma_smooth[i] : 0.0f, floor_sigma);
        if (sg > 0.0f)
            z[i] = (prof[i] - background[i]) / sg;
    }
    std::vector<float> c(nq, NAN);
    for (int i = 0; i < nq; i++) {
        float best = NAN;
        for (int o = -bh; o <= bh; o++) {
            const float v = z[std::clamp(i + o, 0, nq - 1)];
            if (std::isfinite(v) && (!std::isfinite(best) || v > best))
                best = v;
        }
        if (std::isfinite(best))
            c[i] = std::clamp(best - Z_FLOOR, 0.0f, Z_CAP);
    }

    // The null: the same statistic at every bin belonging to no band of either phase.
    std::vector<float> null_c;
    for (int i = 0; i < nq; i++) {
        if (!std::isfinite(c[i]) || q_of(i) < Q_MIN_EVAL)
            continue;
        float nearest = 1e9f;
        for (const float d: ICE_RING_RES_A)
            nearest = std::min(nearest, std::fabs(q_of(i) - TWO_PI / d));
        for (const float d: ICE_RING_CUBIC_RES_A)
            nearest = std::min(nearest, std::fabs(q_of(i) - TWO_PI / d));
        if (nearest > static_cast<float>(bh) * dq + NULL_EXCLUSION_Q)
            null_c.push_back(c[i]);
    }
    if (static_cast<int>(null_c.size()) < NULL_MIN_BINS)
        return 0.0f;

    float mu0 = 0.0f, p0 = 0.0f;
    for (const float v: null_c) {
        mu0 += v;
        p0 += (v > 0.0f) ? 1.0f : 0.0f;
    }
    const float kn = static_cast<float>(null_c.size());
    mu0 /= kn;
    p0 /= kn;
    float v0 = 0.0f;
    for (const float v: null_c)
        v0 += (v - mu0) * (v - mu0);
    v0 /= (kn - 1.0f);

    // Hexagonal ice: the primary triplet above 3 A carries twice the weight of the rest.
    std::vector<float> hex_q, hex_w;
    for (const float d: ICE_RING_RES_A) {
        hex_q.push_back(TWO_PI / d);
        hex_w.push_back(d > 3.0f ? 2.0f : 1.0f);
    }
    std::vector<float> cubic_q, cubic_w;
    for (const float d: ICE_RING_CUBIC_RES_A) {
        cubic_q.push_back(TWO_PI / d);
        cubic_w.push_back(d > 3.0f ? 2.0f : 1.0f);
    }

    const float s = std::max(HypothesisScore(c, hex_q, hex_w, 4, low_q, dq, bh, mu0, v0, p0),
                             HypothesisScore(c, cubic_q, cubic_w, 3, low_q, dq, bh, mu0, v0, p0));
    return s * s / (s * s + S_HALF * S_HALF);
}

float IceScoreSpots(const std::vector<float> &spot_d_A, const std::vector<uint64_t> &profile_count,
                    int32_t q_bins, const AzimuthalIntegrationSettings &settings, float half_width_q) {
    const int nq = std::max<int>(q_bins, 0);
    const float low_q = settings.GetLowQ_recipA();
    const float dq = settings.GetQSpacing_recipA();
    if (nq < 2 || !(dq > 0.0f) || !(half_width_q > 0.0f) || spot_d_A.empty() || profile_count.empty())
        return 0.0f;

    const std::vector<double> count_q = FoldCountToQ(profile_count, nq);

    std::vector<float> spot_q;
    spot_q.reserve(spot_d_A.size());
    for (const float d: spot_d_A)
        if (d > 0.0f)
            spot_q.push_back(TWO_PI / d);
    std::ranges::sort(spot_q);

    std::vector<float> band_q;
    for (const float d: ICE_RING_RES_A)
        band_q.push_back(TWO_PI / d);

    double k0 = 0.0, mu = 0.0, var = 0.0;
    for (const float qb: band_q) {
        const double area_band = AreaAt(count_q, low_q, dq, qb, half_width_q);
        if (!(area_band > 0.0))
            continue;
        std::vector<double> control;
        for (float delta = 2.0f * half_width_q; delta <= SPOT_OFFSET_MAX_Q; delta += SPOT_OFFSET_STEP_Q) {
            // An offset that lands on another ice band is not a control.
            float nearest = 1e9f;
            for (const float qo: band_q)
                nearest = std::min(nearest, std::min(std::fabs(qb + delta - qo), std::fabs(qb - delta - qo)));
            if (nearest < 2.0f * half_width_q)
                continue;
            const double area_hi = AreaAt(count_q, low_q, dq, qb + delta, half_width_q);
            const double area_lo = AreaAt(count_q, low_q, dq, qb - delta, half_width_q);
            if (area_hi < SPOT_MIN_AREA_FRACTION * area_band || area_lo < SPOT_MIN_AREA_FRACTION * area_band)
                continue;
            control.push_back(0.5 * area_band * (CountAt(spot_q, qb + delta, half_width_q) / area_hi
                                                 + CountAt(spot_q, qb - delta, half_width_q) / area_lo));
        }
        if (static_cast<int>(control.size()) < SPOT_MIN_OFFSETS)
            continue;
        double mean = 0.0;
        for (const double v: control)
            mean += v;
        mean /= static_cast<double>(control.size());
        double m2 = 0.0;
        for (const double v: control)
            m2 += (v - mean) * (v - mean);
        k0 += CountAt(spot_q, qb, half_width_q);
        mu += mean;
        var += m2 / static_cast<double>(control.size());
    }

    mu = std::max(mu, static_cast<double>(SPOT_CONTROL_FLOOR));
    if (k0 <= mu)
        return 0.0f;
    // Quasi-Poisson: the controls' own scatter says how much wider than Poisson the count really is.
    const double phi = std::max(1.0, var / mu);
    const double p = PoissonUpperTail(static_cast<int64_t>(std::floor((k0 - 1.0) / phi)) + 1, mu / phi);
    const double s_sig = std::min(1.0, -std::log10(std::max(p, 1e-300)) / SPOT_SIGNIFICANCE_SAT);
    const double s_size = std::min(1.0, std::log(k0 / mu) / std::log(SPOT_RATIO_SAT));
    return static_cast<float>(std::max(0.0, std::min(s_sig, s_size)));
}

float IceScore(const std::vector<float> &profile, const std::vector<float> &profile_std,
               const std::vector<uint64_t> &profile_count, int32_t q_bins,
               const AzimuthalIntegrationSettings &settings, const std::vector<float> &spot_d_A,
               float half_width_q) {
    // The larger of the two, deliberately, and this is a decided question rather than an open one.
    //
    // The radial channel cannot see ice that arrives as large crystallites: that ice diffracts as
    // discrete spots and leaves the profile flat, so there is no radial evidence to read. Removing the
    // spot channel would make the score blind to the regime the second channel exists for, by
    // construction and not merely on this corpus.
    //
    // Measured over the corpus, the maximum takes the fraction of frames called icy in the protein
    // classes labelled clean from 7% (radial alone) to 16%. That looks like a specificity cost and is
    // not yet known to be one: the loops where the spot channel is the ONLY detector are precisely the
    // ones whose ice status the labels do not settle, and several loops labelled clean - and three of
    // four water loops - were independently found to carry a full hexagonal pattern. Until the corpus
    // is relabelled from measurement rather than from directory names, the corpus cannot say whether
    // those frames are false positives or correct detections nobody had recorded.
    //
    // If the rate does have to come down afterwards, the lever is to drop each hypothesis's worst band
    // before scoring it (measured on the prototype: clean protein 20.5 -> 12.2%, ice 88.8 -> ~83%), not
    // to remove a channel.
    return std::max(IceScoreRadial(profile, profile_std, profile_count, q_bins, settings),
                    IceScoreSpots(spot_d_A, profile_count, q_bins, settings, half_width_q));
}