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

#include "BeamCenterFromBackground.h"

#include <algorithm>
#include <cmath>
#include <thread>

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

namespace {

// The band the background is fitted over. The low-resolution end sits outside the beam stop
// and its penumbra, the high-resolution end where the solvent ring has died away.
constexpr float BAND_LOW_RES_A = 12.0f;
constexpr float BAND_HIGH_RES_A = 2.2f;

constexpr int SECTORS = 36;
constexpr int RADIAL_BINS = 120;

// A cell with fewer pixels than this has no usable mean.
constexpr int MIN_PIXELS_PER_CELL = 20;

// A radial bin missing more azimuth than this is a partial ring - it leaves the detector, or a
// module gap eats it - and a partial ring biases the profile it is compared against.
constexpr float MIN_SECTOR_COVERAGE = 0.85f;

// Fractions of a bin's pixels are Bragg peaks. Two rounds of clipping at the upper 2 sigma take
// the mean back to the background without needing the pixel values a second time.
constexpr float CLIP_SIGMA = 2.0f;
constexpr int CLIP_ROUNDS = 2;

// The profile is rebuilt at the trial centre every iteration, so a centre that is off smears the
// solvent ring and flattens g', which over-estimates the shift. Half steps damp that; the fixed
// point is unchanged, only the path to it.
constexpr float DAMPING = 0.5f;
constexpr float CONVERGED_PXL = 0.02f;

// The step is compared with the one before it, and a walk whose step reverses twice running has
// crossed its fixed point rather than walked toward it: with DAMPING = 0.5 that happens where the
// local gain exceeds 2, which is a period-2 limit cycle the walk will not leave. Instrumented, a
// sweep that spent 96 of its 100 iterations in one oscillated by 0.08 px about a centre it knew to
// 1.7 px. Travel does not reverse - a fit walking 339 px kept its step direction through all 22 of
// its travelling iterations - so this ends the dithering without shortening the budget below.
constexpr int REVERSALS_AT_THE_FIXED_POINT = 2;

// A travel budget, not a convergence criterion. The shift a sector's regression can report is
// bounded by the width of the features it reads - a sector whose profile has moved a long way is
// not g + d*g' for any d, and the least-squares projection onto g' returns far less than the true
// d - so the walk advances by a bounded distance per iteration, of the order of twenty pixels,
// however far it still has to go. Ten iterations therefore cap the fit at about two hundred
// pixels of travel and a centre further out than that is left part way there, still walking, with
// the per-iteration precision reported as though it had arrived. The count is set to cross a
// detector instead; a centre that is already close leaves on CONVERGED_PXL, or on the reversal
// test above, after a handful.
constexpr int MAX_ITERATIONS = 100;

// The two passes over the pixels are split into this many row blocks, each accumulating into cells
// of its own, and the blocks are folded in block order. The split is a property of the image and
// not of the machine, so the sums are grouped the same way however many threads ran them.
constexpr int BLOCKS = 64;

// Below these the fit has not seen enough of the detector to be believed at all.
constexpr int MIN_USABLE_SECTORS = SECTORS * 3 / 5;
constexpr int MIN_USABLE_RADIAL_BINS = 15;

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

}   // namespace

std::optional<BeamCenterEstimate>
FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const PixelMask &mask,
                             const std::vector<float> &mean, size_t nthreads) {
    if (nthreads == 0)
        nthreads = std::max(1u, std::thread::hardware_concurrency());
    const auto W = static_cast<int>(experiment.GetXPixelsNumConv());
    const auto H = static_cast<int>(experiment.GetYPixelsNumConv());
    const size_t n_pixels = static_cast<size_t>(W) * H;
    if (mean.size() != n_pixels)
        return {};

    const auto &pixel_mask = mask.GetMask(experiment);
    auto geom = experiment.GetDiffractionGeometry();

    const float wavelength = geom.GetWavelength_A();
    const float sin_high = wavelength / (2.0f * BAND_HIGH_RES_A);
    if (sin_high >= 1.0f)
        return {};
    const float tt_lo = 2.0f * std::asin(wavelength / (2.0f * BAND_LOW_RES_A));
    const float tt_hi = 2.0f * std::asin(sin_high);
    const float d_tt = (tt_hi - tt_lo) / RADIAL_BINS;

    const auto rot = geom.GetDetectorMatrix().arr();   // row major
    const float pixel_size = geom.GetPixelSize_mm();
    const float distance = geom.GetDetectorDistance_mm();

    float beam_x = geom.GetBeamX_pxl();
    float beam_y = geom.GetBeamY_pxl();

    constexpr int n_cells = RADIAL_BINS * SECTORS;
    std::vector<int32_t> cell_of(n_pixels);
    std::vector<double> sum(n_cells), sum_sq(n_cells), sum_jx(n_cells), sum_jy(n_cells);
    std::vector<int32_t> count(n_cells), count_all(n_cells);
    std::vector<float> profile(RADIAL_BINS), d_profile(RADIAL_BINS), clip_limit(n_cells);
    std::vector<char> radial_ok(RADIAL_BINS);

    // One set of cells per block, allocated once and reused by every iteration.
    std::vector<int> block_row(BLOCKS + 1);
    for (int b = 0; b <= BLOCKS; b++)
        block_row[b] = static_cast<int>(static_cast<int64_t>(b) * H / BLOCKS);
    std::vector<double> block_sum(static_cast<size_t>(BLOCKS) * n_cells);
    std::vector<double> block_sum_sq(static_cast<size_t>(BLOCKS) * n_cells);
    std::vector<double> block_jx(static_cast<size_t>(BLOCKS) * n_cells);
    std::vector<double> block_jy(static_cast<size_t>(BLOCKS) * n_cells);
    std::vector<int32_t> block_count(static_cast<size_t>(BLOCKS) * n_cells);

    float step_x = 0.0f, step_y = 0.0f, sigma_x = 0.0f, sigma_y = 0.0f;
    float previous_x = 0.0f, previous_y = 0.0f;
    int reversals = 0;
    for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++) {
        ParallelFor(BLOCKS, nthreads, [&](int b) {
            double *b_sum = block_sum.data() + static_cast<size_t>(b) * n_cells;
            double *b_sum_sq = block_sum_sq.data() + static_cast<size_t>(b) * n_cells;
            double *b_jx = block_jx.data() + static_cast<size_t>(b) * n_cells;
            double *b_jy = block_jy.data() + static_cast<size_t>(b) * n_cells;
            int32_t *b_count = block_count.data() + static_cast<size_t>(b) * n_cells;
            std::fill(b_sum, b_sum + n_cells, 0.0);
            std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0);
            std::fill(b_jx, b_jx + n_cells, 0.0);
            std::fill(b_jy, b_jy + n_cells, 0.0);
            std::fill(b_count, b_count + n_cells, 0);
            for (int y = block_row[b]; y < block_row[b + 1]; y++) {
                for (int x = 0; x < W; x++) {
                    const size_t i = static_cast<size_t>(y) * W + x;
                    cell_of[i] = -1;
                    if (pixel_mask[i] != 0 || !std::isfinite(mean[i]))
                        continue;
                    const float u = (x - beam_x) * pixel_size;
                    const float v = (y - beam_y) * pixel_size;
                    const float lx = rot[0] * u + rot[1] * v + rot[2] * distance;
                    const float ly = rot[3] * u + rot[4] * v + rot[5] * distance;
                    const float lz = rot[6] * u + rot[7] * v + rot[8] * distance;
                    const float rho = std::sqrt(lx * lx + ly * ly);
                    const float two_theta = std::atan2(rho, lz);
                    if (two_theta < tt_lo || two_theta >= tt_hi || rho == 0.0f)
                        continue;

                    const float phi = std::atan2(ly, lx);
                    // Both bins are clamped: a pixel one float ulp below the top of the band divides
                    // to exactly RADIAL_BINS, which is one cell past the end of every accumulator.
                    const int r_bin = std::clamp(static_cast<int>((two_theta - tt_lo) / d_tt), 0, RADIAL_BINS - 1);
                    const int s_bin = std::clamp(static_cast<int>((phi + PI) / (2 * PI) * SECTORS), 0, SECTORS - 1);
                    const int cell = r_bin * SECTORS + s_bin;

                    // d(2theta)/d(beam), through the lab coordinate: the detector coordinate depends
                    // on the centre only as (x - beam_x), so moving the centre is moving the pixel.
                    const float denominator = rho * rho + lz * lz;
                    const float g_x = lz * lx / (rho * denominator);
                    const float g_y = lz * ly / (rho * denominator);
                    const float g_z = -rho / denominator;
                    cell_of[i] = cell;
                    b_count[cell]++;
                    b_sum[cell] += mean[i];
                    b_sum_sq[cell] += static_cast<double>(mean[i]) * mean[i];
                    b_jx[cell] += -pixel_size * (g_x * rot[0] + g_y * rot[3] + g_z * rot[6]);
                    b_jy[cell] += -pixel_size * (g_x * rot[1] + g_y * rot[4] + g_z * rot[7]);
                }
            }
        });
        for (int c = 0; c < n_cells; c++) {
            double s = 0, ss = 0, jx = 0, jy = 0;
            int32_t n = 0;
            for (int b = 0; b < BLOCKS; b++) {
                const size_t k = static_cast<size_t>(b) * n_cells + c;
                s += block_sum[k]; ss += block_sum_sq[k];
                jx += block_jx[k]; jy += block_jy[k];
                n += block_count[k];
            }
            sum[c] = s; sum_sq[c] = ss; sum_jx[c] = jx; sum_jy[c] = jy; count[c] = n;
        }

        count_all = count;   // the Jacobian sums belong to the unclipped pixel set

        for (int round = 0; round < CLIP_ROUNDS; round++) {
            for (int c = 0; c < n_cells; c++) {
                if (count[c] < MIN_PIXELS_PER_CELL) { clip_limit[c] = -1.0f; continue; }
                const double m = sum[c] / count[c];
                const double variance = std::max(sum_sq[c] / count[c] - m * m, 0.0);
                clip_limit[c] = static_cast<float>(m + CLIP_SIGMA * std::sqrt(variance));
            }
            ParallelFor(BLOCKS, nthreads, [&](int b) {
                double *b_sum = block_sum.data() + static_cast<size_t>(b) * n_cells;
                double *b_sum_sq = block_sum_sq.data() + static_cast<size_t>(b) * n_cells;
                int32_t *b_count = block_count.data() + static_cast<size_t>(b) * n_cells;
                std::fill(b_sum, b_sum + n_cells, 0.0);
                std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0);
                std::fill(b_count, b_count + n_cells, 0);
                const size_t lo = static_cast<size_t>(block_row[b]) * W;
                const size_t hi = static_cast<size_t>(block_row[b + 1]) * W;
                for (size_t i = lo; i < hi; i++) {
                    const int32_t c = cell_of[i];
                    if (c < 0 || clip_limit[c] < 0.0f || mean[i] > clip_limit[c])
                        continue;
                    b_count[c]++;
                    b_sum[c] += mean[i];
                    b_sum_sq[c] += static_cast<double>(mean[i]) * mean[i];
                }
            });
            for (int c = 0; c < n_cells; c++) {
                double s = 0, ss = 0;
                int32_t n = 0;
                for (int b = 0; b < BLOCKS; b++) {
                    const size_t k = static_cast<size_t>(b) * n_cells + c;
                    s += block_sum[k]; ss += block_sum_sq[k]; n += block_count[k];
                }
                sum[c] = s; sum_sq[c] = ss; count[c] = n;
            }
        }

        // Radial profile: the median over the sectors that have a mean, on rings that are
        // almost fully covered.
        int usable_radial = 0;
        for (int r = 0; r < RADIAL_BINS; r++) {
            std::vector<float> present;
            for (int s = 0; s < SECTORS; s++)
                if (count[r * SECTORS + s] >= MIN_PIXELS_PER_CELL)
                    present.push_back(static_cast<float>(sum[r * SECTORS + s] / count[r * SECTORS + s]));
            radial_ok[r] = static_cast<float>(present.size()) >= MIN_SECTOR_COVERAGE * SECTORS;
            profile[r] = radial_ok[r] ? median_of(present) : 0.0f;
            usable_radial += radial_ok[r];
        }
        if (usable_radial < MIN_USABLE_RADIAL_BINS)
            return {};
        // Central difference, so a bin next to a gap in the profile drops out with it. The test
        // reads the ring BEFORE it, so it has to read the covered/not-covered flags as they were,
        // not as this same loop has already rewritten them.
        const std::vector<char> covered = radial_ok;
        for (int r = 0; r < RADIAL_BINS; r++) {
            const bool have = r > 0 && r + 1 < RADIAL_BINS && covered[r - 1] && covered[r] && covered[r + 1];
            d_profile[r] = have ? (profile[r + 1] - profile[r - 1]) / (2 * d_tt) : 0.0f;
            radial_ok[r] = have;
        }

        // Per sector: regress (profile of the sector - common profile) on {g, g'}. The first
        // coefficient is the sector's amplitude, the second its radial shift; only the shift
        // is carried on.
        std::vector<float> shift, weight, jacobian_x, jacobian_y;
        for (int s = 0; s < SECTORS; s++) {
            double a11 = 0, a12 = 0, a22 = 0, b1 = 0, b2 = 0;
            double jx = 0, jy = 0;
            int n = 0;
            for (int r = 0; r < RADIAL_BINS; r++) {
                const int c = r * SECTORS + s;
                if (!radial_ok[r] || count[c] < MIN_PIXELS_PER_CELL)
                    continue;
                const double g = profile[r], dg = d_profile[r];
                const double y = sum[c] / count[c] - profile[r];
                a11 += g * g; a12 += g * dg; a22 += dg * dg;
                b1 += g * y; b2 += dg * y;
                jx += sum_jx[c] / count_all[c];
                jy += sum_jy[c] / count_all[c];
                n++;
            }
            const double det = a11 * a22 - a12 * a12;
            if (n < MIN_USABLE_RADIAL_BINS || det <= 0)
                continue;
            const double amplitude = (a22 * b1 - a12 * b2) / det;
            const double this_shift = (a11 * b2 - a12 * b1) / det;
            // Residual sum of squares from the normal equations, without a second pass.
            double residual = 0;
            for (int r = 0; r < RADIAL_BINS; r++) {
                const int c = r * SECTORS + s;
                if (!radial_ok[r] || count[c] < MIN_PIXELS_PER_CELL)
                    continue;
                const double e = sum[c] / count[c] - profile[r] - amplitude * profile[r] - this_shift * d_profile[r];
                residual += e * e;
            }
            const double variance = residual / (n - 2) * (a11 / det);
            if (!(variance > 0))
                continue;
            shift.push_back(static_cast<float>(this_shift));
            weight.push_back(static_cast<float>(1.0 / variance));
            jacobian_x.push_back(static_cast<float>(jx / n));
            jacobian_y.push_back(static_cast<float>(jy / n));
        }
        if (static_cast<int>(shift.size()) < MIN_USABLE_SECTORS)
            return {};

        // shift_k = Jx_k dx + Jy_k dy, robustified so one bad sector cannot carry the answer.
        std::vector<float> w = weight;
        double c11 = 0, c12 = 0, c22 = 0;
        for (int round = 0; round < 3; round++) {
            c11 = c12 = c22 = 0;
            double r1 = 0, r2 = 0;
            for (size_t k = 0; k < shift.size(); k++) {
                c11 += w[k] * jacobian_x[k] * jacobian_x[k];
                c12 += w[k] * jacobian_x[k] * jacobian_y[k];
                c22 += w[k] * jacobian_y[k] * jacobian_y[k];
                r1 += w[k] * jacobian_x[k] * shift[k];
                r2 += w[k] * jacobian_y[k] * shift[k];
            }
            const double det = c11 * c22 - c12 * c12;
            if (det <= 0)
                return {};
            step_x = static_cast<float>((c22 * r1 - c12 * r2) / det);
            step_y = static_cast<float>((c11 * r2 - c12 * r1) / det);
            std::vector<float> residual(shift.size());
            for (size_t k = 0; k < shift.size(); k++)
                residual[k] = shift[k] - jacobian_x[k] * step_x - jacobian_y[k] * step_y;
            std::vector<float> absolute(residual.size());
            for (size_t k = 0; k < residual.size(); k++) absolute[k] = std::abs(residual[k]);
            const float scale = 1.4826f * median_of(absolute) + 1e-30f;
            for (size_t k = 0; k < shift.size(); k++) {
                const float t = residual[k] / (3 * scale);
                w[k] = weight[k] / (1.0f + t * t);
            }
        }

        double chi2 = 0;
        for (size_t k = 0; k < shift.size(); k++) {
            const double e = shift[k] - jacobian_x[k] * step_x - jacobian_y[k] * step_y;
            chi2 += w[k] * e * e;
        }
        chi2 = std::max(chi2 / (shift.size() - 2), 1.0);
        const double det = c11 * c22 - c12 * c12;
        sigma_x = static_cast<float>(std::sqrt(c22 / det * chi2));
        sigma_y = static_cast<float>(std::sqrt(c11 / det * chi2));

        reversals = (step_x * previous_x + step_y * previous_y < 0.0f) ? reversals + 1 : 0;
        previous_x = step_x;
        previous_y = step_y;

        beam_x += DAMPING * step_x;
        beam_y += DAMPING * step_y;
        if (std::hypot(step_x, step_y) < CONVERGED_PXL || reversals >= REVERSALS_AT_THE_FIXED_POINT)
            break;
    }

    // A fit that leaves on the iteration cap has not converged - it was still walking when it ran
    // out - and the precision of its last step is not what it knows the centre to. The step it
    // still wanted to take is a floor under what is left, so report that instead: it turns a
    // confidently wrong answer into one the caller's sigma gate refuses. A fit that stopped on
    // CONVERGED_PXL or at its fixed point takes a step far under any sigma worth reporting, so
    // this never touches those.
    return BeamCenterEstimate{beam_x, beam_y,
                              std::max({sigma_x, sigma_y, std::hypot(step_x, step_y)})};
}