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

#include "HDF5DataFilePluginAzInt.h"

HDF5DataFilePluginAzInt::HDF5DataFilePluginAzInt(const StartMessage &message) :
az_int_bin_to_q(message.az_int_bin_to_q),
az_int_bin_to_two_theta(message.az_int_bin_to_two_theta),
az_int_bin_to_phi(message.az_int_bin_to_phi),
azimuthal_bins(message.az_int_phi_bin_count.value_or(1)),
q_bins(message.az_int_q_bin_count.value_or(1)){
}

void HDF5DataFilePluginAzInt::OpenFile(HDF5File &data_file, const DataMessage &msg, size_t images_per_file) {
    if (az_int_bin_to_q.empty() || q_bins <= 0 || azimuthal_bins <= 0)
        return;

    std::vector<hsize_t> dim = {azimuthal_bins, q_bins};

    HDF5Group(data_file, "/entry/azint").NXClass("NXcollection");
    data_file.SaveVector("/entry/azint/bin_to_q", az_int_bin_to_q, dim);
    if (!az_int_bin_to_two_theta.empty())
        data_file.SaveVector("/entry/azint/bin_to_two_theta", az_int_bin_to_two_theta, dim);
    if (!az_int_bin_to_phi.empty())
        data_file.SaveVector("/entry/azint/bin_to_phi", az_int_bin_to_phi, dim);

    az_int_image.reserve(images_per_file * azimuthal_bins * q_bins);
    if (!msg.az_int_profile_count.empty()) {
        save_profile_count = true;
        az_int_image_count.reserve(images_per_file * azimuthal_bins * q_bins);
    }
    if (!msg.az_int_profile_std.empty()) {
        save_profile_std = true;
        az_int_image_std.reserve(images_per_file * azimuthal_bins * q_bins);
    }
}

void HDF5DataFilePluginAzInt::Write(const DataMessage &msg, uint64_t image_number) {
    if (az_int_bin_to_q.empty() || q_bins <= 0 || azimuthal_bins <= 0)
        return;

    if (image_number >= max_image_number || (max_image_number == 0)) {
        max_image_number = image_number;
        az_int_image.resize((max_image_number + 1) * azimuthal_bins * q_bins, NAN);
        if (save_profile_count)
            az_int_image_count.resize((max_image_number + 1) * azimuthal_bins * q_bins, 0);
        if (save_profile_std)
            az_int_image_std.resize((max_image_number + 1) * azimuthal_bins * q_bins, NAN);
    }

    if (!msg.az_int_profile.empty() && (msg.az_int_profile.size() == azimuthal_bins * q_bins)) {
        for (int i = 0; i < azimuthal_bins * q_bins; i++)
            az_int_image[image_number * azimuthal_bins * q_bins + i] = msg.az_int_profile[i];
    }

    if (save_profile_count
        && !msg.az_int_profile_count.empty()
        && (msg.az_int_profile_count.size() == azimuthal_bins * q_bins)) {
        for (int i = 0; i < azimuthal_bins * q_bins; i++)
            az_int_image_count[image_number * azimuthal_bins * q_bins + i] = msg.az_int_profile_count[i];
    }

    if (save_profile_std
        && !msg.az_int_profile_std.empty()
        && (msg.az_int_profile_std.size() == azimuthal_bins * q_bins)) {
        for (int i = 0; i < azimuthal_bins * q_bins; i++)
            az_int_image_std[image_number * azimuthal_bins * q_bins + i] = msg.az_int_profile_std[i];
    }
}

void HDF5DataFilePluginAzInt::WriteFinal(HDF5File &data_file) {
    if (!az_int_image.empty())
        data_file.SaveVector("/entry/azint/image", az_int_image,
                             {(hsize_t) (max_image_number + 1), azimuthal_bins, q_bins});
    if (!az_int_image_std.empty())
        data_file.SaveVector("/entry/azint/image_std", az_int_image_std,
            {(hsize_t) (max_image_number + 1), azimuthal_bins, q_bins});
    if (!az_int_image_count.empty())
        data_file.SaveVector("/entry/azint/image_count", az_int_image_count,
            {(hsize_t) (max_image_number + 1), azimuthal_bins, q_bins});
}