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

#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) {
        auto spot_q = 2 * PI / s.d_A;
        bool tmp = false;
        for (const auto &q: ice_rings_q)
            tmp |= (fabs(spot_q - q) < tolerance_q_recipA);
        s.ice_ring = tmp;
    }
}

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, float detector_d_min_A) {
    // 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;
    }

    const float d_A = 1.0f / (SPOT_RESOLUTION_MERGE_REACH * std::sqrt(one_over_d2));

    // However far the crystal diffracts, no merge reaches past the corner of the detector.
    return detector_d_min_A > 0.0f ? std::max(d_A, detector_d_min_A) : d_A;
}

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

    // 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, experiment.GetDetectorMaxResolution_A());

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