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

#pragma once

#include <cstdint>
#include <atomic>
#include <future>
#include <memory>
#include <mutex>
#include <optional>
#include <vector>

#include "../../common/CompressedImage.h"
#include "../../common/DiffractionExperiment.h"
#include "../../common/DiffractionGeometry.h"
#include "../../common/JFJochMessages.h"
#include "../../common/PixelMask.h"
#ifdef JFJOCH_USE_CUDA
#include "ShadowAccumulatorGPU.h"
#endif

// Finds the beam-stop shadow - the central disk and the holder arm - from a set of images,
// mirroring the accumulate-then-finalize shape of DarkMaskAnalysis: feed frames with
// AddImage(), then read the mask once with GetMask(). The mask is in converted geometry
// and is 1 where the beam stop shadows the detector.
//
// The shadow is a place where the background is missing, so it is found by comparing each
// pixel's mean against the typical background at the same 2 theta - the median over its ring,
// taken over the pixels not already known to be shadowed. That comparison holds wherever the
// ring still has unshadowed pixels to measure. Where it does not - a ring lying wholly inside
// the stop - there is nothing to compare against, and such a ring is shadow in its entirety.
//
// The background belongs to the beam and the shadow to the stop, and the two are not concentric:
// the stop sits off the beam by a sizeable fraction of its own radius. Only the per-ring
// comparison is used, so nothing here assumes they share a centre.
//
// Frames are chosen by the caller; the detection needs enough of them that the background
// is counted rather than guessed (see MIN_EXPECTED_COUNTS in the .cpp).
// Thread-safe: workers call AddImage concurrently. The projection is split into bands of rows, each
// with a lock of its own, and a worker adding a frame starts at a different band from the one before
// it, so workers meet only when they reach the same band.
class ShadowFinder {
    mutable std::mutex m;

    const int   width;
    const int   height;
    float       beam_x;
    float       beam_y;

    // What the background owes to the source rather than to the hardware. The scattered background
    // is not flat around a ring: a polarized source suppresses it in its own plane by a factor that
    // reaches three at the 2 theta a short detector distance puts in a corner - several times the
    // dip this class is looking for - so the comparison divides it out before it compares. The only
    // other correction that varies along a ring is the solid angle of a tilted detector, divided out
    // with it (ring_position); detector and air absorption are functions of 2 theta alone, nearly
    // enough, and the ring's own median absorbs them.
    // The geometry is kept whole rather than reduced to a distance and a pixel size because the
    // azimuth is the whole point: detector tilt, a quarter-turned image and an in-plane rotation
    // all move the polarization plane across the stored image, and the geometry already knows where
    // it lies. Only the centre is replaced, by the one the caller measured.
    const DiffractionGeometry geometry;
    const std::optional<float> polarization;   // unset leaves the background as it was measured
    // What the last GetMask divided by: the stated degree, in the plane the background shows.
    mutable std::optional<float> applied_polarization;
    // Detector units per photon (BraggIntegrationSettings::PixelGain; 1 for counted photons). The
    // shadow is a ratio and does not see it, but the two tests made in counts do: how significant a
    // deficit is, and the count a pixel must reach to have recorded a reflection.
    const float gain;

    std::vector<uint32_t> pixel_mask;   // pixels already masked carry no background to test
    // The panels the image is made of. With any, the mask is made on the detector plane (GetMask).
    const DetectorModules modules;

    // Per-pixel projection over the frames added so far (converted geometry). The sums and counts are
    // integers and the maximum is a maximum, so the result does not depend on the order the frames
    // arrive in.
    struct Projection {
        std::vector<int64_t>  max_value;
        std::vector<int64_t>  sum_value;
        std::vector<uint32_t> valid_count;
        uint32_t frames = 0;
    };
    // The frames added on the host. Allocated by the first of them: with a GPU there are usually none.
    Projection host;
    std::mutex host_mutex;   // guards the allocation and the frame count, not the sums
    static constexpr size_t BANDS = 64;
    std::mutex band_mutex[BANDS];
    std::atomic<size_t> next_band{0};

#ifdef JFJOCH_USE_CUDA
    // One per card when GPUs are available. Frames a card can decode are accumulated on the card of the
    // thread adding them instead of on the host - only the compressed chunk crosses PCIe - and the
    // cards' projections are folded into one, with the host projection, when the mask is read. Frames
    // a card cannot take (anything but bitshuffle+LZ4) still go to the host, so a run mixing
    // compressions is handled without a second code path.
    // Each is built on a thread of its own: it allocates and clears several hundred megabytes of device
    // memory, and cudaMalloc synchronises the whole device, so doing it in the constructor would
    // stall the caller before it has read its first frame. The first AddImage waits for them, by
    // which time the reads have been running for a while.
    mutable std::vector<std::future<std::unique_ptr<ShadowAccumulatorGPU>>> gpu_pending;
    mutable std::vector<std::unique_ptr<ShadowAccumulatorGPU>> gpus;   // by card; null where none
    mutable std::mutex gpu_mutex;                  // guards the two above
    std::unique_ptr<std::mutex[]> card_mutex;      // a card's frames are added one at a time

    // The accumulator of card `device` once the constructions have finished, or null if it has none.
    ShadowAccumulatorGPU *Gpu(int device) const;
    // The one accumulator every other card's projection has been folded into, or null if there is
    // none. What the projection is read from; frames added after the first call are not seen.
    [[nodiscard]] ShadowAccumulatorGPU *Gpu() const;
#endif

    // Add the pixels [begin, end) of one frame to the host projection.
    template<class T> void Add(const T *ptr, size_t begin, size_t end);

#ifdef JFJOCH_USE_CUDA
    // The device's projection with the host's folded in. max_value is only taken from a projection
    // that actually counted the pixel - one that never saw it holds 0, which would beat a genuinely
    // negative maximum.
    [[nodiscard]] Projection Reduce() const;

    // That projection, made on its first read and kept. The ring-centre fit, the mask and the
    // beam-centre capture all read the same one, and on a 16 Mpx detector each is 360 MB brought back
    // from the device into fresh memory.
    mutable std::optional<Projection> reduced;
#endif
    // The projection the frames added so far make: the host's, or the one above where the device
    // took frames. Called with `m` held.
    [[nodiscard]] const Projection &Reduced() const;

    // The mask of a W x H projection, its known-masked pixels and the centre the rings are drawn about -
    // the image's own, or the plane raster's where the image is made of panels.
    [[nodiscard]] std::vector<uint32_t> MaskOn(const Projection &p, const std::vector<uint32_t> &known_mask,
                                               int W, int H, float centre_x, float centre_y, size_t nthreads) const;

public:
    static constexpr uint32_t SHADOW = 1;
    static constexpr uint32_t TRANSMITTING = 2;

    ShadowFinder(const DiffractionExperiment &experiment, const PixelMask &mask);
    ~ShadowFinder();

    // The beam centre of the geometry the rings are drawn in - the point of normal incidence, which
    // is where the beam lands only on an untilted detector. It starts as the file's, which is the
    // only one there is when the finder is built; a caller that has measured one replaces it before
    // reading the mask. A centre far from the truth draws the rings across the background's own radial
    // fall-off instead of along it, and the comparison then describes the fall-off rather than the
    // hardware. The projection is not centred on anything, so this may be set after the frames.
    void BeamCenter(float x, float y);

    // Accumulate one full converted-geometry image. Gap / masked pixels (the pixel type's sentinel
    // extreme) are skipped. `buffer` is scratch space for decompression, reused across the calls of
    // one worker.
    void AddImage(const DataMessage &data, std::vector<uint8_t> &buffer);

    // Compute the shadow mask (SHADOW, TRANSMITTING or 0 = keep), of the converted pixel count.
    // TRANSMITTING marks the pieces of hardware that let part of the beam through, added after the
    // shadow proper; both are masked, and a consumer that must not see those pieces can tell them apart.
    // Recomputed on each call from the projection, which is put together on the first read of it
    // (GetMask or GetMeanProjection): frames added after that are not seen.
    // nthreads = 0 asks for all hardware threads. The per-pixel passes over a 16M-pixel detector
    // dominate this, and they are all exactly parallel.
    [[nodiscard]] std::vector<uint32_t> GetMask(size_t nthreads = 0) const;

    // The polarization the last GetMask divided the background by: the stated degree, in the plane the
    // background shows where it shows one clearly, and as stated otherwise. The run's other corrections
    // take it from here, so that the plane is decided once. Unset with no stated degree, or before GetMask.
    [[nodiscard]] std::optional<float> GetAppliedPolarization() const;

    // Mean counts per pixel over the frames added, NAN where nothing was counted. This is the
    // projection GetMask() tests, so anything else that wants the background before indexing
    // gets it without reading the frames a second time.
    [[nodiscard]] std::vector<float> GetMeanProjection() const;

    // Let go of the projection the two above read, once the caller has what it wants of it: on a
    // 16 Mpx detector it is 360 MB. A later read puts it together again.
    void ReleaseProjection();

    [[nodiscard]] uint32_t GetFrameCount() const;
};