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

#include "SpotWidth.h"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <utility>

using namespace spot_width;

namespace {

// The engine reads pixels in the INT32_MIN(masked)/INT32_MAX(saturated) convention.
inline bool valid(int32_t v) { return v != INT32_MIN && v != INT32_MAX; }

// Nothing inside this radius of the beam centre: the beam stop and its halo are not spots.
constexpr float MIN_BEAM_DISTANCE_PX = 60.0f;
// A neighbour this close puts its own flux inside the aperture, which would read as extra width.
constexpr float ISOLATION_PX = 28.0f;
// Spots taken per resolution band per image, strongest first.
constexpr int PER_BAND_PER_IMAGE = 40;
// The r <= 4 px sum must be this many sigma above the background before the tail is believed.
constexpr double SNR_MIN = 15.0;
constexpr int R_CENTROID = 4;
constexpr double MAX_CENTROID_OFFSET_PX = 2.0;
// Spots needed before a band, and the crystal, are characterised at all.
constexpr size_t MIN_SPOTS_PER_BAND = 15;
constexpr size_t MIN_SPOTS_TOTAL = 20;

// Resolution bands, A. The quota is per band, so a crystal is characterised over its whole range
// and not wherever its strongest spots happen to sit.
constexpr int N_BAND = 5;
constexpr std::array<std::pair<float, float>, N_BAND> BANDS = {{
        {2.0f, 3.0f}, {3.0f, 4.5f}, {4.5f, 7.0f}, {7.0f, 12.0f}, {12.0f, 30.0f}}};

int band_of(float d_A) {
    for (int b = 0; b < N_BAND; b++)
        if (d_A >= BANDS[b].first && d_A < BANDS[b].second) return b;
    return -1;
}

// The radius at which the curve reaches `frac`, linearly interpolated. prof[i] is the flux inside
// radius i+1.
float interpolate_radius(double frac, const std::array<float, R_MAX> &prof) {
    if (prof[0] >= frac)
        return prof[0] > 0.0f ? static_cast<float>(frac / prof[0]) : 1.0f;
    for (int i = 1; i < R_MAX; i++)
        if (prof[i] >= frac)
            return static_cast<float>(i + (frac - prof[i - 1]) / (prof[i] - prof[i - 1]));
    return static_cast<float>(R_MAX);
}

double median_of(std::vector<double> &v) {
    if (v.empty()) return 0.0;
    const size_t mid = v.size() / 2;
    std::nth_element(v.begin(), v.begin() + mid, v.end());
    const double hi = v[mid];
    if (v.size() % 2 == 1) return hi;
    return 0.5 * (hi + *std::max_element(v.begin(), v.begin() + mid));
}

}   // namespace

void MeasureSpotFluxCurves(const ImagePreprocessorBuffer &image, int width, int height,
                           const DiffractionGeometry &geometry,
                           const std::vector<DiffractionSpot> &spots,
                           std::vector<FluxCurve> &out) {
    if (spots.empty()) return;

    const float beam_x = geometry.GetBeamX_pxl(), beam_y = geometry.GetBeamY_pxl();

    // Where every spot of this image sits, so isolation can be tested against all of them and not
    // only against the ones that survive the gates below.
    std::vector<Coord> centre(spots.size());
    for (size_t i = 0; i < spots.size(); i++)
        centre[i] = spots[i].RawCoord();

    // Isolation on a grid of ISOLATION_PX cells: a neighbour within that distance is in this cell or
    // one of the eight around it.
    const int gw = static_cast<int>(width / ISOLATION_PX) + 1;
    const int gh = static_cast<int>(height / ISOLATION_PX) + 1;
    std::vector<std::vector<uint32_t>> cell(static_cast<size_t>(gw) * gh);
    const auto cell_of = [&](const Coord &c) {
        const int gx = std::clamp(static_cast<int>(c.x / ISOLATION_PX), 0, gw - 1);
        const int gy = std::clamp(static_cast<int>(c.y / ISOLATION_PX), 0, gh - 1);
        return std::pair<int, int>(gx, gy);
    };
    for (size_t i = 0; i < spots.size(); i++) {
        const auto [gx, gy] = cell_of(centre[i]);
        cell[static_cast<size_t>(gy) * gw + gx].push_back(static_cast<uint32_t>(i));
    }
    const auto isolated = [&](size_t i) {
        const auto [gx, gy] = cell_of(centre[i]);
        for (int y = std::max(0, gy - 1); y <= std::min(gh - 1, gy + 1); y++)
            for (int x = std::max(0, gx - 1); x <= std::min(gw - 1, gx + 1); x++)
                for (uint32_t j : cell[static_cast<size_t>(y) * gw + x]) {
                    if (j == i) continue;
                    if (std::hypot(centre[j].x - centre[i].x, centre[j].y - centre[i].y) < ISOLATION_PX)
                        return false;
                }
        return true;
    };

    // Candidates that pass the geometric gates, by band, strongest first.
    struct Candidate { size_t index; int64_t count; float d_A; };
    std::array<std::vector<Candidate>, N_BAND> candidates;
    for (size_t i = 0; i < spots.size(); i++) {
        const Coord &c = centre[i];
        const int cx = static_cast<int>(std::lround(c.x)), cy = static_cast<int>(std::lround(c.y));
        if (cx < R_BKG_OUT || cy < R_BKG_OUT || cx >= width - R_BKG_OUT || cy >= height - R_BKG_OUT)
            continue;
        if (std::hypot(c.x - beam_x, c.y - beam_y) < MIN_BEAM_DISTANCE_PX) continue;
        const float d_A = geometry.PxlToRes(c.x, c.y);
        const int band = band_of(d_A);
        if (band < 0) continue;
        if (!isolated(i)) continue;
        candidates[band].push_back({i, spots[i].Count(), d_A});
    }

    std::vector<double> ring;
    for (int band = 0; band < N_BAND; band++) {
        auto &cand = candidates[band];
        const size_t take = std::min<size_t>(cand.size(), PER_BAND_PER_IMAGE);
        std::partial_sort(cand.begin(), cand.begin() + take, cand.end(),
                          [](const Candidate &a, const Candidate &b) { return a.count > b.count; });
        for (size_t k = 0; k < take; k++) {
            const Coord &c = centre[cand[k].index];
            const int cx = static_cast<int>(std::lround(c.x)), cy = static_cast<int>(std::lround(c.y));

            // The background under the spot, and a check that the whole aperture is readable: a hole
            // in it removes flux from one radius and not another, which is exactly the shape this
            // measures.
            ring.clear();
            bool readable = true;
            for (int dy = -R_BKG_OUT; dy <= R_BKG_OUT && readable; dy++)
                for (int dx = -R_BKG_OUT; dx <= R_BKG_OUT; dx++) {
                    const int d2 = dx * dx + dy * dy;
                    if (d2 > R_BKG_OUT * R_BKG_OUT) continue;
                    const int32_t px = image[static_cast<size_t>(cy + dy) * width + (cx + dx)];
                    if (!valid(px)) { readable = false; break; }
                    if (d2 >= R_BKG_IN * R_BKG_IN) ring.push_back(px);
                }
            if (!readable || ring.size() < 20) continue;
            const size_t n_ring = ring.size();
            const double bkg = median_of(ring);

            // Flux and centroid over the r <= 4 px core, then the signal-to-noise gate. A weak spot's
            // tail is background, and an encircled-flux curve built on it measures the background.
            double core = 0.0, mx = 0.0, my = 0.0;
            int n_core = 0;
            for (int dy = -R_CENTROID; dy <= R_CENTROID; dy++)
                for (int dx = -R_CENTROID; dx <= R_CENTROID; dx++) {
                    if (dx * dx + dy * dy > R_CENTROID * R_CENTROID) continue;
                    const double v = image[static_cast<size_t>(cy + dy) * width + (cx + dx)] - bkg;
                    core += v;
                    mx += v * dx;
                    my += v * dy;
                    ++n_core;
                }
            if (core <= 0.0) continue;
            const double noise = std::sqrt(core + n_core * std::max(bkg, 0.05)
                                                  * (1.0 + static_cast<double>(n_core) / n_ring));
            if (core / noise < SNR_MIN) continue;
            mx /= core;
            my /= core;
            if (std::abs(mx) > MAX_CENTROID_OFFSET_PX || std::abs(my) > MAX_CENTROID_OFFSET_PX)
                continue;

            // The encircled flux about that centroid, out to the fixed aperture.
            FluxCurve curve;
            curve.d_A = cand[k].d_A;
            for (int dy = -R_MAX; dy <= R_MAX; dy++)
                for (int dx = -R_MAX; dx <= R_MAX; dx++) {
                    const double rc = std::hypot(dx - mx, dy - my);
                    if (rc > R_MAX) continue;
                    const double v = image[static_cast<size_t>(cy + dy) * width + (cx + dx)] - bkg;
                    for (int t = std::max(1, static_cast<int>(std::ceil(rc))); t <= R_MAX; t++)
                        curve.c[t - 1] += static_cast<float>(v);
                }
            if (!(curve.c[R_NORM - 1] > 0.0f) || !(curve.c[R_MAX - 1] > 0.0f)) continue;
            const float norm = curve.c[R_NORM - 1];
            for (float &v : curve.c) v /= norm;
            out.push_back(curve);
        }
    }
}

std::optional<float> spot_width::R80AtReference(const std::vector<FluxCurve> &curves) {
    if (curves.size() < MIN_SPOTS_TOTAL) return std::nullopt;

    // One point per band: the median curve of the band, the radius it holds 80 % of its flux at, and
    // the median resolution it was measured at.
    struct Point { double inv_d; double r80; double weight; };
    std::vector<Point> points;
    std::vector<double> values, band_d;
    for (int b = 0; b < N_BAND; b++) {
        band_d.clear();
        for (const auto &c : curves)
            if (c.d_A >= BANDS[b].first && c.d_A < BANDS[b].second) band_d.push_back(c.d_A);
        if (band_d.size() < MIN_SPOTS_PER_BAND) continue;
        std::array<float, R_MAX> profile{};
        for (int t = 0; t < R_MAX; t++) {
            values.clear();
            for (const auto &c : curves)
                if (c.d_A >= BANDS[b].first && c.d_A < BANDS[b].second) values.push_back(c.c[t]);
            profile[t] = static_cast<float>(median_of(values));
        }
        const double d_med = median_of(band_d);
        if (d_med <= 0.0) continue;
        points.push_back({1.0 / d_med, interpolate_radius(0.8, profile),
                          static_cast<double>(band_d.size())});
    }
    if (points.empty()) return std::nullopt;
    if (points.size() == 1) return static_cast<float>(points[0].r80);

    // The mosaic contribution to the detector footprint grows as 1/d, so r80 is linear in 1/d.
    double sw = 0.0, sx = 0.0, sxx = 0.0, sy = 0.0, sxy = 0.0;
    for (const auto &p : points) {
        sw += p.weight;
        sx += p.weight * p.inv_d;
        sxx += p.weight * p.inv_d * p.inv_d;
        sy += p.weight * p.r80;
        sxy += p.weight * p.inv_d * p.r80;
    }
    const double det = sw * sxx - sx * sx;
    double value = sy / sw;
    if (std::abs(det) > 1e-12) {
        const double c1 = (sw * sxy - sx * sy) / det;
        value = (sy - c1 * sx) / sw + c1 / D_REF_A;
    }
    // Never extrapolate outside what the bands actually measured.
    double lo = std::numeric_limits<double>::max(), hi = 0.0;
    for (const auto &p : points) { lo = std::min(lo, p.r80); hi = std::max(hi, p.r80); }
    return static_cast<float>(std::clamp(value, 0.8 * lo, 1.25 * hi));
}

float spot_width::R1ForWidth(float r80) {
    return std::clamp(std::round(2.0f * r80), 4.0f, 6.0f);
}