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

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

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,
                                       const RingBackground &rings, const DiffractionGeometry &geometry,
                                       const DetectorModules &modules, Tally &tally) {
    if (r3.empty())
        return;
    tally.sum_d2.resize(r3.size(), 0.0);
    tally.sum_d2_radial.resize(r3.size(), 0.0);
    tally.sum_dd2.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);

    struct Pixel {
        int32_t v;
        float d2;     // squared distance from the aperture centre
        double b;     // the ring background at this pixel's resolution
    };
    std::vector<Pixel> disk, ring;
    for (size_t a = 0; a < cx.size(); ++a) {
        const int32_t *t = tiles + a * static_cast<size_t>(side) * side;
        const int64_t panel = modules.ModuleOfImagePixel(cx[a], cy[a]);
        bool readable = true;
        disk.clear();
        ring.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 || (d2 >= r1 * r1 && d2 < r2 * r2))
                    continue;
                const int32_t v = t[(dy + half) * side + dx + half];
                if (v == INT32_MIN || v == INT32_MAX) {
                    readable = false;
                    break;
                }
                double b = 0.0;
                if (!rings.empty()) {
                    const auto [px, py] = modules.ImageToPlane(panel, cx[a] + dx, cy[a] + dy);
                    b = rings.At(2.0 * PI / geometry.PxlToRes(static_cast<float>(px), static_cast<float>(py)));
                }
                (d2 < r1 * r1 ? disk : ring).push_back({v, d2, b});
            }
        if (!readable || disk.empty())
            continue;
        double disk_v = 0.0, disk_b = 0.0;
        for (const auto &p : disk) {
            disk_v += p.v;
            disk_b += p.b;
        }
        disk_v /= static_cast<double>(disk.size());
        disk_b /= static_cast<double>(disk.size());
        // The curve at the aperture's centre, which the corrected clip measures the ring's rise from.
        double centre_b = 0.0;
        if (!rings.empty()) {
            const auto [px, py] = modules.ImageToPlane(panel, cx[a], cy[a]);
            centre_b = rings.At(2.0 * PI / geometry.PxlToRes(static_cast<float>(px), static_cast<float>(py)));
        }

        // Every candidate ring as the integrator reads it: the mean, then one high-side clip pass. With the
        // radial correction, as the integrator applies it: the clip follows the curve (each pixel's
        // threshold raised by the curve's rise from the centre to it, in whole counts -
        // RadialClipOffset), and the correction is taken over the ring's geometry, as the kernel is.
        std::vector<double> d_plain(r3.size()), d_radial(r3.size());
        bool measured = true;
        for (size_t k = 0; k < r3.size() && measured; ++k) {
            const float r3_sq = r3[k] * r3[k];
            double sum = 0.0, sum_b = 0.0;
            int n = 0;
            for (const auto &p : ring)
                if (p.d2 < r3_sq) {
                    sum += p.v;
                    sum_b += p.b;
                    ++n;
                }
            if (n == 0) {
                measured = false;
                break;
            }
            const double mean = sum / n;
            const double thr = mean + clip_nsigma * std::sqrt(gain * std::max(mean, static_cast<double>(gain)));
            double kept = 0.0, kept_radial = 0.0;
            int n_kept = 0, n_kept_radial = 0;
            for (const auto &p : ring) {
                if (p.d2 >= r3_sq)
                    continue;
                if (clip_nsigma <= 0.0f || p.v <= thr) {
                    kept += p.v;
                    ++n_kept;
                }
                if (clip_nsigma <= 0.0f || p.v <= thr + static_cast<double>(std::lround(p.b - centre_b))) {
                    kept_radial += p.v;
                    ++n_kept_radial;
                }
            }
            measured = n_kept > 0 && n_kept_radial > 0;
            if (measured) {
                d_plain[k] = disk_v - kept / n_kept;
                d_radial[k] = disk_v - kept_radial / n_kept_radial - (disk_b - sum_b / n);
            }
        }
        if (!measured)
            continue;
        for (size_t k = 0; k < r3.size(); ++k) {
            tally.sum_d2[k] += d_plain[k] * d_plain[k];
            if (!rings.empty()) {
                tally.sum_d2_radial[k] += d_radial[k] * d_radial[k];
                const double dd = d_radial[k] * d_radial[k] - d_plain[k] * d_plain[k];
                tally.sum_dd2[k] += dd * dd;
            }
        }
        ++tally.apertures;
        if (!rings.empty())
            ++tally.apertures_radial;
    }
}

void background_aperture::Tally::Add(const Tally &other) {
    sum_d2.resize(std::max(sum_d2.size(), other.sum_d2.size()), 0.0);
    sum_d2_radial.resize(std::max(sum_d2_radial.size(), other.sum_d2_radial.size()), 0.0);
    for (size_t k = 0; k < other.sum_d2.size(); ++k)
        sum_d2[k] += other.sum_d2[k];
    for (size_t k = 0; k < other.sum_d2_radial.size(); ++k)
        sum_d2_radial[k] += other.sum_d2_radial[k];
    sum_dd2.resize(std::max(sum_dd2.size(), other.sum_dd2.size()), 0.0);
    for (size_t k = 0; k < other.sum_dd2.size(); ++k)
        sum_dd2[k] += other.sum_dd2[k];
    apertures += other.apertures;
    apertures_radial += other.apertures_radial;
}

std::vector<double> background_aperture::Tally::MeanSquare(bool radial) const {
    const size_t n = radial ? apertures_radial : apertures;
    std::vector<double> out;
    if (n == 0)
        return out;
    for (const double v : radial ? sum_d2_radial : sum_d2)
        out.push_back(v / static_cast<double>(n));
    return out;
}

double background_aperture::Tally::RadialZ(size_t k) const {
    if (apertures_radial < 2 || apertures_radial != apertures || k >= sum_dd2.size())
        return 0.0;
    const double n = static_cast<double>(apertures_radial);
    const double mean = (sum_d2_radial[k] - sum_d2[k]) / n;
    const double var = std::max(sum_dd2[k] / n - mean * mean, 0.0);
    return var > 0.0 ? -mean / std::sqrt(var / n) : 0.0;
}

int background_aperture::Best(const std::vector<double> &mean_square) {
    if (mean_square.empty())
        return -1;
    return static_cast<int>(std::min_element(mean_square.begin(), mean_square.end()) - mean_square.begin());
}

bool background_aperture::RadialHelps(const Tally &tally, size_t k, double min_gain) {
    if (tally.apertures == 0 || tally.apertures_radial != tally.apertures || k >= tally.sum_d2.size())
        return false;
    return tally.sum_d2_radial[k] < (1.0 - min_gain) * tally.sum_d2[k];
}