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

#include "SpotFootprint.h"

#include <algorithm>
#include <cmath>
#include <map>
#include <tuple>

namespace {

// The window around a spot is BRAGG_FOOTPRINT_NSIGMA-like: three of its standard deviations, at least
// a few pixels and at most this many, which is wider than any spot the integrator could hold.
constexpr float WINDOW_NSIGMA = 3.0f;
constexpr float WINDOW_MIN_PX = 3.0f;
constexpr float WINDOW_MAX_PX = 24.0f;
// The background is the median of an elliptical annulus between these multiples of the window.
constexpr float BKG_INNER = 1.5f;
constexpr float BKG_OUTER = 2.2f;
constexpr int ITERATIONS = 8;

inline bool valid(int32_t v) { return v != INT32_MIN && v != INT32_MAX; }

float median_of(std::vector<float> &v) {
    const size_t m = v.size() / 2;
    std::nth_element(v.begin(), v.begin() + static_cast<std::ptrdiff_t>(m), v.end());
    return v[m];
}

}   // namespace

void MeasureFootprintSpots(const int32_t *img, int width, int height, float beam_x, float beam_y,
                           const std::vector<float> &x, const std::vector<float> &y,
                           std::vector<FootprintSpot> &out) {
    const int half = static_cast<int>(std::ceil(BKG_OUTER * WINDOW_MAX_PX)) + 1;
    std::vector<float> ring;
    for (size_t s = 0; s < x.size(); ++s) {
        const float rx = x[s] - beam_x, ry = y[s] - beam_y;
        const float r = std::sqrt(rx * rx + ry * ry);
        if (!(r > 1.0f)) continue;
        const float ux = rx / r, uy = ry / r;
        const int ix = static_cast<int>(std::lround(x[s])), iy = static_cast<int>(std::lround(y[s]));
        if (ix - half < 0 || iy - half < 0 || ix + half >= width || iy + half >= height) continue;

        // Start from a compact spot at the prediction; each round re-centres on the signal and takes
        // the window to three of the widths just measured.
        float cx = x[s], cy = y[s];
        float s2r = 1.0f, s2t = 1.0f;
        bool ok = true;
        for (int it = 0; it < ITERATIONS && ok; ++it) {
            const float wr = std::clamp(WINDOW_NSIGMA * std::sqrt(s2r), WINDOW_MIN_PX, WINDOW_MAX_PX);
            const float wt = std::clamp(WINDOW_NSIGMA * std::sqrt(s2t), WINDOW_MIN_PX, WINDOW_MAX_PX);
            const float reach = BKG_OUTER * std::max(wr, wt);
            const int x0 = static_cast<int>(std::floor(cx - reach)), x1 = static_cast<int>(std::ceil(cx + reach));
            const int y0 = static_cast<int>(std::floor(cy - reach)), y1 = static_cast<int>(std::ceil(cy + reach));
            if (x0 < 0 || y0 < 0 || x1 >= width || y1 >= height) { ok = false; break; }

            ring.clear();
            for (int py = y0; py <= y1; ++py)
                for (int px = x0; px <= x1; ++px) {
                    const float dx = px - cx, dy = py - cy;
                    const float rad = dx * ux + dy * uy, tn = -dx * uy + dy * ux;
                    const float e = rad * rad / (wr * wr) + tn * tn / (wt * wt);
                    const int32_t v = img[static_cast<size_t>(py) * width + px];
                    if (e >= BKG_INNER * BKG_INNER && e < BKG_OUTER * BKG_OUTER && valid(v))
                        ring.push_back(static_cast<float>(v));
                }
            if (ring.size() < 10) { ok = false; break; }
            const double bkg = median_of(ring);

            double w = 0.0, mr = 0.0, mt = 0.0, m2r = 0.0, m2t = 0.0;
            for (int py = y0; py <= y1 && ok; ++py)
                for (int px = x0; px <= x1; ++px) {
                    const float dx = px - cx, dy = py - cy;
                    const float rad = dx * ux + dy * uy, tn = -dx * uy + dy * ux;
                    if (rad * rad / (wr * wr) + tn * tn / (wt * wt) >= 1.0f) continue;
                    const int32_t v = img[static_cast<size_t>(py) * width + px];
                    if (!valid(v)) { ok = false; break; }
                    const double net = static_cast<double>(v) - bkg;
                    w += net;
                    mr += net * rad;
                    mt += net * tn;
                    m2r += net * rad * rad;
                    m2t += net * tn * tn;
                }
            if (!ok || !(w > 0.0)) { ok = false; break; }
            const double cr = mr / w, ct = mt / w;
            s2r = static_cast<float>(std::max(0.25, m2r / w - cr * cr));
            s2t = static_cast<float>(std::max(0.25, m2t / w - ct * ct));
            cx += static_cast<float>(cr * ux - ct * uy);
            cy += static_cast<float>(cr * uy + ct * ux);
        }
        if (ok)
            out.push_back({r, std::sqrt(s2r), std::sqrt(s2t)});
    }
}

SpotFootprint FootprintFromSpots(const std::vector<FootprintSpot> &spots, float r_max) {
    SpotFootprint fp;
    if (!(r_max > 0.0f)) return fp;
    const float bin = r_max / FOOTPRINT_BINS;
    std::vector<std::vector<float>> rad(FOOTPRINT_BINS), tan(FOOTPRINT_BINS);
    for (const auto &s : spots) {
        const int b = std::clamp(static_cast<int>(s.r_px / bin), 0, FOOTPRINT_BINS - 1);
        rad[b].push_back(s.sigma_rad);
        tan[b].push_back(s.sigma_tan);
    }
    std::vector<int> filled;
    std::vector<float> mr(FOOTPRINT_BINS, 0.0f), mt(FOOTPRINT_BINS, 0.0f);
    for (int b = 0; b < FOOTPRINT_BINS; ++b)
        if (static_cast<int>(rad[b].size()) >= FOOTPRINT_MIN_SPOTS_PER_BIN) {
            mr[b] = median_of(rad[b]);
            mt[b] = median_of(tan[b]);
            filled.push_back(b);
        }
    if (filled.empty()) return fp;
    fp.bin_px = bin;
    for (int b = 0; b < FOOTPRINT_BINS; ++b) {
        // The nearest bin that has enough spots; the inner one on a tie.
        int best = filled.front();
        for (int f : filled)
            if (std::abs(f - b) < std::abs(best - b)) best = f;
        fp.sigma_rad.push_back(mr[best]);
        fp.sigma_tan.push_back(mt[best]);
    }
    return fp;
}

void MeasureFootprintOffsets(const std::vector<SpotToSave> &spots, const std::vector<Reflection> &reflections,
                             float beam_x, float beam_y, std::vector<FootprintOffset> &out) {
    std::map<std::tuple<int, int, int>, const Reflection *> predicted;
    for (const auto &r : reflections)
        predicted[{r.h, r.k, r.l}] = &r;
    for (const auto &s : spots) {
        if (!s.indexed || s.lattice != 0) continue;
        const auto it = predicted.find({static_cast<int>(s.h), static_cast<int>(s.k), static_cast<int>(s.l)});
        if (it == predicted.end()) continue;
        const float px = it->second->predicted_x, py = it->second->predicted_y;
        const float rx = px - beam_x, ry = py - beam_y;
        const float r = std::sqrt(rx * rx + ry * ry);
        if (!(r > 1.0f)) continue;
        const float ux = rx / r, uy = ry / r;
        const float dx = s.x - px, dy = s.y - py;
        out.push_back({r, dx * ux + dy * uy, -dx * uy + dy * ux});
    }
}

SpotFootprint FootprintWithOffsets(const SpotFootprint &widths, std::vector<FootprintOffset> offsets) {
    if (widths.empty()) return widths;
    // Sorted, so the sums below do not depend on the order the frames were measured in.
    std::sort(offsets.begin(), offsets.end());
    const int n = static_cast<int>(widths.sigma_rad.size());
    std::vector<double> s2r(n, 0.0), s2t(n, 0.0);
    std::vector<int> cnt(n, 0);
    for (const auto &o : offsets) {
        const int b = std::clamp(static_cast<int>(o.r_px / widths.bin_px), 0, n - 1);
        s2r[b] += static_cast<double>(o.off_rad) * o.off_rad;
        s2t[b] += static_cast<double>(o.off_tan) * o.off_tan;
        ++cnt[b];
    }
    std::vector<int> filled;
    for (int b = 0; b < n; ++b)
        if (cnt[b] >= FOOTPRINT_MIN_SPOTS_PER_BIN) filled.push_back(b);
    if (filled.empty()) return widths;
    SpotFootprint fp = widths;
    for (int b = 0; b < n; ++b) {
        int best = filled.front();
        for (int f : filled)
            if (std::abs(f - b) < std::abs(best - b)) best = f;
        fp.sigma_rad[b] = static_cast<float>(std::sqrt(widths.sigma_rad[b] * widths.sigma_rad[b] + s2r[best] / cnt[best]));
        fp.sigma_tan[b] = static_cast<float>(std::sqrt(widths.sigma_tan[b] * widths.sigma_tan[b] + s2t[best] / cnt[best]));
    }
    return fp;
}