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

#pragma once

#include <memory>
#include <mutex>
#include <optional>
#include "DiffractionExperiment.h"
#include "PixelMask.h"

class AzimuthalIntegrationMapping {
protected:
    const AzimuthalIntegrationSettings settings;
    const float wavelength;
    const size_t width, height;

    std::vector<float> bin_to_q;
    std::vector<float> bin_to_2theta;
    std::vector<float> bin_to_d;
    std::vector<float> bin_to_phi;
    std::vector<uint16_t> pixel_to_bin;

    std::vector<float> pixel_resolution;
    std::vector<float> corrections;

    std::optional<float> polarization_factor;

    // Checksums of the two tables the GPU engines upload, taken once here. They are part of the key
    // the shared device-table cache looks them up by (CudaSharedTables.h), so an engine that hands
    // its own checksum in does not have to hash tens of megabytes on its way to a cache hit - and
    // there is one engine per worker per pass. Both vectors are written in the constructor and never
    // touched again, so a checksum taken there stays true for the mapping's whole life.
    uint64_t pixel_to_bin_checksum = 0;
    uint64_t corrections_checksum = 0;

    // Bit-packed "this pixel is outside the resolution limits" mask, memoised for the limits it was
    // last built for. Every worker's spot finder wants the identical mask and building it walks every
    // pixel, so it is built once and shared. The mapping is const and read from all the worker threads
    // at once, hence the mutex; and the mask is handed out as a shared_ptr, so a caller keeps the one
    // it was given even if another thread later replaces the cached one.
    mutable std::mutex res_mask_mutex;
    mutable std::optional<float> res_mask_high, res_mask_low;
    mutable std::shared_ptr<const std::vector<uint32_t>> res_mask_bits;

    size_t nthreads;

    void UpdateMaxBinNumber();

    void SetupRawGeom(const DiffractionExperiment& experiment, const std::vector<uint32_t> &mask);
    void SetupConvGeomRows(const DiffractionGeometry &geom, const std::vector<uint32_t> &mask, size_t row0, size_t row_end);
    void SetupConvGeom(const DiffractionGeometry &geom, const std::vector<uint32_t> &mask);

    void SetupPixel(const DiffractionGeometry &geom, const std::vector<uint32_t> &mask,
                    uint32_t pxl, uint32_t col, uint32_t row);
public:

    AzimuthalIntegrationMapping(const DiffractionExperiment& experiment,
                         const PixelMask& mask,
                         size_t nthreads = 0);

    [[nodiscard]] uint16_t GetBinNumber() const;
    [[nodiscard]] const std::vector<uint16_t>& GetPixelToBin() const;
    [[nodiscard]] const std::vector<float> &GetBinToQ() const;
    [[nodiscard]] const std::vector<float> &GetBinToD() const;
    [[nodiscard]] const std::vector<float> &GetBinToTwoTheta() const;
    [[nodiscard]] const std::vector<float> &GetBinToPhi() const;
    [[nodiscard]] uint16_t QToBin(float q) const;
    [[nodiscard]] const std::vector<float> &Corrections() const;
    [[nodiscard]] const std::vector<float> &Resolution() const;
    [[nodiscard]] uint64_t GetPixelToBinChecksum() const;
    [[nodiscard]] uint64_t GetCorrectionsChecksum() const;
    // Pixels the spot finders must ignore because their resolution falls outside the limits, packed
    // 32 pixels to a word (bit set = ignore), in the finders' own layout.
    [[nodiscard]] std::shared_ptr<const std::vector<uint32_t>>
        ResolutionMaskBits(std::optional<float> high_res, std::optional<float> low_res) const;
    [[nodiscard]] size_t GetWidth() const;
    [[nodiscard]] size_t GetHeight() const;
    [[nodiscard]] const AzimuthalIntegrationSettings& Settings() const;
    [[nodiscard]] int32_t GetAzimuthalBinCount() const;
    [[nodiscard]] int32_t GetQBinCount() const;
    [[nodiscard]] size_t GetNThreads() const;
};