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

#pragma once

// Where one pixel falls in the background beam-centre fit (FindBeamCenterFromBackground): which
// radial-bin x sector cell of the fitted band it lands in about a trial centre, and the derivative
// of its 2theta with respect to that centre. Written once and compiled both for the host fit and for
// its device twin (BeamCenterBackgroundGPU), so the two read the same formula.
//
// The same formula is the same number only when both sides evaluate it the same way. The library
// atan2f differs between glibc and CUDA in the last bit, and the compilers fuse multiply-adds each
// in their own way; either moves a pixel within a rounding of a cell edge into the neighbouring cell
// (measured on 16 Mpx sweeps: ~30 of 6.5 million pixels, enough to move the fitted centre by up to
// 0.05 px). So the angles come from BackgroundAtan2 below and both translation units are compiled
// without contraction (geom_refinement/CMakeLists.txt), and the host and the device then agree to
// the bit.

#include <cmath>
#include <cstdint>

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

#ifdef __CUDACC__
#define BACKGROUND_BAND_HD __host__ __device__ inline
#else
#define BACKGROUND_BAND_HD inline
#endif

struct BackgroundBand {
    static constexpr int SECTORS = 36;
    static constexpr int RADIAL_BINS = 120;
    static constexpr int CELLS = RADIAL_BINS * SECTORS;

    float rot[9];          // detector matrix, row major
    float pixel_size;      // mm
    float distance;        // mm
    float tt_lo, tt_hi;    // the band in 2theta
    float d_tt;            // width of one radial bin
    double tan_lo, tan_hi; // the band in tan(2theta), widened for the quick rejection
};

// A panel of the image and its place on the plane (DetectorModule), in a form the device can hold.
struct BackgroundPanel {
    int x0, y0, width, height;
    double ox, oy, fx, fy, sx, sy;
};

// Plane position of image pixel (x, y) of panel p - DetectorModule::ImageToPlane, written here so that the
// host and the device evaluate it the same way.
BACKGROUND_BAND_HD void BackgroundPlanePoint(const BackgroundPanel &p, int x, int y, float &px, float &py) {
    const double u = x - p.x0, v = y - p.y0;
    px = static_cast<float>(p.ox + u * p.fx + v * p.sx);
    py = static_cast<float>(p.oy + u * p.fy + v * p.sy);
}

// atan2(y, x) from IEEE operations alone - add, multiply, divide, square root, all correctly rounded on
// the host and on the device - so that, compiled without contraction (see CMakeLists.txt), the host
// and the device return the same bits. The library atan2f does not: glibc's and CUDA's differ in the
// last place. Two half-angle reductions take the argument below tan(pi/16), where eleven terms of the
// series leave an error under 1e-17.
BACKGROUND_BAND_HD double BackgroundAtan2(double y, double x) {
    const double ax = x < 0 ? -x : x, ay = y < 0 ? -y : y;
    if (ax == 0.0 && ay == 0.0)
        return 0.0;
    const bool swap = ay > ax;
    double t = swap ? ax / ay : ay / ax;
    t = t / (1.0 + sqrt(1.0 + t * t));
    t = t / (1.0 + sqrt(1.0 + t * t));
    const double t2 = t * t;
    double series = 1.0 / 23.0;
    for (int k = 21; k >= 1; k -= 2)
        series = 1.0 / k - t2 * series;
    double a = 4.0 * t * series;
    if (swap) a = PI / 2 - a;
    if (x < 0) a = PI - a;
    return y < 0 ? -a : a;
}

// The cell of the pixel at (x, y) on the plane about (beam_x, beam_y), or -1 when it is outside the band;
// for a pixel in it, also the two components of the derivative of its 2theta with respect to the centre.
BACKGROUND_BAND_HD int BackgroundBandCell(const BackgroundBand &b, float x, float y, float beam_x, float beam_y,
                                          float &jac_x, float &jac_y) {
    const float *rot = b.rot;
    const float u = (x - beam_x) * b.pixel_size;
    const float v = (y - beam_y) * b.pixel_size;
    const float lx = rot[0] * u + rot[1] * v + rot[2] * b.distance;
    const float ly = rot[3] * u + rot[4] * v + rot[5] * b.distance;
    const float lz = rot[6] * u + rot[7] * v + rot[8] * b.distance;
    const float rho_sq = lx * lx + ly * ly;
    // Most pixels lie outside the band. Those clearly outside it in tan(2theta) = rho / lz - by a
    // margin far above float rounding - skip the square root and both atan2 below; every pixel the
    // exact test would keep still reaches it.
    if (lz > 0.0f) {
        const double lz_sq = static_cast<double>(lz) * lz;
        if (rho_sq < b.tan_lo * b.tan_lo * lz_sq || rho_sq > b.tan_hi * b.tan_hi * lz_sq)
            return -1;
    }
    const float rho = sqrtf(rho_sq);
    const float two_theta = static_cast<float>(BackgroundAtan2(rho, lz));
    if (two_theta < b.tt_lo || two_theta >= b.tt_hi || rho == 0.0f)
        return -1;

    const float phi = static_cast<float>(BackgroundAtan2(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.
    int r_bin = static_cast<int>((two_theta - b.tt_lo) / b.d_tt);
    r_bin = r_bin < 0 ? 0 : (r_bin > BackgroundBand::RADIAL_BINS - 1 ? BackgroundBand::RADIAL_BINS - 1 : r_bin);
    int s_bin = static_cast<int>((phi + PI) / (2 * PI) * BackgroundBand::SECTORS);
    s_bin = s_bin < 0 ? 0 : (s_bin > BackgroundBand::SECTORS - 1 ? BackgroundBand::SECTORS - 1 : 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;
    jac_x = -b.pixel_size * (g_x * rot[0] + g_y * rot[3] + g_z * rot[6]);
    jac_y = -b.pixel_size * (g_x * rot[1] + g_y * rot[4] + g_z * rot[7]);
    return r_bin * BackgroundBand::SECTORS + s_bin;
}