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

#ifndef JUNGFRAUJOCH_HISTOGRAM_H
#define JUNGFRAUJOCH_HISTOGRAM_H

#include <cstdint>
#include <cstddef>
#include <vector>
#include <mutex>

#include "MultiLinePlot.h"
#include "../common/JFJochException.h"

template<class T>
class SetAverage {
    std::vector<T> sum;
    std::vector<uint64_t> count;
    mutable std::mutex m;

public:
    explicit SetAverage(size_t bins) : sum(bins), count(bins) {
    }

    void Add(size_t bin, T val) {
        std::unique_lock ul(m);

        if (bin < sum.size()) {
            sum[bin] += val;
            count[bin] += 1;
        }
    }

    MultiLinePlot GetPlot() const {
        std::unique_lock ul(m);

        MultiLinePlotStruct plot;
        plot.x.resize(sum.size());
        plot.y.resize(sum.size());
        for (int i = 0; i < sum.size(); i++) {
            plot.x[i] = static_cast<float>(i);
            if (count[i] > 0)
                plot.y[i] = static_cast<float>(sum[i]) / count[i];
            else
                plot.y[i] = 0;
        }
        MultiLinePlot ret;
        ret.AddPlot(plot);
        return ret;
    }
};

class Histogram {
    std::vector<uint64_t> count;
    int32_t total_count = 0;
    int32_t max_bin = 0;

public:
    explicit Histogram(size_t bins) : count(bins) {}

    void Add(int32_t val) {
        if (val > 0 && val < count.size()) {
            count[val] += 1;
            if (val > max_bin)
                max_bin = val;
            ++total_count;
        }
    }

    void clear() {
        for (auto &c: count) c = 0;
        total_count = 0;
        max_bin = 0;
    }

    [[nodiscard]] std::vector<float> GetCount() const {
        std::vector<float> ret;
        ret.reserve(max_bin + 1);
        for (size_t i = 0; i < max_bin + 1; i++)
            ret.emplace_back(count[i]);
        return ret;
    }

    [[nodiscard]] uint64_t GetTotalCount() const {
        return total_count;
    }

    // Returns the value x such that approximately `percent`% of samples are <= x.
    // - percent must be in [0, 100]
    // - returns std::nullopt if histogram is empty
    [[nodiscard]] std::optional<int32_t> Percentile(float percent) const {
        if (!std::isfinite(percent) || percent < 0.0f || percent > 100.0f) {
            throw JFJochException(JFJochExceptionCategory::InputParameterBelowMin,
                                  "FloatHistogram Percentile expects percent in [0, 100]");
        }

        if (total_count == 0)
            return std::nullopt;

        // Target rank in [0, total-1]
        const double q = static_cast<double>(percent) / 100.0;
        const auto target = static_cast<int64_t>(std::floor(q * static_cast<double>(total_count - 1)));

        uint64_t cumulative = 0;
        for (int64_t i = 0; i < max_bin + 1; i++) {
            cumulative += count[i];
            if (target < cumulative && count[i] > 0)
                    return i;
        }

        // If due to rounding we didn't return inside the loop, clamp to the last bin's upper edge.
        return count.size() - 1;
    }
};

#endif //JUNGFRAUJOCH_HISTOGRAM_H