Jungfraujoch/common/DiffractionExperiment.h

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

#pragma once

#include <chrono>
#include <exception>
#include <optional>

#include "../compression/CompressionAlgorithmEnum.h"

#include "UnitCell.h"
#include "Coord.h"
#include "Definitions.h"
#include "JFJochMessages.h"
#include "DetectorSetup.h"
#include "../image_analysis/spot_finding/SpotFindingSettings.h"
#include "DatasetSettings.h"
#include "ROIMap.h"
#include "InstrumentMetadata.h"
#include "ImageFormatSettings.h"
#include "DetectorSettings.h"
#include "AzimuthalIntegrationSettings.h"
#include "FileWriterSettings.h"
#include "DiffractionGeometry.h"
#include "CompressedImage.h"
#include "IndexingSettings.h"
#include "BraggIntegrationSettings.h"
#include "ScalingSettings.h"

#include <gemmi/symmetry.hpp>

enum class DetectorMode {
    Standard, PedestalG0, PedestalG1, PedestalG2, DarkMask
};

struct AcquisitionDeviceNetConfig {
    std::string mac_addr;
    std::string ipv4_addr;
    uint64_t udp_port;
};

struct DetectorModuleConfig {
    uint64_t udp_dest_port_1;
    uint64_t udp_dest_port_2;
    std::string ipv4_src_addr_1;
    std::string ipv4_src_addr_2;
    std::string ipv4_dest_addr_1;
    std::string ipv4_dest_addr_2;
    std::string mac_addr_dest_1;
    std::string mac_addr_dest_2;
    uint32_t module_id_in_data_stream;
    uint32_t data_stream;
};

class DiffractionExperiment {

    int64_t ndatastreams;

    std::optional<std::chrono::microseconds> zmq_preview_period = std::chrono::seconds(1);

    DetectorMode mode;

    uint64_t series_id;

    // Dataset settings
    DatasetSettings dataset;
    InstrumentMetadata instrument;
    ImageFormatSettings image_format_settings;
    DetectorSettings detector_settings;
    AzimuthalIntegrationSettings az_integration_settings;
    DetectorSetup detector;
    FileWriterSettings file_writer;
    IndexingSettings indexing;
    BraggIntegrationSettings bragg_integration_settings;
    ScalingSettings scaling_settings;

    DarkMaskSettings dark_mask_settings;

    ROIMap roi_mask;
    int64_t summation;
    bool cpu_summation;

    std::string detector_update_zmq_addr;
public:
    // Public methods are atomic
    DiffractionExperiment();
    DiffractionExperiment(const DetectorSetup& geom);

    // Methods below can be chained together
    DiffractionExperiment& Detector(const DetectorSetup& input);
    DiffractionExperiment& Mode(DetectorMode input);
    DiffractionExperiment& DataStreams(int64_t input);

    DiffractionExperiment& PedestalG0Frames(int64_t input);
    DiffractionExperiment& PedestalG1Frames(int64_t input);
    DiffractionExperiment& PedestalG2Frames(int64_t input);
    DiffractionExperiment& FrameTime(std::chrono::nanoseconds frame_time,
                                     std::chrono::nanoseconds in_count_time = std::chrono::nanoseconds(0));

    DiffractionExperiment& ZMQPreviewPeriod(const std::optional<std::chrono::microseconds> &input);

    DiffractionExperiment& UseInternalPacketGenerator(bool input);
    DiffractionExperiment& InternalPacketGeneratorImages(int64_t input);
    DiffractionExperiment& MaskModuleEdges(bool input);
    DiffractionExperiment& MaskChipEdges(bool input);

    DiffractionExperiment& QRangeForAzimInt_recipA(float low, std::optional<float> high);
    DiffractionExperiment& BkgEstimateQRange_recipA(float low, float high);
    DiffractionExperiment& QSpacingForAzimInt_recipA(float input);

    DiffractionExperiment& StorageCells(int64_t input);
    DiffractionExperiment& StorageCellStart(int64_t input = 15);

    DiffractionExperiment& UsingGainHG0(bool input);
    DiffractionExperiment& FixedGainG1(bool input);
    DiffractionExperiment& IncrementRunNumber();
    DiffractionExperiment& JungfrauConvPhotonCnt(bool input);
    DiffractionExperiment& PulsedSource(bool input);

    DiffractionExperiment& ImagesPerTrigger(int64_t input);
    DiffractionExperiment& NumTriggers(int64_t triggers);
    DiffractionExperiment& IncidentEnergy_keV(float input);
    DiffractionExperiment& BeamX_pxl(float input);
    DiffractionExperiment& BeamY_pxl(float input);
    DiffractionExperiment& DetectorDistance_mm(float input);
    DiffractionExperiment& FilePrefix(std::string input);
    DiffractionExperiment& FilePrefixTrusted(std::string input);   // offline/local: no CheckPath guard (absolute allowed)
    DiffractionExperiment& Compression(CompressionAlgorithm input);
    DiffractionExperiment& SetUnitCell(const std::optional<UnitCell> &cell);
    DiffractionExperiment& SpaceGroupNumber(std::optional<int64_t> input);
    DiffractionExperiment& SampleName(const std::string &input);
    DiffractionExperiment& AttenuatorTransmission(const std::optional<float> &input);
    DiffractionExperiment& TotalFlux(const std::optional<float> &input);
    DiffractionExperiment& Goniometer(const std::optional<GoniometerAxis> &input);
    DiffractionExperiment& Smargon(const std::optional<SmargonPosition> &input);
    DiffractionExperiment& HeaderAppendix(const nlohmann::json& input);
    DiffractionExperiment& ImageAppendix(const nlohmann::json& input);
    DiffractionExperiment& Summation(int64_t input);
    DiffractionExperiment& MaxSpotCount(int64_t input);
    DiffractionExperiment& ImagesPerFile(const std::optional<int64_t> &input);
    DiffractionExperiment& LossyCompressionSerialMX(float input);
    DiffractionExperiment& LossyCompressionPoisson(const std::optional<int64_t> &input);
    DiffractionExperiment& SaveCalibration(const std::optional<bool> &input);
    DiffractionExperiment& ImportDatasetSettings(const DatasetSettings& input);
    DiffractionExperiment& EigerBitDepth(const std::optional<int64_t> &input);
    DiffractionExperiment& ImportInstrumentMetadata(const InstrumentMetadata& input);
    DiffractionExperiment& ApplyPixelMask(bool input);
    DiffractionExperiment& ElectronSource(bool input);
    DiffractionExperiment& OverwriteExistingFiles(bool input);
    DiffractionExperiment& SetFileWriterFormat(FileWriterFormat input);
    DiffractionExperiment& IndexingAlgorithm(IndexingAlgorithmEnum input);
    DiffractionExperiment& GeomRefinementAlgorithm(GeomRefinementAlgorithmEnum input);
    DiffractionExperiment& IndexingTolerance(float input);
    DiffractionExperiment& GridScan(const std::optional<GridScanSettings>& input);
    DiffractionExperiment& RingCurrent_mA(const std::optional<float>& input);
    DiffractionExperiment& SampleTemperature_K(const std::optional<float> &input);
    DiffractionExperiment& PoniRot1_rad(float input);
    DiffractionExperiment& PoniRot2_rad(float input);
    DiffractionExperiment& PoniRot3_rad(float input);
    DiffractionExperiment& FluorescenceSpectrum(const XrayFluorescenceSpectrum& input);
    DiffractionExperiment& DetectIceRings(bool input);
    DiffractionExperiment& RunNumber(uint64_t input);

    IndexingSettings GetIndexingSettings() const;
    DiffractionExperiment& ImportIndexingSettings(const IndexingSettings &input);

    InstrumentMetadata GetInstrumentMetadata() const;

    DiffractionExperiment& ImportImageFormatSettings(const ImageFormatSettings& input);
    ImageFormatSettings GetImageFormatSettings() const;

    DiffractionExperiment& ImportDetectorSettings(const DetectorSettings& input);
    DetectorSettings GetDetectorSettings() const;

    DiffractionExperiment& ImportAzimuthalIntegrationSettings(const AzimuthalIntegrationSettings& input);
    AzimuthalIntegrationSettings GetAzimuthalIntegrationSettings() const;

    // Highest q (2*pi/d) any pixel of the detector reaches, from the current geometry, and the same
    // limit as a resolution in Angstrom. This is what an unset azimuthal-integration high q and an unset
    // spot-finding high-resolution limit resolve to.
    [[nodiscard]] float GetDetectorMaxQ_recipA() const;
    [[nodiscard]] float GetDetectorMaxResolution_A() const;

    DiffractionExperiment& ImportBraggIntegrationSettings(const BraggIntegrationSettings& input);
    BraggIntegrationSettings GetBraggIntegrationSettings() const;

    DiffractionExperiment& ImportFileWriterSettings(const FileWriterSettings& input);
    FileWriterSettings GetFileWriterSettings() const;

    DiffractionExperiment &ImportDarkMaskSettings(const DarkMaskSettings &input);
    DarkMaskSettings GetDarkMaskSettings() const;

    DiffractionExperiment &ImportScalingSettings(const ScalingSettings& input);
    ScalingSettings GetScalingSettings() const;

    DatasetSettings GetDatasetSettings() const;

    void FillMessage(StartMessage &message) const;

    static void CheckDataProcessingSettings(const SpotFindingSettings& settings);
    static SpotFindingSettings DefaultDataProcessingSettings();

    DetectorMode GetDetectorMode() const;

    int64_t GetBitDepthReadout() const; // 12 bit is OK :)

    int64_t GetSaturationLimit() const;
    int64_t GetOverflow() const;
    int64_t GetUnderflow() const;

    int64_t GetPedestalG0Frames() const;
    int64_t GetPedestalG1Frames() const;
    int64_t GetPedestalG2Frames() const;

    int64_t GetImageNum() const;

    int64_t GetFrameNum() const;
    int64_t GetFrameNumPerTrigger() const;

    std::chrono::nanoseconds GetFrameTime() const;
    std::chrono::nanoseconds GetDetectorPeriod() const;
    std::chrono::nanoseconds GetImageTime() const;

    std::chrono::nanoseconds GetImageCountTime() const;
    std::chrono::nanoseconds GetFrameCountTime() const;
    bool GetFrameCountTimeAuto() const;

    DiffractionExperiment& StorageCellDelay(std::chrono::nanoseconds input);
    std::chrono::nanoseconds GetStorageCellDelay() const;

    DiffractionExperiment& DetectorDelay(std::chrono::nanoseconds input);
    std::chrono::nanoseconds GetDetectorDelay() const;

    int64_t GetMaxCompressedSize() const;
    int64_t GetImageBufferLocationSize() const;

    // Room reserved in every image-buffer slot for the per-image CBOR metadata, on top of the
    // compressed image. CBORSerialize_Image_MetadataHeadroom holds it against the largest metadata
    // the serializer can produce, so change the two together.
    static constexpr int64_t kImageMetadataHeadroom = 4 * 1024 * 1024;

    int64_t GetDataStreamsNum() const;
    int64_t GetModulesNum(uint16_t data_stream) const;
    int64_t GetModulesNum() const;
    int64_t GetFirstModuleOfDataStream(uint16_t data_stream) const;

    int64_t GetPixelsNum() const;
    int64_t GetYPixelsNum() const;
    int64_t GetXPixelsNum() const;

    int64_t GetPixelsNumConv() const; // This is actual / converted number
    int64_t GetYPixelsNumConv() const;
    int64_t GetXPixelsNumConv() const;

    int64_t GetPixel0OfModuleConv(uint16_t module_number) const;
    int64_t GetModuleFastDirectionStep(uint16_t module_number) const;
    int64_t GetModuleSlowDirectionStep(uint16_t module_number) const;

    Coord GetModuleFastDirection(uint16_t module_number) const;
    Coord GetModuleSlowDirection(uint16_t module_number) const;

    std::optional<std::chrono::microseconds> GetZMQPreviewPeriod() const;

    int64_t GetDefaultPlotBinning() const;

    bool IsUsingInternalPacketGen() const;
    int64_t GetInternalPacketGeneratorImages() const;

    uint32_t GetSrcIPv4Address(uint32_t data_stream, uint32_t half_module) const;

    bool GetMaskModuleEdges() const;
    bool GetMaskChipEdges() const;

    float GetLowQForAzimInt_recipA() const;
    float GetHighQForAzimInt_recipA() const;
    float GetQSpacingForAzimInt_recipA() const;

    float GetLowQForBkgEstimate_recipA() const;
    float GetHighQForBkgEstimate_recipA() const;

    int64_t GetStorageCellNumber() const;
    int64_t GetStorageCellStart() const;

    int64_t GetMaxSpotCount() const;

    float GetPixelSize_mm() const;
    std::string GetSourceName() const;
    std::string GetSourceType() const;
    std::string GetInstrumentName() const;
    std::string GetDetectorDescription() const;

    std::vector<std::string> GetDetectorModuleHostname() const;

    DiffractionExperiment& ApplySolidAngleCorr(bool input);
    DiffractionExperiment& PolarizationFactor(const std::optional<float> &input);
    DiffractionExperiment& BandwidthFWHM(const std::optional<float> &input);

    bool GetApplySolidAngleCorr() const;
    std::optional<float> GetPolarizationFactor() const;
    std::optional<float> GetBandwidthFWHM() const;

    int64_t GetUDPInterfaceCount() const;
    std::vector<DetectorModuleConfig> GetDetectorModuleConfig(const std::vector<AcquisitionDeviceNetConfig>& net_config) const;

    bool IsFixedGainG1() const;
    bool IsUsingGainHG0() const;

    uint64_t GetRunNumber() const;
    std::string GetRunName() const;
    bool IsJungfrauConvPhotonCnt() const;

    const DetectorSetup& GetDetectorSetup() const;
    DetectorSetup& Detector();

    bool IsPulsedSource() const;
    bool IsElectronSource() const;

    bool IsSpotFindingEnabled() const;

    float GetPhotonEnergyForConversion_keV() const;

    std::optional<float> GetAttenuatorTransmission() const;
    std::optional<float> GetTotalFlux() const;
    std::optional<GoniometerAxis> GetGoniometer() const;
    std::optional<GridScanSettings> GetGridScan() const;

    const nlohmann::json& GetHeaderAppendix() const;
    const nlohmann::json& GetImageAppendix() const;
    std::optional<UnitCell> GetUnitCell() const;
    std::string GetUnitCellString() const;
    std::optional<int64_t> GetSpaceGroupNumber() const;
    bool GetSaveCalibration() const;
    int64_t GetSummation() const;
    int64_t GetFPGASummation() const;
    std::string GetSampleName() const;
    float GetIncidentEnergy_keV() const;
    float GetWavelength_A() const;
    float GetBeamX_pxl() const;
    float GetBeamY_pxl() const;
    float GetDetectorDistance_mm() const;

    Coord GetScatteringVector() const;

    std::string GetFilePrefix() const;

    CompressionAlgorithm GetCompressionAlgorithm() const;
    int64_t GetNumTriggers() const;

    ROIMap& ROI();
    const ROIMap& ROI() const;
    std::vector<uint16_t> ExportROIMap() const;
    int64_t GetImagesPerFile() const;

    float GetLossyCompressionSerialMX() const;
    std::optional<int64_t> GetLossyCompressionPoisson() const;
    std::string GetExperimentGroup() const;

    std::optional<int64_t> GetPixelValueLowThreshold() const;
    DiffractionExperiment &PixelValueLowThreshold(const std::optional<int64_t>& input);

    std::optional<int64_t> GetPixelValueHighThreshold() const;
    DiffractionExperiment &PixelValueHighThreshold(const std::optional<int64_t>& input);

    DiffractionExperiment& BitDepthImage(const std::optional<int64_t> &input);
    DiffractionExperiment& PixelSigned(const std::optional<bool> &input);

    bool IsGeometryTransformed() const;
    DiffractionExperiment& GeometryTransformation(bool input);
    int64_t GetImageFillValue() const;

    DiffractionExperiment& JungfrauConversionFactor_keV(const std::optional<float> &input);
    std::optional<float> GetJungfrauConversionFactor_keV() const;
    DiffractionExperiment& AutoSummation(bool input);
    bool GetAutoSummation() const;

    int64_t GetByteDepthImage() const;
    int64_t GetByteDepthFPGA() const;
    bool IsPixelSigned() const;

    bool IsPedestalRun() const;

    DiffractionExperiment &Raw();
    DiffractionExperiment &Conversion();

    float GetPedestalG0RMSLimit() const;
    uint32_t GetPedestalMinImageCount() const;
    float GetEigerThreshold_keV() const;
    int64_t GetEigerBitDepth() const;
    DetectorTiming GetDetectorTiming() const;

    bool IsDetectorModuleSync() const;
    bool IsDetectorMirroredY() const;
    // The sample transformation chain in mounting order, base first. See TransformationAxis.
    [[nodiscard]] std::vector<TransformationAxis> BuildTransformationChain() const;

    [[nodiscard]] DetectorType GetDetectorType() const;
    [[nodiscard]] bool IsMaskPixelsWithoutG0() const;
    [[nodiscard]] bool IsApplyPixelMask() const;

    DiffractionExperiment& CPUSummation(bool input);
    [[nodiscard]] bool IsCPUSummation() const;

    [[nodiscard]] FileWriterFormat GetFileWriterFormat() const;

    DiffractionGeometry GetDiffractionGeometry() const;
    void CalcAzIntCorrRawCoord(float *output, size_t module_number) const;
    void CalcSpotFinderResolutionMap(float *data, size_t module_number) const;
    CompressedImageMode GetImageMode() const;

    // Resolves the configured algorithm to a concrete one from GPU availability and
    // unit-cell presence: Auto -> FFTW/FFT/FFBIDX, and FFBIDX without a cell -> None.
    // Never returns Auto - the indexer pool relies on this and cannot resolve Auto itself.
    IndexingAlgorithmEnum GetIndexingAlgorithm() const;
    GeomRefinementAlgorithmEnum GetGeomRefinementAlgorithm() const;
    float GetIndexingTolerance() const;

    std::optional<float> GetSampleTemperature_K() const;
    std::optional<float> GetRingCurrent_mA() const;

    float GetPoniRot1_rad() const;
    float GetPoniRot2_rad() const;
    float GetPoniRot3_rad() const;

    std::optional<gemmi::SpaceGroup> GetGemmiSpaceGroup() const;
    gemmi::CrystalSystem GetCrystalSystem() const;
    std::string GetSpaceGroupName() const;
    char GetCentering() const;
    const XrayFluorescenceSpectrum &GetFluorescenceSpectrum() const;

    bool IsDetectIceRings() const;
    int64_t GetDarkMaskNumberOfFrames() const;

    bool IsRotationIndexing() const;

    std::optional<double> GetRotationWedgeForScaling() const;
    bool GetRefineRotationWedgeInScaling() const;
};