Jungfraujoch/common/DetectorModules.h

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

#pragma once

#include <cstddef>
#include <cstdint>
#include <optional>
#include <utility>
#include <vector>

#include "ModuleShift.h"

class DetectorGeometry;

// One detector module: where its pixels are in the image rugnux reads, and where they really are on the
// detector plane. The plane is the one DiffractionGeometry describes (beam centre, distance, tilt), in its
// pixel units, so a module placed exactly as the image has it maps every pixel onto itself.
//
// Image pixel (x, y) of the module, x0 <= x < x0 + width, y0 <= y < y0 + height, lies on the plane at
//   (ox, oy) + (x - x0) (fx, fy) + (y - y0) (sx, sy)
// - the NXdetector_module description (data_origin, data_size, module_offset, fast_pixel_direction,
// slow_pixel_direction) taken in the plane. A pixel's centre is at its index, as everywhere in the image.
struct DetectorModule {
    int64_t x0 = 0, y0 = 0, width = 0, height = 0;
    double ox = 0, oy = 0;
    double fx = 1, fy = 0;
    double sx = 0, sy = 1;

    // The module where the image has it: every pixel maps onto itself.
    static DetectorModule AsImaged(int64_t x0, int64_t y0, int64_t width, int64_t height);

    [[nodiscard]] bool ContainsImagePixel(double x, double y) const;
    [[nodiscard]] std::pair<double, double> ImageToPlane(double x, double y) const;
    [[nodiscard]] std::pair<double, double> PlaneToImage(double px, double py) const;
    // Plane position of the module's centre.
    [[nodiscard]] std::pair<double, double> PlaneCentre() const;
    // Moved by a refinement's rigid step, which is expressed on the plane (ModuleShift: turn about
    // (x0, y0) of the shift, then shift).
    void Apply(const ModuleShift &shift);
};

// A point in the image rugnux reads, on panel `panel` (an index into DetectorModules; -1 when the table is
// empty and the whole image is one flat panel).
struct PanelPoint {
    int64_t panel;
    double x, y;
};

// An axis-aligned box on the detector plane, x0 <= x < x1, y0 <= y < y1.
struct PlaneBox {
    double x0, y0, x1, y1;
};

// The modules of one detector. Empty means the image is one flat detector: every mapping is the identity.
// A module is a panel: one rigid box of the image with one placement. Panel ids are indices into the table.
class DetectorModules {
    std::vector<DetectorModule> modules;
public:
    DetectorModules() = default;
    explicit DetectorModules(std::vector<DetectorModule> modules);

    [[nodiscard]] bool empty() const;
    [[nodiscard]] size_t size() const;
    [[nodiscard]] const DetectorModule &operator[](size_t i) const;
    [[nodiscard]] const std::vector<DetectorModule> &Get() const;

    // True where some module is placed other than as imaged.
    [[nodiscard]] bool Moved() const;

    // Module of an image pixel, or -1 (a gap, or outside every module).
    [[nodiscard]] int64_t ModuleOfImagePixel(double x, double y) const;
    // Module whose placed footprint holds a plane point, or -1.
    [[nodiscard]] int64_t ModuleOfPlanePoint(double px, double py) const;

    // Image -> plane; a point on no module is left as it is.
    [[nodiscard]] std::pair<double, double> ImageToPlane(double x, double y) const;
    // Plane -> image; a point on no placed module is left as it is.
    [[nodiscard]] std::pair<double, double> PlaneToImage(double px, double py) const;

    // Image <-> plane on a known panel, no search. Panel -1 (the empty table) is the identity.
    [[nodiscard]] std::pair<double, double> ImageToPlane(int64_t panel, double x, double y) const;
    [[nodiscard]] std::pair<double, double> PlaneToImage(int64_t panel, double px, double py) const;
    // The panel whose placed footprint holds a plane point, and where that point is in the image; nullopt in
    // a gap or off every panel. Empty table: panel -1 and the point itself, wherever it is - the caller
    // bounds it by the image, as without panels.
    [[nodiscard]] std::optional<PanelPoint> PlaneToPanel(double px, double py) const;
    // Box on the plane enclosing every panel's pixels (their footprints, half a pixel beyond the edge pixel
    // centres): a quick reject before PlaneToPanel. Empty table: [0, width) x [0, height), the image
    // rectangle the bounds tests without panels use.
    [[nodiscard]] PlaneBox PlaneBounds(int64_t width, int64_t height) const;

    // Calls f(panel, index, px, py) for every image pixel on a panel - panel by panel, row by row - with
    // index = y * width + x in the image and (px, py) its centre on the plane. Pixels on no panel (gaps) are
    // not visited. Empty table: one panel -1 that is the whole image, plane == image.
    template<class F>
    void ForEachPixel(int64_t width, int64_t height, F &&f) const {
        if (modules.empty()) {
            for (int64_t y = 0; y < height; y++)
                for (int64_t x = 0; x < width; x++)
                    f(int64_t(-1), static_cast<size_t>(y * width + x), static_cast<double>(x), static_cast<double>(y));
            return;
        }
        for (size_t p = 0; p < modules.size(); p++) {
            const DetectorModule &m = modules[p];
            for (int64_t y = m.y0; y < m.y0 + m.height; y++)
                for (int64_t x = m.x0; x < m.x0 + m.width; x++) {
                    const auto [px, py] = m.ImageToPlane(static_cast<double>(x), static_cast<double>(y));
                    f(static_cast<int64_t>(p), static_cast<size_t>(y * width + x), px, py);
                }
        }
    }

    // Each module moved by the shift whose centre is nearest its own plane centre (and within it).
    void Apply(const std::vector<ModuleShift> &shifts);
    // The modules as refinement starting points: their plane centres, no move.
    [[nodiscard]] std::vector<ModuleShift> Centres() const;
};

// The modules of a modular (JUNGFRAU) detector where its assembled image puts them: one box per module,
// CONVERTED_MODULE_COLS x CONVERTED_MODULE_LINES turned the way the geometry turns it, placed as imaged.
// Empty for a detector of one module.
DetectorModules ModulesAsImaged(const DetectorGeometry &geom);

// The modules of a JUNGFRAU image where only its pixel mask says where they are: every region of pixels
// clear of the module-gap bit (PixelMask::ModuleGapPixelBit) has to be one whole module,
// CONVERTED_MODULE_COLS x CONVERTED_MODULE_LINES either way round, and is taken as placed as imaged.
// Empty where any region is not, or where there are fewer than two. The mask is row-major, width x height.
DetectorModules ModulesFromGapMask(const std::vector<uint32_t> &mask, int64_t width, int64_t height);