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

#ifndef JFJOCH_TOPPIXELS_H
#define JFJOCH_TOPPIXELS_H

#include <array>
#include <algorithm>
#include <cstdint>
#include <stdexcept>
#include <vector>

class TopPixels {
public:
    struct Entry {
        int32_t value = 0;
        int32_t index = 0;
    };

    explicit TopPixels(int capacity = 20)
        : capacity_(std::clamp(capacity, 1, kMaxCapacity)) {}

    void Clear() noexcept { size_ = 0; }

    void Add(int32_t value, int32_t index) noexcept {
        const Entry e{value, index};

        if (size_ < capacity_) {
            InsertSorted_(e);
            return;
        }

        // Quick reject: if not better than the current smallest in top-K
        if (value <= top_[capacity_ - 1].value)
            return;

        ReplaceSmallestAndFixOrder_(e);
    }

    [[nodiscard]] int Capacity() const noexcept { return capacity_; }
    [[nodiscard]] int Size() const noexcept { return size_; }
    [[nodiscard]] bool Empty() const noexcept { return size_ == 0; }

    // Sorted descending by value; valid for i in [0, Size()).
    [[nodiscard]] const Entry& At(int i) const {
        if (i < 0 || i >= size_)
            throw std::out_of_range("TopPixels::At index out of range");
        return top_[i];
    }

    [[nodiscard]] const Entry& operator[](int i) const { return At(i); }

    // Convenience: copy results to a std::vector (still sorted descending).
    [[nodiscard]] std::vector<Entry> ToVector() const {
        return std::vector<Entry>(top_.begin(), top_.begin() + size_);
    }

private:
    static constexpr int kMaxCapacity = 64; // hard cap for stack-friendly storage

    std::array<Entry, kMaxCapacity> top_{};
    int size_ = 0;
    int capacity_ = 20;

    void InsertSorted_(Entry e) noexcept {
        top_[size_] = e;
        ++size_;

        // Bubble up to keep descending order by value
        for (int i = size_ - 1; i > 0 && top_[i].value > top_[i - 1].value; --i)
            std::swap(top_[i], top_[i - 1]);
    }

    void ReplaceSmallestAndFixOrder_(Entry e) noexcept {
        // Replace the smallest (last, because sorted descending)
        top_[capacity_ - 1] = e;

        // Bubble up to restore descending order
        for (int i = capacity_ - 1; i > 0 && top_[i].value > top_[i - 1].value; --i)
            std::swap(top_[i], top_[i - 1]);

        size_ = capacity_;
    }
};

#endif //JFJOCH_TOPPIXELS_H