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

#include "SpotWidth.h"

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <random>
#include <utility>

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

using namespace spot_width;

namespace {

// The engine reads pixels in the INT32_MIN(masked)/INT32_MAX(saturated) convention.
inline bool valid(int32_t v) { return v != INT32_MIN && v != INT32_MAX; }

// Nothing inside this radius of the beam centre: the beam stop and its halo are not spots.
constexpr float MIN_BEAM_DISTANCE_PX = 60.0f;
// A neighbour this close puts its own flux inside the aperture, which would read as extra width.
constexpr float ISOLATION_PX = 28.0f;
// Spots taken per resolution band per image, strongest first.
constexpr int PER_BAND_PER_IMAGE = 40;
// The r <= 4 px sum must be this many sigma above the background before the tail is believed.
constexpr double SNR_MIN = 15.0;
constexpr int R_CENTROID = 4;
constexpr double MAX_CENTROID_OFFSET_PX = 2.0;
// Spots needed before a band, and the crystal, are characterised at all.
constexpr size_t MIN_SPOTS_PER_BAND = 15;
constexpr size_t MIN_SPOTS_TOTAL = 20;

// Resolution bands, A. The quota is per band, so a crystal is characterised over its whole range
// and not wherever its strongest spots happen to sit.
constexpr int N_BAND = 5;
constexpr std::array<std::pair<float, float>, N_BAND> BANDS = {{
        {2.0f, 3.0f}, {3.0f, 4.5f}, {4.5f, 7.0f}, {7.0f, 12.0f}, {12.0f, 30.0f}}};

// Every radius here is compared against an integer pixel offset, so all of it is exact integer
// arithmetic and no square root is needed anywhere in the pixel loops.
constexpr int isqrt_floor(int n) {
    int r = 0;
    while ((r + 1) * (r + 1) <= n) r++;
    return r;
}

// Half-width of the disk of radius R on row dy: the largest |dx| with dx^2 + dy^2 <= R^2. Walking
// the rows by their own extent visits the disk itself rather than its bounding box.
template <int R>
constexpr std::array<int, R + 1> disk_row_half() {
    std::array<int, R + 1> a{};
    for (int dy = 0; dy <= R; dy++) a[dy] = isqrt_floor(R * R - dy * dy);
    return a;
}
constexpr auto HALF_BKG = disk_row_half<R_BKG_OUT>();
constexpr auto HALF_CORE = disk_row_half<R_CENTROID>();

// The largest |dx| on row dy that is still INSIDE the background ring's inner edge, so |dx| beyond
// it is in the ring; -1 where the whole row is.
constexpr std::array<int, R_BKG_OUT + 1> ring_row_inner() {
    std::array<int, R_BKG_OUT + 1> a{};
    for (int dy = 0; dy <= R_BKG_OUT; dy++) {
        const int rem = R_BKG_IN * R_BKG_IN - dy * dy - 1;
        a[dy] = rem < 0 ? -1 : isqrt_floor(rem);
    }
    return a;
}
constexpr auto INNER_BKG = ring_row_inner();

// floor(sqrt(n)) for every squared distance the encircled-flux aperture can produce, so a pixel's
// radial bin - the smallest integer radius that contains it - is a table lookup and a compare.
constexpr std::array<int, R_MAX * R_MAX + 1> isqrt_lookup() {
    std::array<int, R_MAX * R_MAX + 1> a{};
    for (int n = 0; n <= R_MAX * R_MAX; n++) a[n] = isqrt_floor(n);
    return a;
}
constexpr auto ISQRT = isqrt_lookup();

int band_of(float d_A) {
    for (int b = 0; b < N_BAND; b++)
        if (d_A >= BANDS[b].first && d_A < BANDS[b].second) return b;
    return -1;
}

// The radius at which the curve reaches `frac`, linearly interpolated. prof[i] is the flux inside
// radius i+1.
float interpolate_radius(double frac, const std::array<float, R_MAX> &prof) {
    if (prof[0] >= frac)
        return prof[0] > 0.0f ? static_cast<float>(frac / prof[0]) : 1.0f;
    for (int i = 1; i < R_MAX; i++)
        if (prof[i] >= frac)
            return static_cast<float>(i + (frac - prof[i - 1]) / (prof[i] - prof[i - 1]));
    return static_cast<float>(R_MAX);
}

// The disk the second moments are taken over, px.
constexpr int R_MOMENT = 9;

// The spot's second moments along and across its radius, about its own flux-weighted centroid.
// (cx, cy) is the pixel the spot is centred on, (mx, my) the core centroid relative to it, `ring`
// the background ring's counts. The background is the ring's mean with anything 5 sigma above it
// clipped: a median is biased low on sparse Poisson counts, and a bias here adds the same flux at
// every radius of the disk - an extra second moment that does not cancel between the two directions
// once they are weighted by the obliquity.
SpotShape shape_of(const int32_t *centre_px, int width, int cx, int cy, double mx, double my,
                   const std::vector<int32_t> &ring, const DiffractionGeometry &geometry) {
    SpotShape shape;
    double sum = 0.0;
    for (const int32_t v : ring) sum += v;
    const double clip = sum / static_cast<double>(ring.size())
                        + 5.0 * std::sqrt(std::abs(sum / static_cast<double>(ring.size())) + 0.5) + 2.0;
    sum = 0.0;
    size_t n = 0;
    for (const int32_t v : ring)
        if (v < clip) { sum += v; n++; }
    if (n == 0) return shape;
    const double bkg = sum / static_cast<double>(n);

    // Centroid over the moment disk, three times, then the moments about it.
    constexpr double R2 = static_cast<double>(R_MOMENT) * R_MOMENT;
    double flux = 0.0, sxx = 0.0, syy = 0.0, sxy = 0.0;
    for (int iter = 0; iter < 4; iter++) {
        double f = 0.0, sx = 0.0, sy = 0.0, qxx = 0.0, qyy = 0.0, qxy = 0.0;
        for (int dy = -R_MOMENT - 2; dy <= R_MOMENT + 2; dy++) {
            const int32_t *row = centre_px + static_cast<ptrdiff_t>(dy) * width;
            for (int dx = -R_MOMENT - 2; dx <= R_MOMENT + 2; dx++) {
                const double ddx = dx - mx, ddy = dy - my;
                if (ddx * ddx + ddy * ddy >= R2) continue;
                const double w = row[dx] - bkg;
                f += w;
                sx += w * ddx;
                sy += w * ddy;
                qxx += w * ddx * ddx;
                qyy += w * ddy * ddy;
                qxy += w * ddx * ddy;
            }
        }
        if (!(f > 0.0)) return shape;
        if (iter < 3) {
            mx += sx / f;
            my += sy / f;
            if (std::abs(mx) > 2.0 || std::abs(my) > 2.0) return shape;
        } else {
            flux = f;
            // About the centroid of this last pass, not the disk's centre.
            const double ox = sx / f, oy = sy / f;
            sxx = qxx / f - ox * ox;
            syy = qyy / f - oy * oy;
            sxy = qxy / f - ox * oy;
        }
    }
    if (!(flux > 0.0)) return shape;

    // Where one radian of 2theta, and of the angle across the scattering plane, moves the spot.
    const float x = static_cast<float>(cx + mx), y = static_cast<float>(cy + my);
    const Coord s0(0, 0, 1);
    const Coord s1 = geometry.LabCoord(x, y).Normalize();
    const float two_theta = std::acos(std::clamp(s1 * s0, -1.0f, 1.0f));
    if (!(two_theta > 1.0f * static_cast<float>(PI) / 180.0f)) return shape;
    const Coord across = (s0 % s1).Normalize();
    const Coord along = across % s1;   // in the scattering plane, away from the beam
    const float lambda = geometry.GetWavelength_A();
    constexpr float EPS = 1e-3f;
    const auto hit = [&](const Coord &dir) {
        return geometry.RecipToDetector(dir / lambda - s0 / lambda);
    };
    const auto p0 = hit(s1);
    const auto pr = hit(s1 * std::cos(EPS) + along * std::sin(EPS));
    const auto pt = hit(s1 * std::cos(EPS) + across * std::sin(EPS));
    const double jrx = (pr.first - p0.first) / EPS, jry = (pr.second - p0.second) / EPS;
    const double jr = std::hypot(jrx, jry);
    if (!(jr > 0.0)) return shape;
    const double ux = jrx / jr, uy = jry / jr, vx = -uy, vy = ux;
    const double jt = std::abs((pt.first - p0.first) / EPS * vx + (pt.second - p0.second) / EPS * vy);
    if (!(jt > 0.0)) return shape;

    // Parallax: the ray's in-plane direction on the sensor, and tan^2 of its angle to the normal.
    const Coord normal = geometry.GetNormalAxis();
    const float cos_psi = std::abs(s1 * normal);
    const Coord in_plane = s1 - normal * (s1 * normal);
    const double e_len = in_plane.Length();
    double par_u = 0.0, par_v = 0.0;
    if (e_len > 1e-6 && cos_psi > 1e-3f) {
        const double ex = in_plane * geometry.GetFastAxis() / e_len;
        const double ey = in_plane * geometry.GetSlowAxis() / e_len;
        const double tan2 = e_len * e_len / (static_cast<double>(cos_psi) * cos_psi);
        par_u = tan2 * (ex * ux + ey * uy) * (ex * ux + ey * uy);
        par_v = tan2 * (ex * vx + ey * vy) * (ex * vx + ey * vy);
    }

    shape.valid = true;
    shape.two_theta = two_theta;
    shape.jr = static_cast<float>(jr);
    shape.jt = static_cast<float>(jt);
    shape.m_rad = static_cast<float>(ux * ux * sxx + 2.0 * ux * uy * sxy + uy * uy * syy);
    shape.m_tan = static_cast<float>(vx * vx * sxx + 2.0 * vx * vy * sxy + vy * vy * syy);
    shape.cos_psi = cos_psi;
    shape.par_u = static_cast<float>(par_u);
    shape.par_v = static_cast<float>(par_v);
    return shape;
}

template <typename T>
double median_of(std::vector<T> &v) {
    if (v.empty()) return 0.0;
    const size_t mid = v.size() / 2;
    std::nth_element(v.begin(), v.begin() + mid, v.end());
    const double hi = v[mid];
    if (v.size() % 2 == 1) return hi;
    return 0.5 * (hi + *std::max_element(v.begin(), v.begin() + mid));
}

}   // namespace

void MeasureSpotFluxCurves(const ImagePreprocessorBuffer &image, int width, int height,
                           const DiffractionGeometry &geometry,
                           const std::vector<DiffractionSpot> &spots,
                           std::vector<FluxCurve> &out) {
    if (spots.empty()) return;

    const int32_t *pixels = image.data();
    if (pixels == nullptr) return;
    const float beam_x = geometry.GetBeamX_pxl(), beam_y = geometry.GetBeamY_pxl();

    // Where every spot of this image sits, so isolation can be tested against all of them and not
    // only against the ones that survive the gates below.
    std::vector<Coord> centre(spots.size());
    for (size_t i = 0; i < spots.size(); i++)
        centre[i] = spots[i].RawCoord();

    // Isolation on a grid of ISOLATION_PX cells: a neighbour within that distance is in this cell or
    // one of the eight around it. The grid is held as a counting sort - one index array and one
    // offset array - rather than a vector per cell, which on a crowded detector is tens of thousands
    // of allocations per image for a structure that is read once.
    const int gw = static_cast<int>(width / ISOLATION_PX) + 1;
    const int gh = static_cast<int>(height / ISOLATION_PX) + 1;
    const size_t ncell = static_cast<size_t>(gw) * gh;
    const auto cell_of = [&](const Coord &c) {
        const int gx = std::clamp(static_cast<int>(c.x / ISOLATION_PX), 0, gw - 1);
        const int gy = std::clamp(static_cast<int>(c.y / ISOLATION_PX), 0, gh - 1);
        return static_cast<size_t>(gy) * gw + gx;
    };
    std::vector<uint32_t> cell_begin(ncell + 1, 0), cell_item(spots.size()), spot_cell(spots.size());
    for (size_t i = 0; i < spots.size(); i++) {
        spot_cell[i] = static_cast<uint32_t>(cell_of(centre[i]));
        cell_begin[spot_cell[i] + 1]++;
    }
    for (size_t c = 0; c < ncell; c++) cell_begin[c + 1] += cell_begin[c];
    {
        std::vector<uint32_t> cursor(cell_begin.begin(), cell_begin.end() - 1);
        for (size_t i = 0; i < spots.size(); i++)
            cell_item[cursor[spot_cell[i]]++] = static_cast<uint32_t>(i);
    }
    constexpr double ISOLATION_PX2 = static_cast<double>(ISOLATION_PX) * ISOLATION_PX;
    const auto isolated = [&](size_t i) {
        const int gx = static_cast<int>(spot_cell[i] % gw), gy = static_cast<int>(spot_cell[i] / gw);
        for (int y = std::max(0, gy - 1); y <= std::min(gh - 1, gy + 1); y++)
            for (int x = std::max(0, gx - 1); x <= std::min(gw - 1, gx + 1); x++) {
                const size_t c = static_cast<size_t>(y) * gw + x;
                for (uint32_t k = cell_begin[c]; k < cell_begin[c + 1]; k++) {
                    const uint32_t j = cell_item[k];
                    if (j == i) continue;
                    const double ddx = centre[j].x - centre[i].x, ddy = centre[j].y - centre[i].y;
                    if (ddx * ddx + ddy * ddy < ISOLATION_PX2) return false;
                }
            }
        return true;
    };

    // Candidates that pass the geometric gates, by band, strongest first.
    struct Candidate { size_t index; int64_t count; float d_A; };
    std::array<std::vector<Candidate>, N_BAND> candidates;
    constexpr double MIN_BEAM_DISTANCE_PX2 = static_cast<double>(MIN_BEAM_DISTANCE_PX)
                                             * MIN_BEAM_DISTANCE_PX;
    for (size_t i = 0; i < spots.size(); i++) {
        const Coord &c = centre[i];
        const int cx = static_cast<int>(std::lround(c.x)), cy = static_cast<int>(std::lround(c.y));
        if (cx < R_BKG_OUT || cy < R_BKG_OUT || cx >= width - R_BKG_OUT || cy >= height - R_BKG_OUT)
            continue;
        const double bx = c.x - beam_x, by = c.y - beam_y;
        if (bx * bx + by * by < MIN_BEAM_DISTANCE_PX2) continue;
        const float d_A = geometry.PxlToRes(c.x, c.y);
        const int band = band_of(d_A);
        if (band < 0) continue;
        if (!isolated(i)) continue;
        candidates[band].push_back({i, spots[i].Count(), d_A});
    }

    // The ring is gathered as the counts it is - the median of an int list is the same number, and
    // half the bytes move through the partial sort.
    std::vector<int32_t> ring;
    ring.reserve(4 * (R_BKG_OUT + 1) * (R_BKG_OUT - R_BKG_IN + 1));
    for (int band = 0; band < N_BAND; band++) {
        auto &cand = candidates[band];
        const size_t take = std::min<size_t>(cand.size(), PER_BAND_PER_IMAGE);
        std::partial_sort(cand.begin(), cand.begin() + take, cand.end(),
                          [](const Candidate &a, const Candidate &b) { return a.count > b.count; });
        for (size_t k = 0; k < take; k++) {
            const Coord &c = centre[cand[k].index];
            const int cx = static_cast<int>(std::lround(c.x)), cy = static_cast<int>(std::lround(c.y));
            const int32_t *centre_px = pixels + static_cast<size_t>(cy) * width + cx;

            // The background under the spot, and a check that the whole aperture is readable: a hole
            // in it removes flux from one radius and not another, which is exactly the shape this
            // measures.
            ring.clear();
            bool readable = true;
            for (int dy = -R_BKG_OUT; dy <= R_BKG_OUT && readable; dy++) {
                const int half = HALF_BKG[std::abs(dy)], inner = INNER_BKG[std::abs(dy)];
                const int32_t *row = centre_px + static_cast<ptrdiff_t>(dy) * width;
                for (int dx = -half; dx <= half; dx++) {
                    const int32_t px = row[dx];
                    if (!valid(px)) { readable = false; break; }
                    if (dx > inner || dx < -inner) ring.push_back(px);
                }
            }
            if (!readable || ring.size() < 20) continue;
            const size_t n_ring = ring.size();
            const double bkg = median_of(ring);

            // Flux and centroid over the r <= 4 px core, then the signal-to-noise gate. A weak spot's
            // tail is background, and an encircled-flux curve built on it measures the background.
            double core = 0.0, mx = 0.0, my = 0.0;
            int n_core = 0;
            for (int dy = -R_CENTROID; dy <= R_CENTROID; dy++) {
                const int half = HALF_CORE[std::abs(dy)];
                const int32_t *row = centre_px + static_cast<ptrdiff_t>(dy) * width;
                for (int dx = -half; dx <= half; dx++) {
                    const double v = row[dx] - bkg;
                    core += v;
                    mx += v * dx;
                    my += v * dy;
                    ++n_core;
                }
            }
            if (core <= 0.0) continue;
            const double noise = std::sqrt(core + n_core * std::max(bkg, 0.05)
                                                  * (1.0 + static_cast<double>(n_core) / n_ring));
            if (core / noise < SNR_MIN) continue;
            mx /= core;
            my /= core;
            if (std::abs(mx) > MAX_CENTROID_OFFSET_PX || std::abs(my) > MAX_CENTROID_OFFSET_PX)
                continue;

            // The encircled flux about that centroid, out to the fixed aperture. Each pixel is added
            // to the one bin its own radius falls in and the curve is the running total over the
            // bins: the encircled flux at t is everything inside t, so adding every pixel into every
            // bin beyond it instead would sum the same aperture R_MAX/2 times over.
            std::array<double, R_MAX + 1> bin{};
            constexpr double R2_MAX = static_cast<double>(R_MAX) * R_MAX;
            for (int dy = -R_MAX; dy <= R_MAX; dy++) {
                const double ddy = dy - my, ddy2 = ddy * ddy;
                if (ddy2 > R2_MAX) continue;
                const double span = std::sqrt(R2_MAX - ddy2);
                const int lo = std::max(-R_MAX, static_cast<int>(std::floor(mx - span)));
                const int hi = std::min(R_MAX, static_cast<int>(std::ceil(mx + span)));
                const int32_t *row = centre_px + static_cast<ptrdiff_t>(dy) * width;
                for (int dx = lo; dx <= hi; dx++) {
                    const double ddx = dx - mx, rc2 = ddx * ddx + ddy2;
                    if (rc2 > R2_MAX) continue;
                    const int s = ISQRT[static_cast<int>(rc2)];
                    const int t = std::max(1, static_cast<double>(s) * s == rc2 ? s : s + 1);
                    bin[t] += row[dx] - bkg;
                }
            }
            FluxCurve curve;
            curve.d_A = cand[k].d_A;
            double encircled = 0.0;
            for (int t = 1; t <= R_MAX; t++) {
                encircled += bin[t];
                curve.c[t - 1] = static_cast<float>(encircled);
            }
            if (!(curve.c[R_NORM - 1] > 0.0f) || !(curve.c[R_MAX - 1] > 0.0f)) continue;
            const float norm = curve.c[R_NORM - 1];
            for (float &v : curve.c) v /= norm;
            curve.shape = shape_of(centre_px, width, cx, cy, mx, my, ring, geometry);
            out.push_back(curve);
        }
    }
}

float spot_width::R80Fit::At(double d_A) const {
    return static_cast<float>(std::clamp(c0 + c1 / d_A, lo, hi));
}

std::optional<spot_width::R80Fit> spot_width::FitR80(const std::vector<FluxCurve> &curves) {
    if (curves.size() < MIN_SPOTS_TOTAL) return std::nullopt;

    // One point per band: the median curve of the band, the radius it holds 80 % of its flux at, and
    // the median resolution it was measured at.
    struct Point { double inv_d; double r80; double weight; };
    std::vector<Point> points;
    std::vector<double> values, band_d;
    std::vector<uint32_t> members;
    for (int b = 0; b < N_BAND; b++) {
        members.clear();
        band_d.clear();
        for (uint32_t i = 0; i < curves.size(); i++)
            if (curves[i].d_A >= BANDS[b].first && curves[i].d_A < BANDS[b].second) {
                members.push_back(i);
                band_d.push_back(curves[i].d_A);
            }
        if (band_d.size() < MIN_SPOTS_PER_BAND) continue;
        std::array<float, R_MAX> profile{};
        for (int t = 0; t < R_MAX; t++) {
            values.clear();
            for (uint32_t i : members) values.push_back(curves[i].c[t]);
            profile[t] = static_cast<float>(median_of(values));
        }
        const double d_med = median_of(band_d);
        if (d_med <= 0.0) continue;
        points.push_back({1.0 / d_med, interpolate_radius(0.8, profile),
                          static_cast<double>(band_d.size())});
    }
    if (points.empty()) return std::nullopt;

    // Never extrapolate outside what the bands actually measured. A single band measures no slope, so
    // its own value is the whole law; the bounds then bracket it and At() returns it unchanged.
    R80Fit fit;
    double lo = std::numeric_limits<double>::max(), hi = 0.0;
    for (const auto &p : points) { lo = std::min(lo, p.r80); hi = std::max(hi, p.r80); }
    fit.lo = 0.8 * lo;
    fit.hi = 1.25 * hi;
    if (points.size() == 1) {
        fit.c0 = points[0].r80;
        return fit;
    }

    // The mosaic contribution to the detector footprint grows as 1/d, so r80 is linear in 1/d.
    double sw = 0.0, sx = 0.0, sxx = 0.0, sy = 0.0, sxy = 0.0;
    for (const auto &p : points) {
        sw += p.weight;
        sx += p.weight * p.inv_d;
        sxx += p.weight * p.inv_d * p.inv_d;
        sy += p.weight * p.r80;
        sxy += p.weight * p.inv_d * p.r80;
    }
    const double det = sw * sxx - sx * sx;
    fit.c0 = sy / sw;
    if (std::abs(det) > 1e-12) {
        fit.c1 = (sw * sxy - sx * sy) / det;
        fit.c0 = (sy - fit.c1 * sx) / sw;
    }
    return fit;
}

std::optional<float> spot_width::R80AtReference(const std::vector<FluxCurve> &curves) {
    const auto fit = FitR80(curves);
    if (!fit) return std::nullopt;
    return fit->At(D_REF_A);
}

float spot_width::R1ForWidth(float r80) {
    return std::clamp(std::round(2.0f * r80), 4.0f, 6.0f);
}

bool spot_width::WidthSettled(float r80, float r80_before) {
    return std::abs(r80 - r80_before) < SETTLED_STEP_PX
           && std::abs(r80 - 2.25f) > SWITCH_CLEARANCE_PX
           && std::abs(r80 - 2.75f) > SWITCH_CLEARANCE_PX;
}

namespace {

struct BandwidthPoint { double x, y; };

// numpy's default (linear) quantile of sorted values.
double quantile_sorted(const std::vector<double> &v, double p) {
    const double pos = p * static_cast<double>(v.size() - 1);
    const size_t lo = static_cast<size_t>(std::floor(pos));
    const size_t hi = std::min(lo + 1, v.size() - 1);
    return v[lo] + (pos - static_cast<double>(lo)) * (v[hi] - v[lo]);
}

// The line y = a + s2 x through eight equal-count bins of x, each a 20 %-trimmed mean weighted by its
// own scatter; returns {s2, reduced chi^2}. Nothing where fewer than two bins could be formed.
std::optional<std::pair<double, double>> fit_bandwidth_line(std::vector<BandwidthPoint> points) {
    constexpr size_t N_BIN = 8;
    constexpr size_t MIN_PER_BIN = 10;
    constexpr double TRIM = 0.2;
    std::stable_sort(points.begin(), points.end(),
                     [](const BandwidthPoint &a, const BandwidthPoint &b) { return a.x < b.x; });
    std::vector<double> bx, by, bw, ys;
    size_t begin = 0;
    for (size_t b = 0; b < N_BIN; b++) {
        const size_t len = points.size() / N_BIN + (b < points.size() % N_BIN ? 1 : 0);
        const size_t end = begin + len;
        if (len >= MIN_PER_BIN) {
            ys.clear();
            for (size_t i = begin; i < end; i++) ys.push_back(points[i].y);
            std::sort(ys.begin(), ys.end());
            const double lo = quantile_sorted(ys, TRIM), hi = quantile_sorted(ys, 1.0 - TRIM);
            double n = 0.0, sx = 0.0, sy = 0.0, syy = 0.0;
            for (size_t i = begin; i < end; i++)
                if (points[i].y >= lo && points[i].y <= hi) {
                    n += 1.0;
                    sx += points[i].x;
                    sy += points[i].y;
                    syy += points[i].y * points[i].y;
                }
            const double mean = sy / n, var = syy / n - mean * mean;
            bx.push_back(sx / n);
            by.push_back(mean);
            bw.push_back(n / std::max(var, 1e-6));
        }
        begin = end;
    }
    if (bx.size() < 2) return std::nullopt;

    double sw = 0.0, sx = 0.0, sy = 0.0, sxx = 0.0, sxy = 0.0;
    for (size_t i = 0; i < bx.size(); i++) {
        sw += bw[i];
        sx += bw[i] * bx[i];
        sy += bw[i] * by[i];
        sxx += bw[i] * bx[i] * bx[i];
        sxy += bw[i] * bx[i] * by[i];
    }
    const double det = sw * sxx - sx * sx;
    if (!(std::abs(det) > 0.0)) return std::nullopt;
    const double slope = (sw * sxy - sx * sy) / det;
    const double intercept = (sy - slope * sx) / sw;
    double chi2 = 0.0;
    for (size_t i = 0; i < bx.size(); i++) {
        const double r = by[i] - intercept - slope * bx[i];
        chi2 += bw[i] * r * r;
    }
    chi2 /= std::max<double>(1.0, static_cast<double>(bx.size()) - 2.0);
    return std::make_pair(slope, chi2);
}

}   // namespace

std::optional<spot_width::BandwidthEstimate> spot_width::EstimateBandwidth(const std::vector<FluxCurve> &curves,
                                                                           double attenuation_um,
                                                                           double thickness_um,
                                                                           double pixel_um) {
    constexpr double PIXEL_VAR = 1.0 / 12.0;
    constexpr double FWHM_PER_SIGMA = 2.3548;
    constexpr int N_BOOTSTRAP = 200;

    std::vector<BandwidthPoint> points;
    for (const auto &c : curves) {
        const SpotShape &s = c.shape;
        if (!s.valid) continue;
        // The conversion depth's variance for this ray: its depth length is L cos(psi).
        const double depth_var_px2 = sensor_absorption::ConversionDepthVariance_um2(
                attenuation_um * s.cos_psi, thickness_um) / (pixel_um * pixel_um);
        const double obliquity = (static_cast<double>(s.jr) / s.jt) * (static_cast<double>(s.jr) / s.jt);
        const double x = 2.0 * s.jr * std::tan(0.5 * s.two_theta);
        points.push_back({x * x, (s.m_rad - PIXEL_VAR - depth_var_px2 * s.par_u)
                                 - obliquity * (s.m_tan - PIXEL_VAR - depth_var_px2 * s.par_v)});
    }
    if (points.size() < BANDWIDTH_MIN_SPOTS) return std::nullopt;

    const auto fit = fit_bandwidth_line(points);
    if (!fit) return std::nullopt;

    // The slope's spread over re-draws of the spots. The generator is seeded, so the answer is the
    // same on every run of the same data.
    std::mt19937 rng(1);
    std::vector<BandwidthPoint> redraw(points.size());
    double sum = 0.0, sum2 = 0.0;
    int n = 0;
    for (int b = 0; b < N_BOOTSTRAP; b++) {
        for (auto &p : redraw) p = points[rng() % points.size()];
        if (const auto f = fit_bandwidth_line(redraw)) {
            sum += f->first;
            sum2 += f->first * f->first;
            n++;
        }
    }
    const double mean = n > 0 ? sum / n : 0.0;
    const double sd = n > 1 ? std::sqrt(std::max(0.0, sum2 / n - mean * mean)) : 0.0;
    const double se = sd * std::sqrt(std::max(1.0, fit->second));

    BandwidthEstimate e;
    e.spots = points.size();
    e.chi2 = fit->second;
    e.fwhm = (fit->first < 0.0 ? -1.0 : 1.0) * FWHM_PER_SIGMA * std::sqrt(std::abs(fit->first));
    e.fwhm_floor = FWHM_PER_SIGMA * std::sqrt(se);
    e.z = se > 0.0 ? fit->first / se : 0.0;
    return e;
}