// Copyright (2019-2023) Paul Scherrer Institute

#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;

    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 Clear() {
        std::unique_lock<std::mutex> 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<std::mutex> 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<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 {
       return mean;
    }

    [[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_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];
                }

                if (count_bin > 0)
                    ret[bin] = static_cast<float>(sum_bin / static_cast<double>(count_bin));
                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);
    }
};

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) {
            MultiLinePlot tmp = value->GetMeanPlot(bin_size);
            tmp.at(0).title = key;
            ret.emplace_back(tmp.at(0));
        }
        return ret;
    }
};

#endif //JUNGFRAUJOCH_STATUSVECTOR_H