Jungfraujoch/common/PixelMask.h

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

#pragma once

#include "CompressedImage.h"
#include "DetectorSetup.h"
#include "DiffractionExperiment.h"
#include "../jungfrau/JFCalibration.h"

struct PixelMaskStatistics {
    uint32_t user_mask;
    uint32_t noisy_pixel;
    uint32_t error_pixel;
    uint32_t chip_gap_pixel;
    uint32_t total_masked;
    uint32_t module_gap_pixel;
};

class PixelMask {
    std::vector<uint32_t> mask;
    std::vector<uint32_t> raw_mask;
    // One byte per pixel, 1 where the pixel is masked at all - the form the GPU image preprocessor
    // uploads - and the checksum the shared device-table cache keys on (CudaSharedTables.h). Both are
    // pure functions of `mask`, and an analysis engine is built per worker per pass, so they are
    // derived here once instead of in every one of those engines.
    std::vector<uint8_t> binary_mask;
    uint64_t binary_mask_checksum = 0;
    // The same flag packed 32 pixels to a word (pixel i is bit i % 32 of word i / 32) - the form the CPU
    // image preprocessor reads, derived here for the same reason.
    std::vector<uint32_t> packed_mask;

    uint32_t LoadMask(const std::vector<uint32_t>& mask, uint8_t bit);
    // Everything that follows from the mask, recomputed wherever the mask changes.
    void UpdateDerived(const DiffractionExperiment& experiment);
    void UpdateBinaryMask();

    void CalcEdgePixels_i(const DiffractionExperiment& experiment);

public:
    // NXmx bits
    constexpr static const uint8_t ModuleGapPixelBit = 0;
    constexpr static const uint8_t ErrorPixelBit = 1;
    constexpr static const uint8_t NoisyPixelBit = 4;
    constexpr static const uint8_t UserMaskedPixelBit = 8;
    constexpr static const uint8_t BeamStopPixelBit = 9;
    // Found bad on the run's own frames (rugnux pre-scan, HotPixelFinder): lit above its ring on too
    // many frames, or holding the error value on most of them.
    constexpr static const uint8_t HotPixelBit = 10;
    constexpr static const uint8_t ChipGapPixelBit = 31;
    constexpr static const uint8_t ModuleEdgePixelBit = 30;

    PixelMask();
    explicit PixelMask(size_t width, size_t height);
    explicit PixelMask(const DiffractionExperiment& experiment);
    explicit PixelMask(const std::vector<uint32_t> &mask);

    void CalcEdgePixels(const DiffractionExperiment& experiment);
    void LoadUserMask(const DiffractionExperiment& experiment, const std::vector<uint32_t>& mask);
    void LoadUserMask(const DiffractionExperiment& experiment, const CompressedImage& image);
    void LoadBeamStopMask(const DiffractionExperiment& experiment, const std::vector<uint32_t>& mask);
    // The beam-stop shadow belongs to the run that found it, not to the dataset, so a mask read back
    // from a file that carries one starts clear. The user mask (bit 8) is deliberately left alone.
    void ClearBeamStopMask(const DiffractionExperiment& experiment);
    void LoadHotPixelMask(const DiffractionExperiment& experiment, const std::vector<uint32_t>& mask);
    void LoadDECTRISBadPixelMask(const std::vector<uint32_t>& mask);
    void LoadDarkBadPixelMask(const DiffractionExperiment& experiment, const std::vector<uint32_t>& mask);
    void LoadDetectorBadPixelMask(const DiffractionExperiment& experiment, const JFCalibration *calib);

    [[nodiscard]] const std::vector<uint32_t> &GetMaskRaw() const;
    [[nodiscard]] const std::vector<uint32_t> &GetMask(const DiffractionExperiment& experiment) const;
    [[nodiscard]] const std::vector<uint32_t> &GetMask() const;
    [[nodiscard]] const std::vector<uint8_t> &GetBinaryMask() const;
    [[nodiscard]] const std::vector<uint32_t> &GetPackedMask() const;
    [[nodiscard]] uint64_t GetBinaryMaskChecksum() const;

    [[nodiscard]] std::vector<uint32_t> GetUserMask(const DiffractionExperiment& experiment) const;
    [[nodiscard]] std::vector<uint32_t> GetUserMask() const;
    [[nodiscard]] PixelMaskStatistics GetStatistics() const;
};