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

#include "HDF5DataFilePluginMX.h"
#define RESERVE_IMAGES 1000

HDF5DataFilePluginMX::HDF5DataFilePluginMX(size_t in_max_spots) : max_spots(in_max_spots) {}

void HDF5DataFilePluginMX::OpenFile(HDF5File &data_file, const DataMessage &msg) {
    if (max_spots == 0)
        return;

    spot_x.reserve(max_spots * RESERVE_IMAGES);
    spot_y.reserve(max_spots * RESERVE_IMAGES);
    spot_int.reserve(max_spots * RESERVE_IMAGES);
    npeaks.reserve(RESERVE_IMAGES);
    strong_pixel_count.reserve(RESERVE_IMAGES);
    indexed.reserve(RESERVE_IMAGES);
    indexed_lattice.reserve(9 * RESERVE_IMAGES);
    spot_count_rings.reserve(RESERVE_IMAGES);
}

void HDF5DataFilePluginMX::Write(const DataMessage &msg, uint64_t image_number) {
    if (max_spots == 0)
        return;

    if (image_number >= max_image_number) {
        max_image_number = image_number;

        spot_x.resize(max_spots * (max_image_number + 1));
        spot_y.resize(max_spots * (max_image_number + 1));
        spot_int.resize(max_spots * (max_image_number + 1));

        npeaks.resize(max_image_number + 1);
        strong_pixel_count.resize(max_image_number + 1);
        indexed.resize(max_image_number + 1);
        indexed_lattice.resize((max_image_number + 1) * 9);

        if (msg.spot_count_in_rings)
            spot_count_rings.resize(max_image_number + 1);
    }

    uint32_t spot_cnt = std::min(msg.spots.size(), max_spots);

    for (int i = 0; i < spot_cnt; i++) {
        spot_x[max_spots * image_number + i] = msg.spots[i].x;
        spot_y[max_spots * image_number + i] =  msg.spots[i].y;
        spot_int[max_spots * image_number + i] = msg.spots[i].intensity;
    }

    npeaks[image_number] = spot_cnt;
    strong_pixel_count[image_number] = msg.strong_pixel_count;
    indexed[image_number] = msg.indexing_result;
    if (msg.spot_count_in_rings)
        spot_count_rings[image_number] = msg.spot_count_in_rings.value();

    if (msg.indexing_lattice.size() == 9) {
        for (int i = 0; i < 9; i++)
            indexed_lattice[image_number * 9 + i] = msg.indexing_lattice[i];
    }
}

void HDF5DataFilePluginMX::WriteFinal(HDF5File &data_file) {
    HDF5Group(data_file, "/entry/MX").NXClass("NXcollection");

    if (!spot_x.empty()) {
        data_file.SaveVector("/entry/MX/peakXPosRaw", spot_x, {(hsize_t) (max_image_number + 1), max_spots});
        data_file.SaveVector("/entry/MX/peakYPosRaw", spot_y, {(hsize_t) (max_image_number + 1), max_spots});
        data_file.SaveVector("/entry/MX/peakTotalIntensity", spot_int, {(hsize_t) (max_image_number + 1), max_spots});

        data_file.SaveVector("/entry/MX/nPeaks", npeaks);
        data_file.SaveVector("/entry/MX/strongPixels", strong_pixel_count);
    }

    if (!spot_count_rings.empty())
        data_file.SaveVector("/entry/MX/nPeaksRingFiltered", spot_count_rings);

    if (!indexed.empty())
        data_file.SaveVector("/entry/MX/imageIndexed", indexed);
    if (!indexed_lattice.empty())
        data_file.SaveVector("/entry/MX/latticeIndexed", indexed_lattice, {(hsize_t) (max_image_number + 1), 9});
}