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

#include "BackgroundGain.h"

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

namespace background_gain {

namespace {
constexpr int PAIR_DISTANCE = 2;     // columns between the two pixels of a pair
constexpr size_t ROW_STEP = 4;       // every fourth row: a quarter of a frame is ample for one number
constexpr double CLIP_NSIGMA = 4.0;
constexpr int CLIP_ROUNDS = 3;
const double LOG_BIN = std::log(1.25);
}

Levels MeasureFrame(const std::vector<int32_t> &image, size_t width, size_t height,
                    const std::vector<uint32_t> &mask, int64_t saturation) {
    const auto usable = [&](size_t i) {
        return mask[i] == 0 && image[i] >= 1 && image[i] < saturation;
    };
    // The level of every usable pair, worked out once; -1 where the pair is not usable.
    std::vector<int8_t> bin(width * height, -1);
    for (size_t y = 0; y < height; y += ROW_STEP)
        for (size_t x = 0; x + PAIR_DISTANCE < width; x++) {
            const size_t i = y * width + x;
            if (!usable(i) || !usable(i + PAIR_DISTANCE))
                continue;
            const double m = 0.5 * (static_cast<double>(image[i]) + image[i + PAIR_DISTANCE]);
            bin[i] = static_cast<int8_t>(std::min<int>(N_BINS - 1, static_cast<int>(std::log(m) / LOG_BIN)));
        }

    // Clip about zero at CLIP_NSIGMA of the level's own scatter, starting from no clip at all: a few
    // rounds take the Bragg peaks and zingers out without needing any idea of the gain beforehand.
    std::array<double, N_BINS> limit2;
    limit2.fill(std::numeric_limits<double>::infinity());
    Levels out;
    for (int round = 0; round <= CLIP_ROUNDS; round++) {
        out = Levels{};
        for (size_t i = 0; i < bin.size(); i++) {
            const int b = bin[i];
            if (b < 0)
                continue;
            const double d = static_cast<double>(image[i]) - image[i + PAIR_DISTANCE];
            if (d * d > limit2[b])
                continue;
            out.n[b]++;
            out.sum_mean[b] += 0.5 * (static_cast<double>(image[i]) + image[i + PAIR_DISTANCE]);
            out.sum_diff2[b] += d * d;
        }
        for (int b = 0; b < N_BINS; b++)
            if (out.n[b] > 0)
                limit2[b] = CLIP_NSIGMA * CLIP_NSIGMA * out.sum_diff2[b] / static_cast<double>(out.n[b]);
    }
    return out;
}

std::optional<float> Gain(const std::vector<Levels> &frames) {
    // (variance/mean, pairs) per level
    std::vector<std::pair<double, uint64_t>> ratio;
    uint64_t total = 0;
    for (int b = 0; b < N_BINS; b++) {
        uint64_t n = 0;
        double sum_mean = 0.0, sum_diff2 = 0.0;
        for (const auto &f : frames) {
            n += f.n[b];
            sum_mean += f.sum_mean[b];
            sum_diff2 += f.sum_diff2[b];
        }
        // var(a - b) = 2 var(pixel)
        if (n >= MIN_PAIRS_PER_BIN && sum_mean > 0.0) {
            ratio.emplace_back(0.5 * sum_diff2 / sum_mean, n);
            total += n;
        }
    }
    if (ratio.size() < MIN_BINS)
        return std::nullopt;
    // The median over PAIRS, not over levels: the background holds most of the pixels in a few levels,
    // and the Bragg peaks spread a few pixels over many more - an unweighted median over the levels of
    // a weak-background marCCD frame lands on the peaks and reads 5.7 where the background reads 0.3.
    std::sort(ratio.begin(), ratio.end());
    uint64_t cum = 0;
    for (const auto &[r, n] : ratio) {
        cum += n;
        if (2 * cum >= total)
            return static_cast<float>(r);
    }
    return static_cast<float>(ratio.back().first);
}

}  // namespace background_gain