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

#pragma once

// =============================================================================
// SpotWidth - how wide the recorded spots actually are, in detector pixels
// =============================================================================
//
// The integrator's signal radius r1 is not the integration domain (the profile fit runs over a
// (2*ceil(r2)+1)^2 grid), but it IS the aperture the profile WIDTH is learned over, and a second
// moment taken over a disk of radius a saturates at a^2/4 - so at the shipped r1 = 4 the learned
// Gaussian sigma can never exceed 2 px whatever the spot is. A crystal whose spots are wider than
// that is integrated with a profile the model cannot represent, and it is exactly those crystals
// that gain from a larger r1. Measured over the rotation battery, one threshold on the recorded
// width separates the two populations cleanly.
//
// This measures that width WITHOUT asking the integrator, so the answer can be used to set r1:
// r80, the radius holding 80 % of a spot's flux, is read off the encircled-flux curve of isolated
// strong spots over a FIXED 14 px aperture normalised at 8 px. Nothing here depends on r1, r2 or
// r3, or on any quantity the integrator computes; the aperture is a constant of this file. The
// curve is built about the spot's own flux-weighted centroid, no shape is assumed, and the
// candidates are stratified by resolution - the strongest spots of a weak crystal sit at high
// angle and of a strong one at low angle, and the mosaic footprint grows as 1/d, so an unstratified
// sample compares crystals at different resolutions.
// =============================================================================

#include <array>
#include <functional>
#include <optional>
#include <vector>

#include "../../common/DetectorModules.h"
#include "../../common/DiffractionGeometry.h"
#include "../../common/DiffractionSpot.h"
#include "../image_preprocessing/ImagePreprocessorBuffer.h"
#include "SpotFootprint.h"

namespace spot_width {
// The encircled-flux curve runs to R_MAX px and is normalised at R_NORM, clear of the noisy tail.
// R_NORM = 8 also fixes what this estimator can see: a spot needs r80 <= 8 px to be measured at
// all, which is five times the width the decision below is taken at.
constexpr int R_MAX = 14;
constexpr int R_NORM = 8;
// The background under a spot is the median of the R_MAX+2 .. R_MAX+6 px ring around it.
constexpr int R_BKG_IN = R_MAX + 2;
constexpr int R_BKG_OUT = R_MAX + 6;
// Resolution the widths are compared at. Inside the measured range of essentially every crystal.
constexpr float D_REF_A = 5.0f;

// One spot's second moments along and across its own radius, and the exact geometry that turns
// them into a bandwidth (EstimateBandwidth). jr and jt are the pixels one radian of 2theta, and of
// the angle across the scattering plane, moves the spot by; cos_psi is the angle of the ray to the
// sensor normal, and par_u/par_v are tan^2(psi) times the square of the in-plane ray direction's
// component along the radius and across it - the conversion depth's lateral smear, per um^2 of depth
// variance, in the two directions.
struct SpotShape {
    bool valid = false;
    float two_theta = 0.0f, jr = 0.0f, jt = 0.0f;
    float m_rad = 0.0f, m_tan = 0.0f;       // px^2, about the flux-weighted centroid
    float cos_psi = 1.0f, par_u = 0.0f, par_v = 0.0f;
};

// One spot's encircled flux C(1), C(2), ... C(R_MAX), divided by C(R_NORM), with the resolution it
// was recorded at, and its shape.
struct FluxCurve {
    float d_A = 0.0f;
    std::array<float, R_MAX> c{};
    SpotShape shape;
};

// The fitted width law: r80(d) = c0 + c1/d over resolution bands, weighted by the spots in each, so
// crystals whose strong spots sit at different resolutions are compared like with like. The fit is
// over 1/d because the mosaic contribution to the detector footprint grows that way, so c1 is the
// crystal's own answer to how much wider its spots get at high angle - a quantity the fit has always
// computed and nothing has ever been able to read, because only the value at one resolution was
// returned.
struct R80Fit {
    // Kept in double, and evaluated in double, because At() has to reproduce the single number the
    // caller used to get bit for bit: rounding the coefficients to float first and adding them after
    // moves the last bit, and r1 = round(2*r80) can turn on it.
    double c0 = 0.0, c1 = 0.0;
    // What the bands actually measured, widened the way this fit has always widened it. The law is
    // HELD at these outside the measured range rather than extrapolated - it has two coefficients and
    // as few as two points, so a slope read off a short lever arm must not be allowed to run away.
    double lo = 0.0, hi = 0.0;
    [[nodiscard]] float At(double d_A) const;
};

// The width law from a pool of curves, or nothing where too few spots were measurable.
[[nodiscard]] std::optional<R80Fit> FitR80(const std::vector<FluxCurve> &curves);

// r80 at D_REF_A, which is the law above read at one resolution.
[[nodiscard]] std::optional<float> R80AtReference(const std::vector<FluxCurve> &curves);

// The signal radius the measured width asks for: r1 = clamp(round(2*r80), 4, 6). The floor is the
// shipped default and the ceiling is set by pattern density - r2 = r1 + 2 also drives the
// neighbour-ownership radius and the inner edge of the background ring, and past 6 px a dense
// pattern starts losing reflections whose ring falls below six clean pixels.
[[nodiscard]] float R1ForWidth(float r80);

// When a growing sample has told us all it is going to about the radius, so the measurement can
// stop. R1ForWidth is a three-way choice and it moves only where 2*r80 crosses 4.5 and 5.5, i.e. at
// r80 = 2.25 and 2.75 px, so "settled" is two things and not one:
//
//   * the estimate has STOPPED MOVING - this sample is within SETTLED_STEP_PX of what the smaller
//     sample before it said; and
//   * it is not sitting ON a switch - it is SWITCH_CLEARANCE_PX clear of both of them.
//
// Neither carries the rotation battery on its own. Clearance without the step test loses the
// crystal whose r80 is 0.07 px from a switch and whose smaller samples read across it; the step test
// without the clearance lets a sample that has settled AT a switch stop there. Both bounds are set
// inside the range that is right on 38 of 38 crystals rather than at its edge: over every phase of
// the sub-sample the answer is exact for a clearance of 0.20-0.30 px and a step of 0.20-0.60 px.
//
// A quarter of a pixel is also half the width of the whole r1 = 5 band, so a sample landing inside
// that band never stops early - which is deliberate, because that band and its edges are where a
// reduced sample gets a crystal wrong.
constexpr float SWITCH_CLEARANCE_PX = 0.25f;
constexpr float SETTLED_STEP_PX = 0.40f;

// True where this sample's r80 has settled against the smaller sample's before it.
[[nodiscard]] bool WidthSettled(float r80, float r80_before);

// Most of the predicted reflections a widened radius may leave without a background ring, measured
// (BraggIntegrationCounts::bkg_starved_by_neighbour) rather than predicted. Above this the widening
// has cost the pattern more than it can be worth and the shipped radius is restored.
//
// The bound is read off the rotation battery, not chosen. At the radius the rule picks, the eleven
// crystals whose observation count survives the widening sit between 0.000 % and 0.315 %; the one
// crystal that loses a quarter of its observations sits at 4.082 %. That is the only gap in the
// distribution, and this is its midpoint in log space - the maximum-margin split, a factor 3.6
// clear of the nearest measurement on either side.
//
// It has to be the NEIGHBOUR count and not the total: the total also carries the detector's own
// floor (module gaps, the beam stop, the resolution mask), which reaches 2.3 % on some geometries,
// does not move with r1, and would put four of the rule's six wins within a factor of two of the
// loss. Splitting the two apart takes the separation from 2x to 13x.
constexpr double BKG_STARVED_MAX_FRACTION = 0.0113;

// -----------------------------------------------------------------------------------------------
// The X-ray bandwidth, read off the same spots.
//
// A spread dlambda/lambda of sigma smears a reflection along its own radius by 2 tan(theta) sigma
// radians of 2theta and not at all across it. Per spot, with u along the radius and v across it,
//
//     m_rad = 1/12 + par_u + jr^2 (s_ang^2 + 4 tan^2(theta) sigma^2)
//     m_tan = 1/12 + par_v + jt^2  s_ang^2
//
// (1/12: the pixel; par: the sensor parallax, fixed from the sensor's physics; s_ang: everything
// isotropic in angle - divergence, crystal size, mosaic spread seen edge-on). So
//
//     y = (m_rad - 1/12 - par_u) - (jr/jt)^2 (m_tan - 1/12 - par_v) = a + sigma^2 (2 jr tan theta)^2
//
// is a straight line whose slope is the bandwidth. It is fitted over eight equal-count bins of the
// abscissa, each a 20 %-trimmed mean, weighted by its own scatter, and the slope's error is the
// spread of 200 bootstrap re-draws of the spots, inflated by the fit's reduced chi^2 where the line
// fits worse than the scatter says. The slope is significant at BANDWIDTH_Z standard errors - over
// the whole battery that separates the multilayer beamlines from every monochromatic one.
//
// It is a LOWER bound: mosaic spread seen along the radius subtracts, and a spread of cell edges
// (strain) is exactly degenerate with it. What it measures is the effective radial broadening, and
// that is what the consumers of the bandwidth need. The run uses it where it is significant and
// neither the file nor --bandwidth states one.
//
// The width settles on fewer spots than the slope needs, so the pre-scan keeps measuring spots for the
// bandwidth alone until the pool holds BANDWIDTH_POOL_SPOTS or the sample runs out.
constexpr size_t BANDWIDTH_MIN_SPOTS = 80;
constexpr size_t BANDWIDTH_POOL_SPOTS = 1000;
constexpr double BANDWIDTH_Z = 3.0;

struct BandwidthEstimate {
    double fwhm = 0.0;        // relative FWHM of dlambda/lambda; negative where the slope is
    double fwhm_floor = 0.0;  // the FWHM one standard error of sigma^2 amounts to
    double z = 0.0;           // slope / its standard error
    double chi2 = 0.0;        // reduced chi^2 of the binned line
    size_t spots = 0;
};

// The bandwidth from the shapes of a pool of spots, or nothing where too few were measured.
// attenuation_um is the sensor's attenuation length at this wavelength, thickness_um its thickness
// (0 = no depth, no parallax), pixel_um the pixel size.
[[nodiscard]] std::optional<BandwidthEstimate> EstimateBandwidth(const std::vector<FluxCurve> &curves,
                                                                 double attenuation_um, double thickness_um,
                                                                 double pixel_um);
}   // namespace spot_width

// Every isolated, strong, fully readable spot of one image, as a normalised encircled-flux curve
// with its shape, appended to `out`. Called on one worker's own vector, so the pool is independent of the order the
// workers finish in. The spots are in the image; `modules` places its panels on the plane the geometry
// describes (empty: the image is the plane), and a spot's aperture has to lie on its own panel. The pixel
// values are `gain` to a photon (BraggIntegrationSettings::pixel_gain): "strong" and the background's clip
// are judged in photons.
void MeasureSpotFluxCurves(const ImagePreprocessorBuffer &image, int width, int height,
                           const DiffractionGeometry &geometry, const DetectorModules &modules,
                           const std::vector<DiffractionSpot> &spots,
                           std::vector<spot_width::FluxCurve> &out, float gain = 1.0f);

// The same, with spot i's pixels read about centre_px(i) - the pixel its position rounds to - with
// rows `stride` apart: R_BKG_OUT pixels of them every way that the image itself has. The frame is
// width x height.
void MeasureSpotFluxCurves(const std::function<const int32_t *(size_t)> &centre_px, int stride, int width, int height,
                           const DiffractionGeometry &geometry, const DetectorModules &modules,
                           const std::vector<DiffractionSpot> &spots,
                           std::vector<spot_width::FluxCurve> &out, float gain = 1.0f);

// The squares MeasureSpotsOnTiles reads about each spot. The flux curves stay within
// spot_width::R_BKG_OUT of their spot; the footprint's window re-centres without a bound, so a spot
// whose window walks off its square is measured again from the start on a square of
// FOOTPRINT_TILE_HALF, and on the whole frame where even that is left. Over the rotation battery's
// pre-scans a window leaves the first square on a few spots of most frames, and the second on none.
constexpr int SPOT_TILE_HALF = 24;
constexpr int FOOTPRINT_TILE_HALF = 80;

// Spots whose footprint had to be measured again on a larger square, and on the whole frame.
struct SpotTileMisses {
    size_t larger = 0;
    size_t whole_frame = 0;
};

// MeasureSpotFluxCurves and MeasureFootprintSpots of one frame, reading the (2*half+1)^2 pixels about
// each spot (ImagePreprocessorBuffer::GatherTiles) rather than the frame - on the GPU only those cross
// back to the host. The answers are the whole frame's to the bit: each measurement reads the same
// pixels in the same order. The whole frame is read where the squares would be more pixels than it.
SpotTileMisses MeasureSpotsOnTiles(ImagePreprocessorBuffer &image, int width, int height,
                                   const DiffractionGeometry &geometry, const DetectorModules &modules,
                                   const std::vector<DiffractionSpot> &spots,
                                   std::vector<spot_width::FluxCurve> &curves, std::vector<FootprintSpot> &footprint,
                                   float gain = 1.0f, int half = SPOT_TILE_HALF, int footprint_half = FOOTPRINT_TILE_HALF);