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


#include "StatusVector.h"

void StatusVector::Clear() {
    std::unique_lock ul(m);
    content.clear();

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

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

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

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

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

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

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

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

[[nodiscard]] std::vector<float> StatusVector::ExportArray() const {
    std::unique_lock ul(m);
    return content;
}

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

MultiLinePlotStruct StatusVector::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 (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 (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;
}

MultiLinePlotStruct StatusVector::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 (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 (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 StatusVector::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 StatusVector::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;
}

void StatusMultiVector::Clear() {
    std::unique_lock ul(m);
    status.clear();
}

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

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

MultiLinePlotStruct StatusMultiVector::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;
}

MultiLinePlot StatusMultiVector::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;
}

std::vector<float> StatusMultiVector::ExportArray(const std::string &s) const {
    auto iter = status.find(s);
    if (iter == status.end() || !iter->second)
        return {};
    return iter->second->ExportArray();
}