// 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 <numeric>
#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);
}

void AddToHistogram(const std::vector<int32_t> &image, const std::vector<uint32_t> &mask, int64_t saturation,
                    std::vector<uint64_t> &histogram) {
    const int64_t end = std::min<int64_t>(saturation, PEAK_HISTOGRAM_BINS);
    for (size_t i = 0; i < image.size(); i++)
        if (mask[i] == 0 && image[i] >= 0 && image[i] < end)
            histogram[image[i]]++;
}

std::optional<GainEstimate> PhotonPeak(const std::vector<uint64_t> &histogram) {
    // A float file that holds whole numbers, read at a fraction of its unit, fills only every q-th bin:
    // the histogram is read at that step.
    size_t q = 0;
    for (size_t i = 1; i < histogram.size(); i++)
        if (histogram[i] > 0)
            q = std::gcd(q, i);
    q = std::max<size_t>(q, 1);
    std::vector<uint64_t> h;
    for (size_t i = 0; i < histogram.size(); i += q)
        h.push_back(histogram[i]);

    // A peak stands out of a valley where it is higher by five times its own counting noise.
    const auto stands = [&](size_t valley, size_t peak) {
        return static_cast<double>(h[peak]) - static_cast<double>(h[valley])
               > 5.0 * std::sqrt(static_cast<double>(h[peak]));
    };
    // Most pixels hold no photon. Down the zero peak to the first valley.
    const size_t p0 = std::max_element(h.begin(), h.end()) - h.begin();
    size_t v0 = p0;
    while (v0 + 1 < h.size() && h[v0 + 1] <= h[v0])
        v0++;
    if (v0 + 1 >= h.size())
        return std::nullopt;
    const size_t p1 = std::max_element(h.begin() + v0, h.end()) - h.begin();
    const size_t g = p1 - p0;
    // The zero-photon peak is the detector's noise, so it has a width: a single bin with nothing beside it is
    // a CCD's pixels outside the active area, all at zero, and what follows is its background, not a photon.
    if (h[p0 + 1] == 0 || !stands(v0, p1) || p1 + 3 * g / 2 >= h.size())
        return std::nullopt;
    // A photon is a peak narrower than the spacing: it falls below half its height before the midpoint to
    // the next. A CCD's background is a hump as wide as where it lies.
    if (2 * h[p1 + g / 2] >= h[p1])
        return std::nullopt;
    // The two-photon peak: the histogram rises again, near one more gain on, past the valley after p1.
    const size_t v2 = std::min_element(h.begin() + p1, h.begin() + p1 + g + 1) - h.begin();
    const size_t p2 = std::max_element(h.begin() + v2, h.begin() + p1 + 3 * g / 2 + 1) - h.begin();
    if (!stands(v2, p2) || 4 * (p2 - p1) < 3 * g || 4 * (p2 - p1) > 5 * g)
        return std::nullopt;
    // Within the bin: the vertex of a parabola through the log counts of the peak and its neighbours. A
    // zero peak in the first bin is a subtracted baseline, cut where the values turn negative: zero.
    const auto vertex = [&](size_t p) {
        if (p == 0)
            return 0.0;
        const double a = std::log(std::max<double>(1.0, h[p - 1]));
        const double b = std::log(std::max<double>(1.0, h[p]));
        const double c = std::log(std::max<double>(1.0, h[p + 1]));
        const double curvature = a - 2.0 * b + c;
        return static_cast<double>(p) + (curvature < 0.0 ? 0.5 * (a - c) / curvature : 0.0);
    };
    const double pedestal = vertex(p0);
    return GainEstimate{static_cast<float>(q * (vertex(p1) - pedestal)), static_cast<float>(q * pedestal)};
}

std::optional<GainEstimate> Estimate(const std::vector<Levels> &frames, const std::vector<uint64_t> &histogram) {
    if (const auto peaks = PhotonPeak(histogram))
        return peaks;
    if (const auto g = Gain(frames))
        return GainEstimate{*g, 0.0f};
    return std::nullopt;
}

}  // namespace background_gain