// 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 <memory>

#include "JFJochException.h"
#include "MultiLinePlot.h"

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();

        mean = NAN;
        count = 0;
        sum = 0;
    }

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

    std::optional<T> GetElement(uint32_t id) const {
        std::unique_lock ul(m);
        if (id < content.size() && (present.at(id) > 0))
            return content.at(id);
        return {};
    }

    void ImportArray(std::vector<T> &in_content,
                     std::vector<uint8_t> &in_present) {
        std::unique_lock ul(m);
        if (in_content.size() != in_present.size())
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Content and presence arrays must be equal in size");

        content = in_content;
        present = in_present;

        sum = 0;
        count = 0;

        for (int i = 0; i < content.size(); i++) {
            if (present[i]) {
                sum += content[i];
                count += 1;
            }
        }

        if (count > 0)
            mean = sum / count;
        else
            mean = 0;
    }

    size_t GetImageNumber() const {
        return content.size();
    }

    bool empty() const {
        return count == 0;
    }

    int32_t GetActualBinning(int32_t bin_size) const {
        if (content.size() < bin_size)
            return 1;
        return bin_size;
    }


    void ImportArray(std::vector<T> &in_content) {
        std::vector<uint8_t> in_present(in_content.size(), 1);
        ImportArray(in_content, in_present);
    }

    [[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, float x_start, float x_incr,
                                      const std::optional<float> &fill_value = {}) 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] && std::isfinite(content[i])) {
                        ret.x.push_back(x_start + x_incr * i);
                        ret.y.push_back(content[i]);
                    } else if (fill_value) {
                        ret.x.push_back(x_start + x_incr * i);
                        ret.y.push_back(fill_value.value());
                    }
                }
            } 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++) {
                        if (present[i] && std::isfinite(content[i])) {
                            sum_bin += 1.0 * content[i];
                            count_bin += 1.0;
                        }
                    }

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

                    if (count_bin > 0) {
                        ret.x.push_back(x_start + x_incr * bin_x);
                        ret.y.push_back(static_cast<float>(sum_bin / static_cast<double>(count_bin)));
                    } else if (fill_value) {
                        ret.x.push_back(x_start + x_incr * bin_x);
                        ret.y.push_back(fill_value.value());
                    }
                }
            }
        }
        return ret;
    }

    [[nodiscard]] MultiLinePlotStruct GetMaxPerBin(int32_t bin_size, float x_start, float x_incr,
                                                   const std::optional<float> &fill_value = {}) 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] && std::isfinite(content[i])) {
                        ret.x.push_back(x_start + x_incr * i);
                        ret.y.push_back(content[i]);
                    } else if (fill_value) {
                        ret.x.push_back(x_start + x_incr * i);
                        ret.y.push_back(fill_value.value());
                    }
                }
            } 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] && std::isfinite(content[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(x_start + x_incr * bin_x);
                        ret.y.push_back(max_in_bin);
                    } else if (fill_value) {
                        ret.x.push_back(x_start + x_incr * bin_x);
                        ret.y.push_back(fill_value.value());
                    }
                }
            }
        }
        return ret;
    }

    MultiLinePlot GetMeanPlot(int64_t bin_size, float x_start, float x_incr,
                              const std::optional<float> &fill_value = {}) const {
        MultiLinePlot ret;
        ret.AddPlot(GetMeanPerBin(bin_size, x_start, x_incr, fill_value));
        return ret;
    }

    MultiLinePlot GetMaxPlot(int64_t bin_size, float x_start, float x_incr,
                             const std::optional<float> &fill_value = {}) const {
        MultiLinePlot ret;
        ret.AddPlot(GetMaxPerBin(bin_size, x_start, x_incr, fill_value));
        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]] MultiLinePlotStruct GetMeanPerBin(const std::string& in_key, int64_t bin_size, float x_start, float x_incr,
                                                    const std::optional<float> &fill_value = {}) const {
        MultiLinePlotStruct ret{};
        for (const auto &[key, value]: status) {
            if (key == in_key) {
                ret = value->GetMeanPerBin(bin_size, x_start, x_incr, fill_value);
                ret.title = key;
            }
        }
        return ret;
    }

    [[nodiscard]] MultiLinePlot GetMeanPlot(int64_t bin_size, float x_start, float x_incr,
                                            const std::optional<float> &fill_value = {}) const {
        MultiLinePlot ret;
        for (const auto &[key, value]: status) {
            auto tmp = value->GetMeanPerBin(bin_size, x_start, x_incr, fill_value);
            tmp.title = key;
            ret.AddPlot(tmp);
        }
        return ret;
    }

    [[nodiscard]] std::vector<T> ExportArray(const std::string& s, T def_value) const {
        auto iter = status.find(s);
        if (iter == status.end() || !iter->second)
            return {};
        return iter->second->ExportArray(def_value);
    }
};

#endif //JUNGFRAUJOCH_STATUSVECTOR_H