// Copyright (2019-2023) Paul Scherrer Institute

#ifndef JUNGFRAUJOCH_STATUSVECTOR_H
#define JUNGFRAUJOCH_STATUSVECTOR_H

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

#include "Plot.h"

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

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

    MultiLinePlot GeneratePlot(const std::vector<float> &status, int32_t bin_size) const {
        MultiLinePlot ret(1);
        if (status.size() == 1) {
            ret[0].x = {(content.size() - 1) / 2.0f};
            ret[0].y = {status[0]};
        } else if (!status.empty()) {
            ret[0].y = status;
            ret[0].x.resize(status.size());
            if (bin_size == 1) {
                for (int i = 0; i < status.size(); i++)
                    ret[0].x[i] = i;
            } else {
                for (int i = 0; i < status.size(); i++)
                    ret[0].x[i] = bin_size * (i + 0.5f);
            }
        }
        return ret;
    }
public:
    void AddElement(uint32_t id, T val) {
        std::unique_lock<std::mutex> ul(m);
        if (id >= content.size()) {
            content.resize(id + 1);
            present.resize(id + 1);
        }
        content[id] = val;
        present[id] = 1;
    }

    [[nodiscard]] std::vector<T> ExportArray(T def_value) const {
        std::unique_lock<std::mutex> 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 {
        std::unique_lock<std::mutex> ul(m);
        size_t count = 0;
        double sum = 0;
        for (int i = 0; i < content.size(); i++) {
            sum += present[i] * content[i];
            count += present[i];
        }

        if (count > 0)
            return static_cast<float>(sum / count);
        else
            return NAN;
    }

    [[nodiscard]] std::vector<float> GetMeanPerBin(int32_t bin_size) const {
        std::unique_lock<std::mutex> ul(m);

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

        std::vector<float> ret;

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

        if (!content.empty()) {
            for (int bin = 0; bin < elems; bin++) {
                double sum = 0;
                int64_t count = 0;

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

                if (count > 0)
                    ret[bin] = static_cast<float>(sum / static_cast<double>(count));
                else
                    ret[bin] = 0.0;
            }
        }
        return ret;
    }

    [[nodiscard]] std::vector<float> GetMaxPerBin(int32_t bin_size, float default_val = 0.0) const {
        std::unique_lock<std::mutex> ul(m);

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

        std::vector<float> ret;

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

        if (!content.empty()) {
            for (int bin = 0; bin < elems; bin++) {
                ret[bin] = default_val;

                for (int i = bin * bin_size; (i < (bin + 1) * bin_size) && (i < content.size()); i++) {
                    if (present[i] && (content[i] > ret[bin]))
                        ret[bin] = content[i];
                }
            }
        }
        return ret;
    }

    MultiLinePlot GetMeanPlot(int64_t bin_size) const {
        // GetStatus has mutex, no need to lock again
        return GeneratePlot(GetMeanPerBin(bin_size), bin_size);
    }

    MultiLinePlot GetMaxPlot(int64_t bin_size) const {
        return GeneratePlot(GetMaxPerBin(bin_size), bin_size);
    }
};


#endif //JUNGFRAUJOCH_STATUSVECTOR_H