// Copyright (2019-2023) Paul Scherrer Institute

#include "HDF5DataFile.h"
#include "../compression/JFJochCompressor.h"

HDF5DataFile::HDF5DataFile(const std::string &in_filename, int64_t width, int64_t height,
                           int64_t pixel_depth_byte, bool is_signed, const std::vector<float>& in_rad_int_bin_to_q,
                           CompressionAlgorithm compression,
                           size_t in_max_spots) : data_type(pixel_depth_byte, is_signed),
                           max_spots(in_max_spots), rad_int_bin_to_q(in_rad_int_bin_to_q) {
    max_image_number = 0;
    nimages = 0;

    if (max_spots > 256)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Max spot count must be less than 256 for performance reasons");

    filename = in_filename;

    xpixel = width;
    ypixel = height;

    dcpl.SetCompression(compression, pixel_depth_byte, JFJochBitShuffleCompressor::DefaultBlockSize);
    dcpl.SetChunking( {1, ypixel, xpixel});
    if (is_signed) {
        if (pixel_depth_byte == 2)
            dcpl.SetFillValue16(INT16_MIN);
        else
            dcpl.SetFillValue32(INT32_MIN);
    }
}

HDF5DataFile::~HDF5DataFile() {
    if (data_file) {
        HDF5Group group_exp(*data_file, "/entry/jungfrau");
        group_exp.NXClass("NXcollection");

        group_exp.SaveVector("info", jf_info);
        group_exp.SaveVector("timestamp", timestamp);
        group_exp.SaveVector("storage_cell", storage_cell);
        group_exp.SaveVector("exptime", exptime);
        group_exp.SaveVector("receiverAqDevDelay", receiver_aq_dev_delay);

        data_file->SaveVector("/entry/result/strongPixelCount", strong_pixel_count);
        if (!resolution_estimation.empty())
            data_file->SaveVector("/entry/result/resolutionEstimation", resolution_estimation);
    }
}

void HDF5DataFile::CreateFile() {
    SetupSWMRFile();

    data_file = std::make_unique<HDF5File>(filename, false, true);
    data_set = std::make_unique<HDF5DataSet>(*data_file, "/entry/data/data");
    data_set_spot_x = std::make_unique<HDF5DataSet>(*data_file, "/entry/result/peakXPosRaw");
    data_set_spot_y = std::make_unique<HDF5DataSet>(*data_file, "/entry/result/peakYPosRaw");
    data_set_spot_int = std::make_unique<HDF5DataSet>(*data_file, "/entry/result/peakTotalIntensity");
    data_set_spot_indexed = std::make_unique<HDF5DataSet>(*data_file, "/entry/result/peakIndexed");
    data_set_spot_npeaks = std::make_unique<HDF5DataSet>(*data_file, "/entry/result/nPeaks");
    data_set_xfel_event_code = std::make_unique<HDF5DataSet>(*data_file, "/entry/xfel/eventCode");
    data_set_xfel_pulseid = std::make_unique<HDF5DataSet>(*data_file, "/entry/xfel/pulseID");
    data_set_indexed = std::make_unique<HDF5DataSet>(*data_file, "/entry/result/imageIndexed");
    data_set_indexed_lattice = std::make_unique<HDF5DataSet>(*data_file, "/entry/result/latticeIndexed");
    if (!rad_int_bin_to_q.empty())
        data_set_az_int = std::make_unique<HDF5DataSet>(*data_file, "/entry/result/azimIntegration");
}

void HDF5DataFile::Write(const DataMessage &msg, uint64_t image_number) {
    std::lock_guard<std::mutex> lock(hdf5_mutex);
    bool new_file = false;

    if (!data_file) {
        CreateFile();
        new_file = true;
    }

    if (new_file || (static_cast<int64_t>(image_number) > max_image_number)) {
        max_image_number = image_number;
        data_set->SetExtent({max_image_number+1, ypixel, xpixel});
        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_xfel_pulseid->SetExtent({max_image_number+1});
        data_set_xfel_event_code->SetExtent({max_image_number+1});
        data_set_indexed->SetExtent({max_image_number + 1});
        data_set_indexed_lattice->SetExtent({max_image_number + 1, 9});
        if (!rad_int_bin_to_q.empty())
            data_set_az_int->SetExtent({max_image_number + 1, rad_int_bin_to_q.size()});

        strong_pixel_count.resize(max_image_number + 1);
        jf_info.resize(max_image_number + 1);
        receiver_aq_dev_delay.resize(max_image_number + 1);
        timestamp.resize(max_image_number + 1);
        exptime.resize(max_image_number + 1);
        storage_cell.resize(max_image_number + 1);

        if (msg.resolution_estimation)
            resolution_estimation.resize(max_image_number + 1);
    }

    nimages++;
    data_set->WriteDirectChunk(msg.image.data, msg.image.size, {image_number, 0, 0});

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

    std::vector<float> spot_x(max_spots), spot_y(max_spots), spot_intensity(max_spots);
    std::vector<uint8_t> spot_indexed(max_spots);
    for (int i = 0; i < 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(cnt, {image_number});

    data_set_xfel_event_code->WriteScalar(msg.xfel_event_code, {image_number});
    data_set_xfel_pulseid->WriteScalar(msg.xfel_bunch_id, {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});

    strong_pixel_count[image_number] = msg.strong_pixel_count;
    jf_info[image_number] = msg.jf_info;
    receiver_aq_dev_delay[image_number] = msg.receiver_aq_dev_delay;
    timestamp[image_number] = msg.timestamp;
    storage_cell[image_number] = msg.storage_cell;
    exptime[image_number] = msg.exptime;

    if (msg.resolution_estimation)
        resolution_estimation[image_number] = msg.resolution_estimation.value();

    if (!msg.az_int_profile.empty() && (msg.az_int_profile.size() == rad_int_bin_to_q.size()))
        data_set_az_int->WriteVec(msg.az_int_profile, {image_number, 0}, {1, rad_int_bin_to_q.size()});

    auto now_time = std::chrono::system_clock::now();
    auto time_diff = now_time - (last_flush + swmr_flush_period);
    if (time_diff.count() >= 0) {
        data_set->Flush();
        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_xfel_event_code->Flush();
        data_set_xfel_pulseid->Flush();
        data_set_indexed->Flush();
        data_set_indexed_lattice->Flush();
        if (!rad_int_bin_to_q.empty())
            data_set_az_int->Flush();

        last_flush = now_time;
    }
}

HDF5DataFileStatistics HDF5DataFile::GetStatistics() const {
    HDF5DataFileStatistics ret;
    ret.max_image_number = max_image_number;
    ret.total_images = nimages;
    ret.filename = filename;
    return ret;
}

size_t HDF5DataFile::GetNumImages() const {
    return nimages;
}

void HDF5DataFile::SetupSWMRFile() {
    HDF5File data_file_local(filename, true, true);
    HDF5Group(data_file_local, "/entry").NXClass("NXentry");
    HDF5Group(data_file_local, "/entry/data").NXClass("NXdata");

    HDF5Group(data_file_local, "/entry/xfel").NXClass("NXcollection");
    HDF5Group(data_file_local, "/entry/result").NXClass("NXcollection");

    HDF5DataSpace data_space({1, ypixel, xpixel}, {H5S_UNLIMITED, ypixel, xpixel});
    HDF5DataSet(data_file_local, "/entry/data/data", data_type, data_space, dcpl);

    HDF5DataSpace data_space_spots({1, max_spots}, {H5S_UNLIMITED, max_spots});
    HDF5Dcpl dcpl_spots;
    dcpl_spots.SetChunking({1, max_spots});
    HDF5DataSet(data_file_local, "/entry/result/peakXPosRaw", HDF5DataType(0.0f), data_space_spots, dcpl_spots);
    HDF5DataSet(data_file_local, "/entry/result/peakYPosRaw", HDF5DataType(0.0f), data_space_spots, dcpl_spots);
    HDF5DataSet(data_file_local, "/entry/result/peakTotalIntensity", HDF5DataType(0.0f), data_space_spots, dcpl_spots);
    HDF5DataSet(data_file_local, "/entry/result/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_local, "/entry/result/nPeaks", HDF5DataType(4, false), data_space_linear, dcpl_linear).SetExtent({1});
    HDF5DataSet(data_file_local, "/entry/xfel/pulseID", HDF5DataType(8, false), data_space_linear, dcpl_linear).SetExtent({1});
    HDF5DataSet(data_file_local, "/entry/xfel/eventCode", HDF5DataType(4, false), data_space_linear, dcpl_linear).SetExtent({1});
    HDF5DataSet(data_file_local, "/entry/result/imageIndexed", HDF5DataType(1, false), 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_local, "/entry/result/latticeIndexed", HDF5DataType(0.0f), data_space_lattice, dcpl_lattice);

    if (!rad_int_bin_to_q.empty()) {
        data_file_local.SaveVector("/entry/result/azimIntegrationBinToQ", rad_int_bin_to_q);

        HDF5DataSpace data_space_rad_int({1, rad_int_bin_to_q.size()}, {H5S_UNLIMITED, rad_int_bin_to_q.size()});
        HDF5Dcpl dcpl_rad_int;
        dcpl_rad_int.SetChunking({1, rad_int_bin_to_q.size()});
        HDF5DataSet(data_file_local, "/entry/result/azimIntegration", HDF5DataType(0.0f), data_space_rad_int,
                    dcpl_rad_int);
    }
    last_flush = std::chrono::system_clock::now();
}