Jungfraujoch/common/DetectorModules.cpp

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

#include "DetectorModules.h"
#include "DetectorGeometry.h"
#include "Definitions.h"
#include "PixelMask.h"

#include <algorithm>
#include <cmath>
#include <tuple>

DetectorModule DetectorModule::AsImaged(int64_t x0, int64_t y0, int64_t width, int64_t height) {
    DetectorModule m;
    m.x0 = x0;
    m.y0 = y0;
    m.width = width;
    m.height = height;
    m.ox = static_cast<double>(x0);
    m.oy = static_cast<double>(y0);
    return m;
}

bool DetectorModule::ContainsImagePixel(double x, double y) const {
    // A pixel's centre is at its index, so the module covers [x0 - 0.5, x0 + width - 0.5).
    return x >= static_cast<double>(x0) - 0.5 && x < static_cast<double>(x0 + width) - 0.5
           && y >= static_cast<double>(y0) - 0.5 && y < static_cast<double>(y0 + height) - 0.5;
}

std::pair<double, double> DetectorModule::ImageToPlane(double x, double y) const {
    const double u = x - static_cast<double>(x0), v = y - static_cast<double>(y0);
    return {ox + u * fx + v * sx, oy + u * fy + v * sy};
}

std::pair<double, double> DetectorModule::PlaneToImage(double px, double py) const {
    // Inverse of the 2x2 [f s]
    const double det = fx * sy - sx * fy;
    const double dx = px - ox, dy = py - oy;
    const double u = (sy * dx - sx * dy) / det;
    const double v = (-fy * dx + fx * dy) / det;
    return {static_cast<double>(x0) + u, static_cast<double>(y0) + v};
}

std::pair<double, double> DetectorModule::PlaneCentre() const {
    return ImageToPlane(static_cast<double>(x0) + 0.5 * static_cast<double>(width - 1),
                        static_cast<double>(y0) + 0.5 * static_cast<double>(height - 1));
}

void DetectorModule::Apply(const ModuleShift &shift) {
    const double c = std::cos(shift.rot_rad), s = std::sin(shift.rot_rad);
    const double ux = ox - shift.x0, uy = oy - shift.y0;
    ox = shift.x0 + c * ux - s * uy + shift.dx;
    oy = shift.y0 + s * ux + c * uy + shift.dy;
    const double nfx = c * fx - s * fy, nfy = s * fx + c * fy;
    const double nsx = c * sx - s * sy, nsy = s * sx + c * sy;
    fx = nfx; fy = nfy;
    sx = nsx; sy = nsy;
}

DetectorModules::DetectorModules(std::vector<DetectorModule> in) : modules(std::move(in)) {}

bool DetectorModules::empty() const {
    return modules.empty();
}

size_t DetectorModules::size() const {
    return modules.size();
}

const DetectorModule &DetectorModules::operator[](size_t i) const {
    return modules.at(i);
}

const std::vector<DetectorModule> &DetectorModules::Get() const {
    return modules;
}

bool DetectorModules::Moved() const {
    // Not an exact comparison: a module read back from a file went through a length and a unit vector
    // in single precision, which puts a module as imaged a few 1e-4 px off its box when it is far from
    // pixel (0, 0) on both axes.
    constexpr double tolerance = 1e-3;
    for (const auto &m: modules)
        if (std::fabs(m.ox - static_cast<double>(m.x0)) > tolerance || std::fabs(m.oy - static_cast<double>(m.y0)) > tolerance
            || std::fabs(m.fx - 1.0) > tolerance / CONVERTED_MODULE_COLS || std::fabs(m.fy) > tolerance / CONVERTED_MODULE_COLS
            || std::fabs(m.sx) > tolerance / CONVERTED_MODULE_COLS || std::fabs(m.sy - 1.0) > tolerance / CONVERTED_MODULE_COLS)
            return true;
    return false;
}

int64_t DetectorModules::ModuleOfImagePixel(double x, double y) const {
    for (size_t i = 0; i < modules.size(); i++)
        if (modules[i].ContainsImagePixel(x, y))
            return static_cast<int64_t>(i);
    return -1;
}

int64_t DetectorModules::ModuleOfPlanePoint(double px, double py) const {
    for (size_t i = 0; i < modules.size(); i++) {
        const auto [x, y] = modules[i].PlaneToImage(px, py);
        if (modules[i].ContainsImagePixel(x, y))
            return static_cast<int64_t>(i);
    }
    return -1;
}

std::pair<double, double> DetectorModules::ImageToPlane(double x, double y) const {
    const int64_t m = ModuleOfImagePixel(x, y);
    if (m < 0)
        return {x, y};
    return modules[m].ImageToPlane(x, y);
}

std::pair<double, double> DetectorModules::PlaneToImage(double px, double py) const {
    const int64_t m = ModuleOfPlanePoint(px, py);
    if (m < 0)
        return {px, py};
    return modules[m].PlaneToImage(px, py);
}

std::pair<double, double> DetectorModules::ImageToPlane(int64_t panel, double x, double y) const {
    if (panel < 0)
        return {x, y};
    return modules.at(panel).ImageToPlane(x, y);
}

std::pair<double, double> DetectorModules::PlaneToImage(int64_t panel, double px, double py) const {
    if (panel < 0)
        return {px, py};
    return modules.at(panel).PlaneToImage(px, py);
}

std::optional<PanelPoint> DetectorModules::PlaneToPanel(double px, double py) const {
    if (modules.empty())
        return PanelPoint{-1, px, py};
    for (size_t i = 0; i < modules.size(); i++) {
        const auto [x, y] = modules[i].PlaneToImage(px, py);
        if (modules[i].ContainsImagePixel(x, y))
            return PanelPoint{static_cast<int64_t>(i), x, y};
    }
    return std::nullopt;
}

PlaneBox DetectorModules::PlaneBounds(int64_t width, int64_t height) const {
    if (modules.empty())
        return {0.0, 0.0, static_cast<double>(width), static_cast<double>(height)};
    PlaneBox box{INFINITY, INFINITY, -INFINITY, -INFINITY};
    for (const auto &m: modules) {
        const double xa = static_cast<double>(m.x0) - 0.5, xb = static_cast<double>(m.x0 + m.width) - 0.5;
        const double ya = static_cast<double>(m.y0) - 0.5, yb = static_cast<double>(m.y0 + m.height) - 0.5;
        for (const auto &[x, y]: {std::pair{xa, ya}, std::pair{xb, ya}, std::pair{xa, yb}, std::pair{xb, yb}}) {
            const auto [px, py] = m.ImageToPlane(x, y);
            box.x0 = std::min(box.x0, px);
            box.y0 = std::min(box.y0, py);
            box.x1 = std::max(box.x1, px);
            box.y1 = std::max(box.y1, py);
        }
    }
    return box;
}

void DetectorModules::Apply(const std::vector<ModuleShift> &shifts) {
    for (auto &m: modules) {
        const auto [cx, cy] = m.PlaneCentre();
        const double reach = 0.5 * static_cast<double>(std::min(m.width, m.height));
        const ModuleShift *best = nullptr;
        for (const auto &s: shifts)
            if (std::hypot(s.x0 - cx, s.y0 - cy) < reach
                && (!best || std::hypot(s.x0 - cx, s.y0 - cy) < std::hypot(best->x0 - cx, best->y0 - cy)))
                best = &s;
        if (best)
            m.Apply(*best);
    }
}

std::vector<ModuleShift> DetectorModules::Centres() const {
    std::vector<ModuleShift> out;
    for (const auto &m: modules) {
        ModuleShift s;
        std::tie(s.x0, s.y0) = m.PlaneCentre();
        out.push_back(s);
    }
    return out;
}

DetectorModules ModulesAsImaged(const DetectorGeometry &geom) {
    if (geom.GetModulesNum() < 2)
        return {};
    std::vector<DetectorModule> out;
    for (int64_t m = 0; m < geom.GetModulesNum(); m++) {
        const Coord f = geom.GetFastDirection(m), s = geom.GetSlowDirection(m);
        const double cols = CONVERTED_MODULE_COLS - 1, lines = CONVERTED_MODULE_LINES - 1;
        const double xa = static_cast<double>(geom.GetX0(m)), ya = static_cast<double>(geom.GetY0(m));
        const double xb = xa + cols * f.x + lines * s.x, yb = ya + cols * f.y + lines * s.y;
        const auto x0 = static_cast<int64_t>(std::lround(std::min(xa, xb)));
        const auto y0 = static_cast<int64_t>(std::lround(std::min(ya, yb)));
        out.push_back(DetectorModule::AsImaged(x0, y0, std::lround(std::abs(xb - xa)) + 1,
                                               std::lround(std::abs(yb - ya)) + 1));
    }
    return DetectorModules(out);
}

DetectorModules ModulesFromGapMask(const std::vector<uint32_t> &mask, int64_t width, int64_t height) {
    if (width <= 0 || height <= 0 || mask.size() != static_cast<size_t>(width * height))
        return {};
    const auto clear = [&](int64_t x, int64_t y) {
        return x >= 0 && y >= 0 && x < width && y < height
               && (mask[y * width + x] & (1u << PixelMask::ModuleGapPixelBit)) == 0;
    };
    std::vector<uint8_t> taken(mask.size(), 0);
    std::vector<DetectorModule> out;
    for (int64_t y = 0; y < height; y++)
        for (int64_t x = 0; x < width; x++) {
            if (!clear(x, y) || taken[y * width + x])
                continue;
            // The first clear pixel of a region not yet taken is a module's corner: the module runs to the
            // gap along the row and down the column, and has to be whole and bounded by gap all round.
            int64_t w = 0, h = 0;
            while (clear(x + w, y)) w++;
            while (clear(x, y + h)) h++;
            if (!((w == CONVERTED_MODULE_COLS && h == CONVERTED_MODULE_LINES)
                  || (w == CONVERTED_MODULE_LINES && h == CONVERTED_MODULE_COLS)))
                return {};
            for (int64_t yy = y - 1; yy <= y + h; yy++)
                for (int64_t xx = x - 1; xx <= x + w; xx++) {
                    const bool inside = xx >= x && xx < x + w && yy >= y && yy < y + h;
                    if (clear(xx, yy) != inside)
                        return {};
                    if (inside)
                        taken[yy * width + xx] = 1;
                }
            out.push_back(DetectorModule::AsImaged(x, y, w, h));
        }
    if (out.size() < 2)
        return {};
    return DetectorModules(out);
}