// 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 "../../common/JFJochMath.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 int MAX_ITERATIONS = 10;
constexpr float CONVERGED_PXL = 0.02f;

// 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) {
    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.GetPoniRotMatrix().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);

    float step_x = 0.0f, step_y = 0.0f, sigma_x = 0.0f, sigma_y = 0.0f;
    for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++) {
        std::fill(sum.begin(), sum.end(), 0.0);
        std::fill(sum_sq.begin(), sum_sq.end(), 0.0);
        std::fill(sum_jx.begin(), sum_jx.end(), 0.0);
        std::fill(sum_jy.begin(), sum_jy.end(), 0.0);
        std::fill(count.begin(), count.end(), 0);

        for (int y = 0; y < H; 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;
                count[cell]++;
                sum[cell] += mean[i];
                sum_sq[cell] += static_cast<double>(mean[i]) * mean[i];
                sum_jx[cell] += -pixel_size * (g_x * rot[0] + g_y * rot[3] + g_z * rot[6]);
                sum_jy[cell] += -pixel_size * (g_x * rot[1] + g_y * rot[4] + g_z * rot[7]);
            }
        }

        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));
            }
            std::fill(sum.begin(), sum.end(), 0.0);
            std::fill(sum_sq.begin(), sum_sq.end(), 0.0);
            std::fill(count.begin(), count.end(), 0);
            for (size_t i = 0; i < n_pixels; i++) {
                const int32_t c = cell_of[i];
                if (c < 0 || clip_limit[c] < 0.0f || mean[i] > clip_limit[c])
                    continue;
                count[c]++;
                sum[c] += mean[i];
                sum_sq[c] += static_cast<double>(mean[i]) * mean[i];
            }
        }

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

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

    return BeamCenterEstimate{beam_x, beam_y, std::max(sigma_x, sigma_y)};
}