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

#pragma once

#include <cstdint>
#include <memory>
#include <mutex>
#include <vector>

#include "../indexing/CUDAMemHelpers.h"

// What a frame's levels and lit thresholds are made of (HotPixelFinder's constants), handed to the
// device so the two halves cannot drift apart.
struct HotPixelLevelRules {
    int min_sector_pixels;
    int min_ring_pixels;
    float lit_nsigma;
    float lit_offset;   // in the frames' units
    float gain;         // the frames' units per photon
};

// The device half of HotPixelFinder, for frames already preprocessed on the GPU. Everything a frame adds
// stays on the device: the per-frame order statistics (each ring-sector's median, each ring's median
// and median absolute deviation), the levels and thresholds made from them, and the per-pixel and
// per-key sums. Every statistic is an exact order statistic of integers, the threshold is computed
// with the same rounding steps as the host's (see HotPixelFinder::AddLevels) and every sum is an
// integer, so the sums are identical to what HotPixelFinder::AddImage produces. When the mask is read
// only what can decide it comes back: per-ring counts for the chance rate, then the few pixels that can
// be persistent or carry the error value.
class HotPixelFinderGPU {
    int device = 0;   // the card it was built on, where everything below lives
    const size_t npixels;
    const size_t nkeys;
    const int nrings;
    const int sectors;
    const HotPixelLevelRules rules;

    CudaDevicePtr<int32_t> key;               // ring * sectors + sector of each pixel, -1 masked
    CudaDevicePtr<uint32_t> pixels_by_key;    // the unmasked pixels, grouped by key, in pixel order
    CudaDevicePtr<uint32_t> key_begin;        // where each key's pixels start in pixels_by_key

    // The per-pixel and per-key sums, exactly those of HotPixelFinder.
    CudaDevicePtr<uint16_t> n_lit, n_error, n_error_ring_ok;
    CudaDevicePtr<int64_t> sum_value, error_level_sum;
    CudaDevicePtr<uint32_t> key_frames;
    CudaDevicePtr<int64_t> key_level_sum, key_local_level_sum;

    // Each pixel's sums are plain read-modify-writes, so the frames add to them one after another:
    // every accumulation waits on the device for the one before it (an event, not the host).
    std::mutex accumulate_mutex;
    cudaEvent_t last_accumulate = nullptr;

public:
    // One worker's buffers, on the stream its frames are preprocessed on.
    struct Frame {
        explicit Frame(std::shared_ptr<CudaStream> stream) : stream(std::move(stream)) {}
        // Its frames are queued, not waited for (Add), so the buffers below must not be freed before
        // the stream has used them.
        ~Frame() { if (stream) cudaStreamSynchronize(*stream); }
        Frame(Frame &&) = default;
        std::shared_ptr<CudaStream> stream;
        CudaDevicePtr<uint32_t> count;           // valid pixels per key
        CudaDevicePtr<int32_t> sector_median;    // per key
        CudaDevicePtr<int32_t> ring_median, ring_mad;
        CudaDevicePtr<int32_t> level;
        CudaDevicePtr<float> threshold;
        CudaDevicePtr<char> ring_ok;
    };

    // The keys of HotPixelFinder: one per pixel, and where each key's pixels start among the unmasked
    // pixels sorted by key.
    HotPixelFinderGPU(const int32_t *key, size_t npixels, const std::vector<uint32_t> &key_begin, int nrings,
                      int sectors, const HotPixelLevelRules &rules);
    ~HotPixelFinderGPU();
    HotPixelFinderGPU(const HotPixelFinderGPU &) = delete;
    HotPixelFinderGPU &operator=(const HotPixelFinderGPU &) = delete;

    [[nodiscard]] int GetDevice() const { return device; }

    // Add the sums another card accumulated to this one's. Every sum is an integer, so the result does
    // not depend on how the frames were split between the cards or on the order they are folded in.
    // Called with this finder's card current.
    void Fold(HotPixelFinderGPU &other);

    // Add one frame, preprocessed on frame.stream, as HotPixelFinder::AddImage does. Queued on that
    // stream; nothing is waited for on the host.
    void Add(const int32_t *device_image, Frame &frame);

    // Per ring, over the pixels lit on no more than half of their valid frames: the lit frames and the
    // valid frames, summed (HotPixelFinder::GetMask's chance rate). Waits for every frame added.
    void ChanceCounts(std::vector<int64_t> &lit, std::vector<int64_t> &seen);

    // The pixels that can be masked: those holding the error value on more than half of `frames`;
    // where spacing_ok, those lit on at least min(max(2, k_chance[ring]), valid frames) of at least
    // min_valid valid frames - a persistent pixel is lit on at least that many, because one reflection
    // explains at least two; and those with no error frame, at least min_valid valid frames and a sum
    // at most half of their sector's level sum, which a dead pixel has. For each, its index and
    // per-pixel sums; and every key's sums.
    struct Candidates {
        std::vector<uint32_t> index;
        std::vector<uint16_t> n_lit, n_error, n_error_ring_ok;
        std::vector<int64_t> sum_value, error_level_sum;
        std::vector<uint32_t> key_frames;
        std::vector<int64_t> key_level_sum, key_local_level_sum;
    };
    Candidates GetCandidates(uint32_t frames, int min_valid, bool spacing_ok, const std::vector<int> &k_chance);
};