Jungfraujoch/common/StatusVector.h

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

#ifndef JUNGFRAUJOCH_STATUSVECTOR_H
#define JUNGFRAUJOCH_STATUSVECTOR_H

#include <vector>
#include <map>
#include <mutex>
#include <cmath>
#include <optional>

#include "Plot.h"

#include "JFJochException.h"
#include <memory>

template <class T> class StatusVector {
    mutable std::mutex m;
    std::vector<T> content;
    std::vector<uint8_t> present;

    float mean = NAN;
    size_t count = 0;
    double sum = 0;
public:
    void Clear() {
        std::unique_lock ul(m);
        content.clear();
        present.clear();
    }

    void AddElement(uint32_t id, std::optional<T> val) {
        if (val.has_value())
            AddElement(id, val.value());
    }

    void AddElement(uint32_t id, T val) {
        std::unique_lock ul(m);
        if (id >= content.size()) {
            content.resize(id + 1);
            present.resize(id + 1);
        }
        content[id] = val;
        present[id] = 1;

        sum += val;
        count += 1;
        mean = sum / count;
    }

    [[nodiscard]] std::vector<T> ExportArray(T def_value) const {
        std::unique_lock ul(m);
        std::vector<T> ret(content.size(), def_value);
        for (int i = 0; i < content.size(); i++) {
            if (present[i])
                ret[i] = content[i];
        }
        return ret;
    }

    [[nodiscard]] float Mean() const {
       return mean;
    }

    MultiLinePlotStruct GetMeanPerBin(int32_t bin_size) const {
        std::unique_lock ul(m);

        MultiLinePlotStruct ret;

        if (bin_size <= 0)
            throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                                  "Bin number must be greater than zero");

        if (!content.empty()) {
            size_t elems;

            if (content.size() < bin_size) {
                // don't bin if less samples than bin size
                bin_size = 1;
                elems = content.size();
            } else
                elems = content.size() / bin_size + ((content.size() % bin_size > 0) ? 1 : 0);

            ret.x.reserve(elems);
            ret.y.reserve(elems);

            if (bin_size == 1) {
                for (int i = 0; i < content.size(); i++) {
                    if (present[i]) {
                        ret.x.push_back(i);
                        ret.y.push_back(content[i]);
                    }
                }
            } else {
                for (int bin = 0; bin < elems; bin++) {
                    double sum_bin = 0;
                    int64_t count_bin = 0;

                    for (int i = bin * bin_size; (i < (bin + 1) * bin_size) && (i < content.size()); i++) {
                        sum_bin += present[i] * content[i];
                        count_bin += present[i];
                    }

                    float bin_x = static_cast<float>(bin_size) * (bin + 0.5f);

                    if (count_bin > 0) {
                        ret.x.push_back(bin_x);
                        ret.y.push_back(static_cast<float>(sum_bin / static_cast<double>(count_bin)));
                    }
                }
            }
        }
        return ret;
    }

    [[nodiscard]] MultiLinePlotStruct GetMaxPerBin(int32_t bin_size) const {
        std::unique_lock ul(m);

        MultiLinePlotStruct ret;

        if (bin_size <= 0)
            throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                                  "Bin number must be greater than zero");

        size_t elems = content.size() / bin_size + ((content.size() % bin_size > 0) ? 1 : 0);

        if (!content.empty()) {
            ret.x.reserve(elems);
            ret.y.reserve(elems);

            if (bin_size == 1) {
                for (int i = 0; i < content.size(); i++) {
                    if (present[i]) {
                        ret.x.push_back(i);
                        ret.y.push_back(content[i]);
                    }
                }
            } else {
                for (int bin = 0; bin < elems; bin++) {
                    float max_in_bin = 0;
                    bool max_bin_set = false;

                    for (int i = bin * bin_size; (i < (bin + 1) * bin_size) && (i < content.size()); i++) {
                        if (present[i] && (content[i] > max_in_bin || !max_bin_set)) {
                            max_bin_set = true;
                            max_in_bin = content[i];
                        }
                    }

                    float bin_x = static_cast<float>(bin_size) * (bin + 0.5f);

                    if (max_bin_set) {
                        ret.x.push_back(bin_x);
                        ret.y.push_back(max_in_bin);
                    }
                }
            }
        }
        return ret;
    }

    MultiLinePlot GetMeanPlot(int64_t bin_size) const {
        MultiLinePlot ret(1);
        ret[0] = GetMeanPerBin(bin_size);
        return ret;
    }

    MultiLinePlot GetMaxPlot(int64_t bin_size) const {
        MultiLinePlot ret(1);
        ret[0] = GetMaxPerBin(bin_size);
        return ret;
    }
};

template <class T> class StatusMultiVector {
    std::mutex m;
    std::map<std::string, std::unique_ptr<StatusVector<T>>> status;
public:
    void Clear() {
        std::unique_lock ul(m);
        status.clear();
    }

    void AddElement(const std::string& s, uint32_t id, T val) {
        std::unique_lock ul(m);
        if (!status.contains(s))
            status[s] = std::make_unique<StatusVector<T>>();
        status[s]->AddElement(id, val);
    }

    void AddElement(const std::string& s, uint32_t id, std::optional<T> val) {
        // no need to lock, as AddElement(string, u32, T) has lock already
        if (val.has_value())
            AddElement(s, id, val.value());
    }

    [[nodiscard]] MultiLinePlot GetMeanPlot(int64_t bin_size) const {
        MultiLinePlot ret;
        for (const auto &[key, value]: status) {
            auto tmp = value->GetMeanPerBin(bin_size);
            tmp.title = key;
            ret.emplace_back(tmp);
        }
        return ret;
    }
};

#endif //JUNGFRAUJOCH_STATUSVECTOR_H