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

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

#include "../../common/JFJochMath.h"
#include "SpotUtils.h"
#include "../../common/ResolutionShells.h"

void CountSpots(DataMessage &msg,
                const std::vector<SpotToSave> &spots,
                float d_min_A) {
    int64_t low_res = 0;
    int64_t ice_ring = 0;
    for (auto &s: spots) {
        if (s.ice_ring)
            ice_ring++;

        if (s.d_A > d_min_A)
            low_res++;
    }
    msg.spot_count = spots.size();
    msg.spot_count_low_res = low_res;
    msg.spot_count_ice_rings = ice_ring;
}

// Spots in the ice-free control flanks either side of the hexagonal rings, rescaled to the ring bands'
// own q width. The control for one ring is the two intervals [w, 2w) beside it - same total width as
// the ring band, and symmetric, so the fall-off of spot density with resolution cancels to first
// order. A flank that lands on another ring is not a control and is dropped, its width with it; the
// three rings at 1.947/1.916/1.882 A are 0.05-0.06 apart in q and usually lose both.
float CountIceRingControlSpots(const std::vector<SpotToSave> &spots, float w) {
    if (!(w > 0.0f))
        return 0.0f;
    float control = 0.0f;
    for (const float d : ICE_RING_RES_A) {
        const float q_ring = 2 * PI / d;
        bool lo_free = true, hi_free = true;
        for (const float other : ICE_RING_RES_A) {
            const float q_other = 2 * PI / other;
            if (q_other > q_ring && q_other < q_ring + 3 * w) hi_free = false;
            if (q_other < q_ring && q_other > q_ring - 3 * w) lo_free = false;
        }
        const int free_flanks = (lo_free ? 1 : 0) + (hi_free ? 1 : 0);
        if (free_flanks == 0)
            continue;
        int64_t n = 0;
        for (const auto &s: spots) {
            if (!(s.d_A > 0.0f)) continue;
            const float dq = 2 * PI / s.d_A - q_ring;
            if (hi_free && dq >= w && dq < 2 * w) n++;
            if (lo_free && dq <= -w && dq > -2 * w) n++;
        }
        // One free flank covers half the ring band's width, so it counts double.
        control += static_cast<float>(n) * 2.0f / static_cast<float>(free_flanks);
    }
    return control;
}

void MarkIceRings(std::vector<SpotToSave> &spots, float tolerance_q_recipA) {
    std::vector<float> ice_rings_q;

    for (const auto &i: ICE_RING_RES_A)
        ice_rings_q.push_back(2 * PI / i);

    for (auto &s: spots)
        s.ice_ring = false;

    MarkRings(spots, ice_rings_q, tolerance_q_recipA);
}

void MarkRings(std::vector<SpotToSave> &spots, const std::vector<float> &rings_q_recipA,
               float tolerance_q_recipA) {
    if (rings_q_recipA.empty())
        return;

    for (auto &s: spots) {
        if (!(s.d_A > 0.0f))
            continue;
        const float spot_q = 2 * PI / s.d_A;
        for (const float q: rings_q_recipA)
            if (fabs(spot_q - q) < tolerance_q_recipA) {
                s.ice_ring = true;
                break;
            }
    }
}

namespace {
    // A bin has to hold this many spots before its excess is looked at, so a pair of noise detections
    // in an otherwise empty bin cannot pose as a ring.
    constexpr int64_t RING_MIN_SPOTS_PER_BIN = 8;
    // ... and the excess over the local baseline has to be this many times its Poisson noise, AND
    // this large a fraction of the baseline. Both, because either alone fails at one end of the
    // range: the Poisson test alone calls a 5% rise a ring where the pool is large, and the fractional
    // test alone calls a two-spot bin a ring where it is small.
    constexpr float RING_MIN_SIGMA = 4.0f;
    constexpr float RING_MIN_EXCESS_FRACTION = 1.0f;
    // Bins of the running median that sets the baseline. Wide compared with the 2-3 bins a ring covers,
    // so the rings themselves do not pull the median up, and narrow compared with the whole q range,
    // so it still follows the fall-off of spot density with resolution.
    constexpr size_t RING_BASELINE_BINS = 41;
    // Past the q where this share of a baseline window's bins are ring bins, the rings have merged:
    // the median that the excess is measured against is itself a ring value, and nothing further out
    // can be told from the crystal. Half, because that is what "the median is a ring" means - it is
    // the definition of the running median failing, not a tuned number.
    constexpr float RING_MERGED_BIN_FRACTION = 0.5f;
    // The pool a measurement needs at all. Below this the histogram is counting statistics.
    constexpr size_t RING_MIN_POOLED_SPOTS = 500;
    // The share of the spots the rings have to hold before they are a PHASE rather than the crystal's
    // own rows. Every pattern has some q bins fuller than their neighbours - a crystal with a short
    // axis puts its reflections in sheets - and calling those a contaminant would be wrong. Calibrated
    // on the 100-dataset open battery, where the measured fractions run as a continuum from 83% down:
    // the sets with a visible powder sit at 8-83%, and below about a twentieth the "rings" are two or
    // three bins holding a percent of the spots, which every clean crystal in that battery also shows.
    constexpr float RING_MIN_SPOT_FRACTION = 0.05f;
}

PowderRings MeasurePowderRings(const std::vector<float> &spot_q_recipA, float half_width_q_recipA) {
    PowderRings out;
    if (!(half_width_q_recipA > 0.0f) || spot_q_recipA.size() < RING_MIN_POOLED_SPOTS)
        return out;

    float q_min = std::numeric_limits<float>::max(), q_max = 0.0f;
    for (const float q : spot_q_recipA) {
        if (!(q > 0.0f)) continue;
        q_min = std::min(q_min, q);
        q_max = std::max(q_max, q);
    }
    if (!(q_max > q_min))
        return out;

    // One bin per ring half-width, so a ring covers two or three of them.
    const size_t nbins = static_cast<size_t>((q_max - q_min) / half_width_q_recipA) + 1;
    if (nbins < RING_BASELINE_BINS)
        return out;

    std::vector<int64_t> count(nbins, 0);
    int64_t total = 0;
    for (const float q : spot_q_recipA) {
        if (!(q > 0.0f)) continue;
        const auto bin = static_cast<size_t>((q - q_min) / half_width_q_recipA);
        if (bin < nbins) { count[bin]++; total++; }
    }
    if (total == 0)
        return out;

    // Running median of the counts. The window is clipped at the ends rather than padded, so the first
    // and last few bins are judged against the baseline of the range they have.
    std::vector<float> baseline(nbins, 0.0f);
    std::vector<int64_t> window;
    window.reserve(RING_BASELINE_BINS);
    for (size_t i = 0; i < nbins; i++) {
        const size_t lo = i > RING_BASELINE_BINS / 2 ? i - RING_BASELINE_BINS / 2 : 0;
        const size_t hi = std::min(nbins, i + RING_BASELINE_BINS / 2 + 1);
        window.assign(count.begin() + static_cast<long>(lo), count.begin() + static_cast<long>(hi));
        std::ranges::nth_element(window, window.begin() + static_cast<long>(window.size() / 2));
        baseline[i] = static_cast<float>(window[window.size() / 2]);
    }

    std::vector<char> is_ring(nbins, 0);
    for (size_t i = 0; i < nbins; i++) {
        const float excess = static_cast<float>(count[i]) - baseline[i];
        is_ring[i] = count[i] >= RING_MIN_SPOTS_PER_BIN
                  && excess > RING_MIN_SIGMA * std::sqrt(std::max(baseline[i], 1.0f))
                  && excess > RING_MIN_EXCESS_FRACTION * baseline[i];
    }

    // Where the rings have merged. Read outwards over the same window the baseline uses: the first q
    // at which ring bins are the majority of that window is where the measurement stops meaning
    // anything, and everything past it is left alone.
    for (size_t i = RING_BASELINE_BINS / 2; i + RING_BASELINE_BINS / 2 < nbins; i++) {
        size_t n_ring = 0;
        for (size_t j = i - RING_BASELINE_BINS / 2; j <= i + RING_BASELINE_BINS / 2; j++)
            n_ring += is_ring[j] ? 1 : 0;
        if (static_cast<float>(n_ring) > RING_MERGED_BIN_FRACTION * static_cast<float>(RING_BASELINE_BINS)) {
            out.resolved_to_d_A = 2 * PI / (q_min + static_cast<float>(i) * half_width_q_recipA);
            break;
        }
    }

    // Contiguous runs of ring bins are one ring, placed at their count-weighted centre.
    double excess_total = 0.0;
    size_t run_start = nbins;
    const auto close_run = [&](size_t run_end) {
        double num = 0.0, den = 0.0;
        for (size_t j = run_start; j < run_end; j++) {
            const auto w = static_cast<double>(count[j]) - baseline[j];
            num += w * (q_min + (static_cast<double>(j) + 0.5) * half_width_q_recipA);
            den += w;
        }
        if (den > 0.0) {
            out.rings_q_recipA.push_back(static_cast<float>(num / den));
            excess_total += den;
        }
        run_start = nbins;
    };
    for (size_t i = 0; i < nbins; i++) {
        if (is_ring[i] && run_start == nbins)
            run_start = i;
        else if (!is_ring[i] && run_start != nbins)
            close_run(i);
    }
    if (run_start != nbins)
        close_run(nbins);

    out.spot_fraction = static_cast<float>(excess_total / static_cast<double>(total));
    if (out.spot_fraction < RING_MIN_SPOT_FRACTION)
        return {};   // measured, and it is not a phase
    return out;
}

std::optional<float> SpotResolutionQuantile(const std::vector<float> &spot_q_recipA, float fraction) {
    if (spot_q_recipA.size() < RING_MIN_POOLED_SPOTS || !(fraction > 0.0f) || !(fraction < 1.0f))
        return std::nullopt;
    std::vector<float> q;
    q.reserve(spot_q_recipA.size());
    for (const float v : spot_q_recipA)
        if (v > 0.0f) q.push_back(v);
    if (q.size() < RING_MIN_POOLED_SPOTS)
        return std::nullopt;
    const auto rank = static_cast<long>(static_cast<float>(q.size()) * fraction);
    std::ranges::nth_element(q, q.begin() + rank);
    return q[rank] > 0.0f ? std::optional<float>(2 * PI / q[rank]) : std::nullopt;
}

void FilterSpotsByCount(std::vector<SpotToSave> &input, int64_t count, bool deprioritise_ice) {
    size_t output_size = std::min<size_t>(input.size(), count);

    std::ranges::partial_sort(input, input.begin() + output_size,
                              std::ranges::less{}, // comparator on the projected key
                              [deprioritise_ice](const SpotToSave &s) {
                                  // projection: non-ice first (false < true), then strongest intensity
                                  // first. Where the run has no measurable ice the flag marks ordinary
                                  // reflections that happen to lie in the fixed bands, so ordering on it
                                  // would discard a fifth of the strongest spots for nothing.
                                  return std::tuple{deprioritise_ice && s.ice_ring, -s.intensity};
                              });
    input.resize(output_size);
}

void FilterSpuriousHighResolutionSpots(std::vector<SpotToSave> &spots, float threshold) {
    std::ranges::sort(spots, [](SpotToSave &a, SpotToSave &b) {
        return a.d_A > b.d_A;
    });

    // Apply 1/d gap threshold: find first gap in q = 1/d exceeding dist_threshold and ignore spots after it
    if (spots.size() >= 2 && threshold > 0.0f) {
        size_t cut_index = spots.size(); // default: keep all
        // d_A sorted descending → q = 1/d_A sorted ascending
        // We check consecutive q gaps: Δq_i = (1/d_i) - (1/d_{i+1})
        for (size_t i = 0; i + 1 < spots.size(); ++i) {
            float d1 = spots[i].d_A;
            float d2 = spots[i + 1].d_A;
            // Avoid division by zero; d_A should be > 0 in valid data
            if (d1 <= 0.0f || d2 <= 0.0f)
                continue;
            float q1 = 2 * PI / d1;
            float q2 = 2 * PI / d2;
            float dq = q2 - q1; // should be >= 0 due to sorting
            if (dq > threshold) {
                cut_index = i + 1; // keep up to i inclusive
                break;
            }
        }
        if (cut_index < spots.size())
            spots.resize(cut_index);
    }
}

namespace {
    // Fraction of the image's weighted spot signal that is allowed to lie beyond the quantile read
    // off below. A quantile near the middle of the distribution measures the shape of the fall-off,
    // which is the crystal's own; the extreme end of it measures the detection threshold and how many
    // reflections the unit cell puts on the frame, which are not.
    constexpr float SPOT_RESOLUTION_TAIL_FRACTION = 0.30f;
    // How much further in 1/d the merged data reach than that quantile. Merging averages many
    // observations of each reflection, so intensities go on being measurable well past the point where
    // one image's spot finder still detects them. Calibrated on rotation data against the resolution at
    // which per-shell CC1/2 falls through 0.30.
    constexpr float SPOT_RESOLUTION_MERGE_REACH = 2.25f;
    // Fewer spots than this and the quantile is not a fall-off, it is a handful of points.
    constexpr size_t SPOT_RESOLUTION_MIN_SPOTS = 4;
}

std::optional<float> GetResolution(const std::vector<SpotToSave> &spots) {
    // Each spot enters weighted by its own signal-to-noise. The intensity is a summed photon count, so
    // it is Poisson and its significance is sqrt(I): that keeps a marginal high-resolution detection
    // from counting for as much as a real reflection, without letting the handful of very strong
    // low-resolution reflections - which say nothing about how far the crystal diffracts - decide the
    // answer, as weighting by intensity itself would.
    std::vector<std::pair<float, float>> spot_1_over_d2_weight;   // (1/d^2, sqrt(intensity))
    spot_1_over_d2_weight.reserve(spots.size());
    float total_weight = 0.0f;
    for (const auto &spot: spots) {
        if (spot.ice_ring || !(spot.d_A > 0.0f) || !(spot.intensity > 0.0f))
            continue;
        const float weight = std::sqrt(spot.intensity);
        spot_1_over_d2_weight.emplace_back(1.0f / (spot.d_A * spot.d_A), weight);
        total_weight += weight;
    }

    if (spot_1_over_d2_weight.size() < SPOT_RESOLUTION_MIN_SPOTS || !(total_weight > 0.0f))
        return std::nullopt;

    // Walk in from the highest-resolution spot until the tail fraction of the weight is behind us.
    std::ranges::sort(spot_1_over_d2_weight, std::ranges::greater{},
                      [](const std::pair<float, float> &s) { return s.first; });
    float walked = 0.0f;
    float one_over_d2 = spot_1_over_d2_weight.front().first;
    for (const auto &[s, weight]: spot_1_over_d2_weight) {
        walked += weight;
        one_over_d2 = s;
        if (walked >= SPOT_RESOLUTION_TAIL_FRACTION * total_weight)
            break;
    }

    // Not clamped at the corner of the detector. The quantile is read from the middle of the
    // fall-off, so it still measures the crystal where the detector cuts that fall-off short;
    // clamping reported where the detector stops instead, which is the one thing this is not for.
    return 1.0f / (SPOT_RESOLUTION_MERGE_REACH * std::sqrt(one_over_d2));
}

void GenerateSpotPlot(DataMessage &msg, const std::vector<SpotToSave> &spots, float d_min_A) {
    const int nshells = 20;
    // The geometry gives no usable high-resolution corner (no distance or no wavelength), so there is
    // no resolution axis to plot the spots against. ResolutionShells would throw on it, once per image.
    if (d_min_A <= 0.0f || d_min_A >= 50.0f)
        return;
    ResolutionShells shells(d_min_A, 50.0, nshells);

    std::vector<float> intensity(nshells);
    std::vector<float> count(nshells);

    for (const auto &s: spots) {
        if (s.ice_ring)
            continue;

        if (auto shell = shells.GetShell(s.d_A)) {
            intensity[*shell] += s.intensity;
            count[*shell] += 1.0f;
        }
    }

    std::vector<float> result(nshells);
    for (int i = 0; i < nshells; ++i) {
        if (count[i] > 0)
            result[i] = intensity[i] / count[i];
        else
            result[i] = 0.0f;
    }

    msg.spot_plot_one_over_d_square = shells.GetShellMeanOneOverResSq();
    msg.spot_plot_intensity = result;
    msg.spot_plot_count = count;
}

void SpotAnalyze(const DiffractionExperiment &experiment,
                 const SpotFindingSettings &spot_finding_settings,
                 const std::vector<DiffractionSpot> &spots,
                 DataMessage &output) {
    auto geom = experiment.GetDiffractionGeometry();

    std::vector<SpotToSave> spots_out;

    for (const auto &spot: spots) {
        if (auto s = spot.Export(geom, output.number); s.has_value())
            spots_out.push_back(s.value());
    }

    if (spot_finding_settings.high_res_gap_Q_recipA.has_value())
        FilterSpuriousHighResolutionSpots(spots_out, spot_finding_settings.high_res_gap_Q_recipA.value());

    if (experiment.GetDatasetSettings().IsDetectIceRings() && spot_finding_settings.ice_ring_width_Q_recipA > 0.0f) {
        MarkIceRings(spots_out, spot_finding_settings.ice_ring_width_Q_recipA);
        // Before FilterSpotsByCount below, which orders ice spots LAST and would throw them away first.
        output.spot_count_ice_control =
            CountIceRingControlSpots(spots_out, spot_finding_settings.ice_ring_width_Q_recipA);
    }

    CountSpots(output, spots_out, spot_finding_settings.cutoff_spot_count_low_res);

    // The rings this run measured for itself, flagged after the ice counters above so those go on
    // reporting hexagonal ice and only hexagonal ice, and before the resolution estimate and the spot
    // budget below, which both want the contaminant out of the way: a powder ring reaching the corner
    // of the detector otherwise sets the estimate, and the budget otherwise spends itself on it.
    MarkRings(spots_out, spot_finding_settings.measured_ring_q_recipA,
              spot_finding_settings.ice_ring_width_Q_recipA);

    // 0 spells "no limit" everywhere else the limit is read (value_or(0) then compares against it), so it
    // has to mean the same here - passing it on as a resolution makes ResolutionShells throw per image.
    const auto &spot_d_min = spot_finding_settings.high_resolution_limit;
    GenerateSpotPlot(output, spots_out,
                     spot_d_min.value_or(0.0f) > 0 ? *spot_d_min : experiment.GetDetectorMaxResolution_A());

    output.resolution_estimate = GetResolution(spots_out);

    // One decision drives both: if indexing is to use the ice-band spots, the spot budget must not
    // throw them away before it gets the chance.
    FilterSpotsByCount(spots_out, experiment.GetMaxSpotCount(),
                       !experiment.GetIndexingSettings().GetIndexIceRings());

    output.spots = spots_out;
}