// Copyright (2019-2024) Paul Scherrer Institute

#include "HDF5DataFilePluginMX.h"

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

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

    HDF5DataSpace data_space_spots({1, max_spots}, {H5S_UNLIMITED, max_spots});
    HDF5Dcpl dcpl_spots;
    dcpl_spots.SetChunking({1, max_spots});
    HDF5DataSet(data_file, "/entry/MX/peakXPosRaw", HDF5DataType(0.0f), data_space_spots, dcpl_spots);
    HDF5DataSet(data_file, "/entry/MX/peakYPosRaw", HDF5DataType(0.0f), data_space_spots, dcpl_spots);
    HDF5DataSet(data_file, "/entry/MX/peakTotalIntensity", HDF5DataType(0.0f), data_space_spots, dcpl_spots);
    HDF5DataSet(data_file, "/entry/MX/peakIndexed", HDF5DataType(1, false), data_space_spots, dcpl_spots);

    HDF5DataSpace data_space_linear({2}, {H5S_UNLIMITED});
    HDF5Dcpl dcpl_linear;
    dcpl_linear.SetChunking({1});
    HDF5DataSet(data_file, "/entry/MX/nPeaks", HDF5DataType(4, false), data_space_linear, dcpl_linear).SetExtent({1});
    HDF5DataSet(data_file, "/entry/MX/strongPixels", HDF5DataType(4, false), data_space_linear, dcpl_linear).SetExtent({1});
    HDF5DataSet(data_file, "/entry/MX/imageIndexed", HDF5DataType(1, false), data_space_linear, dcpl_linear).SetExtent({1});

    if (msg.resolution_estimation) {
        res_estimation = true;
        HDF5DataSet(data_file, "/entry/MX/resolutionEstimation", HDF5DataType(0.0f), data_space_linear,
                    dcpl_linear).SetExtent({1});
    }
    HDF5DataSpace data_space_lattice({1, 9}, {H5S_UNLIMITED, 9});
    HDF5Dcpl dcpl_lattice;
    dcpl_lattice.SetChunking({1, 9});
    HDF5DataSet(data_file, "/entry/MX/latticeIndexed", HDF5DataType(0.0f), data_space_lattice, dcpl_lattice);
}

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

    data_set_spot_x = std::make_unique<HDF5DataSet>(data_file, "/entry/MX/peakXPosRaw");
    data_set_spot_y = std::make_unique<HDF5DataSet>(data_file, "/entry/MX/peakYPosRaw");
    data_set_spot_int = std::make_unique<HDF5DataSet>(data_file, "/entry/MX/peakTotalIntensity");
    data_set_spot_indexed = std::make_unique<HDF5DataSet>(data_file, "/entry/MX/peakIndexed");
    data_set_spot_npeaks = std::make_unique<HDF5DataSet>(data_file, "/entry/MX/nPeaks");

    data_set_strong_pixel_count = std::make_unique<HDF5DataSet>(data_file, "/entry/MX/strongPixels");

    data_set_indexed = std::make_unique<HDF5DataSet>(data_file, "/entry/MX/imageIndexed");
    data_set_indexed_lattice = std::make_unique<HDF5DataSet>(data_file, "/entry/MX/latticeIndexed");
    if (res_estimation)
        data_set_resolution_estimation = std::make_unique<HDF5DataSet>(data_file, "/entry/MX/resolutionEstimation");
}

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

    size_t image_number = msg.number;

    if (image_number >= max_image_number) {
        max_image_number = image_number;
        data_set_spot_x->SetExtent({max_image_number+1, max_spots});
        data_set_spot_y->SetExtent({max_image_number+1, max_spots});
        data_set_spot_int->SetExtent({max_image_number+1, max_spots});
        data_set_spot_indexed->SetExtent({max_image_number+1, max_spots});
        data_set_spot_npeaks->SetExtent({max_image_number+1});
        data_set_strong_pixel_count->SetExtent({max_image_number + 1});

        data_set_indexed->SetExtent({max_image_number + 1});
        data_set_indexed_lattice->SetExtent({max_image_number + 1, 9});
        if (res_estimation)
            data_set_resolution_estimation->SetExtent({max_image_number + 1});
    }

    std::vector<float> spot_x(max_spots), spot_y(max_spots), spot_intensity(max_spots);
    std::vector<uint8_t> spot_indexed(max_spots);

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

    for (int i = 0; i < spot_cnt; i++) {
        spot_x[i] = msg.spots[i].x;
        spot_y[i] =  msg.spots[i].y;
        spot_intensity[i] = msg.spots[i].intensity;
        spot_indexed[i] = msg.spots[i].indexed;
    }
    data_set_spot_x->WriteVec(spot_x, {image_number, 0}, {1, max_spots});
    data_set_spot_y->WriteVec(spot_y, {image_number, 0}, {1, max_spots});
    data_set_spot_int->WriteVec(spot_intensity, {image_number, 0}, {1, max_spots});
    data_set_spot_indexed->WriteVec(spot_indexed, {image_number, 0}, {1, max_spots});
    data_set_spot_npeaks->WriteScalar(spot_cnt, {image_number});
    data_set_strong_pixel_count->WriteScalar(static_cast<uint32_t>(msg.strong_pixel_count), {image_number});

    data_set_indexed->WriteScalar((uint8_t) msg.indexing_result, {image_number});
    if (msg.indexing_lattice.size() == 9)
        data_set_indexed_lattice->WriteVec(msg.indexing_lattice, {image_number, 0}, {1, 9});

    if (res_estimation)
        data_set_resolution_estimation->WriteScalar(msg.resolution_estimation.value(), {image_number});
}

void HDF5DataFilePluginMX::Flush() {
    if (max_spots == 0)
        return;

    data_set_spot_x->Flush();
    data_set_spot_y->Flush();
    data_set_spot_int->Flush();
    data_set_spot_indexed->Flush();
    data_set_spot_npeaks->Flush();
    data_set_strong_pixel_count->Flush();
    data_set_indexed->Flush();
    data_set_indexed_lattice->Flush();
    if (res_estimation)
        data_set_resolution_estimation->Flush();
}

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

}