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

#include "BackgroundAperture.h"

#include <algorithm>
#include <cmath>

std::vector<float> background_aperture::CandidateR3(float r2, float max_r3) {
    std::vector<float> out;
    for (float r3 = r2 + 2.0f; r3 <= max_r3 + 1e-3f; r3 += 1.0f)
        out.push_back(r3);
    return out;
}

void background_aperture::MeasureFrame(ImagePreprocessorBuffer &image, int width, int height,
                                       const std::vector<DiffractionSpot> &spots, float r1, float r2,
                                       const std::vector<float> &r3, float clip_nsigma, float gain,
                                       Tally &tally) {
    if (r3.empty())
        return;
    tally.sum_d2.resize(r3.size(), 0.0);
    const float r_max = r3.back();
    const int half = static_cast<int>(std::ceil(r_max));
    const int side = 2 * half + 1;

    // A grid of centres, at least one aperture apart, kept clear of every found spot.
    const double area = static_cast<double>(width) * height;
    const int step = std::max(side, static_cast<int>(std::ceil(std::sqrt(area / MAX_APERTURES_PER_FRAME))));
    const float clear = r_max + SPOT_CLEARANCE_PX;
    std::vector<int32_t> cx, cy;
    for (int y = half; y + half < height; y += step)
        for (int x = half; x + half < width; x += step) {
            bool near_spot = false;
            for (const auto &s : spots) {
                const Coord c = s.RawCoord();
                if ((c.x - x) * (c.x - x) + (c.y - y) * (c.y - y) < clear * clear) {
                    near_spot = true;
                    break;
                }
            }
            if (!near_spot) {
                cx.push_back(x);
                cy.push_back(y);
            }
        }
    if (cx.empty())
        return;
    const int32_t *tiles = image.GatherTiles(cx, cy, width, half);

    std::vector<double> ring_sum(r3.size());
    std::vector<int> ring_n(r3.size());
    std::vector<int32_t> ring_px;
    std::vector<float> ring_d2;
    for (size_t a = 0; a < cx.size(); ++a) {
        const int32_t *t = tiles + a * static_cast<size_t>(side) * side;
        double disk_sum = 0.0;
        int disk_n = 0;
        bool readable = true;
        ring_px.clear();
        ring_d2.clear();
        for (int dy = -half; dy <= half && readable; ++dy)
            for (int dx = -half; dx <= half; ++dx) {
                const float d2 = static_cast<float>(dx * dx + dy * dy);
                if (d2 >= r_max * r_max)
                    continue;
                const int32_t v = t[(dy + half) * side + dx + half];
                if (v == INT32_MIN || v == INT32_MAX) {
                    readable = false;
                    break;
                }
                if (d2 < r1 * r1) {
                    disk_sum += v;
                    ++disk_n;
                } else if (d2 >= r2 * r2) {
                    ring_px.push_back(v);
                    ring_d2.push_back(d2);
                }
            }
        if (!readable || disk_n == 0)
            continue;
        // Every candidate ring as the integrator reads it: its mean, then one high-side clip pass.
        bool measured = true;
        for (size_t k = 0; k < r3.size() && measured; ++k) {
            double sum = 0.0;
            int n = 0;
            for (size_t p = 0; p < ring_px.size(); ++p)
                if (ring_d2[p] < r3[k] * r3[k]) {
                    sum += ring_px[p];
                    ++n;
                }
            if (n == 0) {
                measured = false;
                break;
            }
            const double mean = sum / n;
            const double thr = mean + clip_nsigma * std::sqrt(gain * std::max(mean, 1.0));
            ring_sum[k] = 0.0;
            ring_n[k] = 0;
            for (size_t p = 0; p < ring_px.size(); ++p)
                if (ring_d2[p] < r3[k] * r3[k] && (clip_nsigma <= 0.0f || ring_px[p] <= thr)) {
                    ring_sum[k] += ring_px[p];
                    ++ring_n[k];
                }
            measured = ring_n[k] > 0;
        }
        if (!measured)
            continue;
        const double disk = disk_sum / disk_n;
        for (size_t k = 0; k < r3.size(); ++k) {
            const double d = disk - ring_sum[k] / ring_n[k];
            tally.sum_d2[k] += d * d;
        }
        ++tally.apertures;
    }
}

int background_aperture::Best(const Tally &tally) {
    if (tally.apertures == 0 || tally.sum_d2.empty())
        return -1;
    return static_cast<int>(std::min_element(tally.sum_d2.begin(), tally.sum_d2.end()) - tally.sum_d2.begin());
}