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

#include "HDF5DataFilePluginROI.h"

#define RESERVE_IMAGES 1000

void HDF5DataFilePluginROI::OpenFile(HDF5File &data_file, const DataMessage &msg) {}

void HDF5DataFilePluginROI::Write(const DataMessage &msg, uint64_t image_number) {
    for (const auto &r: msg.roi) {
        if (roi_data.contains(r.first)) {
            roi_data[r.first].max.reserve(RESERVE_IMAGES);
            roi_data[r.first].sum.reserve(RESERVE_IMAGES);
            roi_data[r.first].sum_sq.reserve(RESERVE_IMAGES);
            roi_data[r.first].npixel.reserve(RESERVE_IMAGES);
            roi_data[r.first].x.reserve(RESERVE_IMAGES);
            roi_data[r.first].y.reserve(RESERVE_IMAGES);
        }
        roi_data[r.first].max[image_number] = r.second.max_count;
        roi_data[r.first].sum[image_number] = r.second.sum;
        roi_data[r.first].npixel[image_number] = r.second.pixels;
        roi_data[r.first].sum_sq[image_number] = r.second.sum_square;

        roi_data[r.first].x[image_number] = static_cast<double>(r.second.x_weighted) / static_cast<double>(r.second.sum);
        roi_data[r.first].y[image_number] = static_cast<double>(r.second.y_weighted) / static_cast<double>(r.second.sum);
    }
}

void HDF5DataFilePluginROI::WriteFinal(HDF5File &data_file) {
    if (roi_data.empty())
        return;

    HDF5Group(data_file, "/entry/roi").NXClass("NXcollection");

    for (const auto &r: roi_data) {
        HDF5Group group_roi(data_file, "/entry/roi/" + r.first);
        group_roi.SaveVector("max", r.second.max.vec());
        group_roi.SaveVector("sum", r.second.sum.vec());
        group_roi.SaveVector("sum_sq", r.second.sum_sq.vec());
        group_roi.SaveVector("npixel", r.second.npixel.vec());
        group_roi.SaveVector("x", r.second.x.vec());
        group_roi.SaveVector("y", r.second.y.vec());
    }
}