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

#include <cmath>
#include <sys/stat.h>

#include "HDF5NXmx.h"

#include "../common/GitInfo.h"
#include "../include/spdlog/fmt/fmt.h"
#include "MakeDirectory.h"
#include "../common/time_utc.h"

NXmx::NXmx(const StartMessage &start)
        : start_message(start),
          filename(start.file_prefix + "_master.h5") {

    tmp_filename = fmt::format("{}.{:x}.tmp", filename, std::chrono::system_clock::now().time_since_epoch().count());

    MakeDirectory(filename);

    hdf5_file = std::make_unique<HDF5File>(tmp_filename);

    hdf5_file->Attr("file_name", filename);
    hdf5_file->Attr("HDF5_Version", hdf5_version());
    HDF5Group(*hdf5_file, "/entry").NXClass("NXentry").SaveScalar("definition", "NXmx");
    HDF5Group(*hdf5_file, "/entry/result").NXClass("NXcollection");
    hdf5_file->SaveScalar("/entry/start_time", start.arm_date);

    Facility(start);
    Detector(start);
    Metrology(start);
    Beam(start);
    Attenuator(start);
    UserData(start);
}

NXmx::~NXmx() {
    if (!std::filesystem::exists(filename.c_str()))
        std::rename(tmp_filename.c_str(), filename.c_str());
}

std::string HDF5Metadata::DataFileName(const StartMessage &msg, int64_t file_number) {
    if (file_number < 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "File number cannot be negative");

    if (msg.source_name == "SwissFEL") {
        if (file_number >= 10000)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Format doesn't allow for 10'000 or more files");
        else if (msg.detector_serial_number.empty())
            return fmt::format("{:s}{:04d}.JF.h5", msg.file_prefix, file_number + 1);
        else
            return fmt::format("{:s}{:04d}.{:s}.h5", msg.file_prefix, file_number + 1, msg.detector_serial_number);
    } else {
        if (file_number >= 1000000)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Format doesn't allow for 1 million or more files");
        else
            return fmt::format("{:s}_data_{:06d}.h5", msg.file_prefix, file_number + 1);
    }
}

void NXmx::LinkToData(const StartMessage &start, const EndMessage &end) {
    hsize_t total_images = end.max_image_number;
    hsize_t images_per_file = start.images_per_file;
    hsize_t file_count = 0;
    if (start.images_per_file > 0) {
        file_count = total_images / images_per_file;
        if (total_images % images_per_file > 0)
            file_count++;
    }

    HDF5Group(*hdf5_file, "/entry/data").NXClass("NXdata");

    for (uint32_t file_id = 0; file_id < file_count; file_id++) {
        char buff[32];
        snprintf(buff,32,"/entry/data/data_%06d", file_id+1);
        hdf5_file->ExternalLink(HDF5Metadata::DataFileName(start, file_id),
                                "/entry/data/data",
                                std::string(buff));
    }
}

void NXmx::Detector(const StartMessage &start) {
    HDF5Group group(*hdf5_file, "/entry/instrument/detector");
    group.NXClass("NXdetector");
    SaveScalar(group, "beam_center_x", start.beam_center_x)->Units("pixel");
    SaveScalar(group, "beam_center_y", start.beam_center_y)->Units("pixel");
    SaveScalar(group, "distance", start.detector_distance)->Units("m");
    SaveScalar(group, "detector_distance", start.detector_distance)->Units("m");

    SaveScalar(group, "count_time", start.count_time)->Units("s");
    SaveScalar(group, "frame_time", start.frame_time)->Units("s");

    SaveScalar(group, "sensor_thickness", start.sensor_thickness)->Units("m");
    SaveScalar(group, "x_pixel_size", start.pixel_size_x)->Units("m");
    SaveScalar(group, "y_pixel_size", start.pixel_size_y)->Units("m");
    SaveScalar(group, "sensor_material", start.sensor_material);
    SaveScalar(group, "description", start.detector_description);

    SaveScalar(group, "bit_depth_image", start.bit_depth_image);
    if (start.bit_depth_readout)
        SaveScalar(group, "bit_depth_readout", start.bit_depth_readout.value());
    SaveScalar(group, "saturation_value", start.saturation_value);
    if (start.error_value)
        SaveScalar(group, "error_value", start.error_value.value()); // this is not NXmx
    SaveScalar(group, "flatfield_applied", start.flatfield_enabled);
    SaveScalar(group, "pixel_mask_applied", start.pixel_mask_enabled);

    if (start.jungfrau_conversion_enabled)
        SaveScalar(group, "jungfrau_conversion_applied", start.jungfrau_conversion_enabled.value());
    if (start.jungfrau_conversion_factor)
        SaveScalar(group, "jungfrau_conversion_factor", start.jungfrau_conversion_factor.value())->Units("eV");

    SaveScalar(group, "geometry_transformation_applied", start.geometry_transformation_enabled.value_or(true));

    SaveScalar(group, "acquisition_type", "triggered");
    SaveScalar(group, "countrate_correction_applied", start.countrate_correction_enabled);
    SaveScalar(group, "number_of_cycles", start.summation);

    HDF5Group det_specific(group, "detectorSpecific");
    det_specific.NXClass("NXcollection");

    if (!start.jfjoch_release.empty())
        SaveScalar(det_specific, "jfjoch_release", start.jfjoch_release);
    SaveScalar(det_specific, "ntrigger", 1);

    SaveScalar(det_specific, "x_pixels_in_detector", static_cast<uint32_t>(start.image_size_x));
    SaveScalar(det_specific, "y_pixels_in_detector", static_cast<uint32_t>(start.image_size_y));
    SaveScalar(det_specific, "software_git_commit", jfjoch_git_sha1());
    SaveScalar(det_specific, "software_git_date", jfjoch_git_date());
    if (start.storage_cell_number) {
        SaveScalar(det_specific, "storage_cell_number", static_cast<uint32_t>(start.storage_cell_number.value()));
        if (start.storage_cell_number.value() > 1)
            SaveScalar(det_specific, "storage_cell_delay", static_cast<uint32_t>(start.storage_cell_delay_ns))->Units(
                    "ns");
    }

    if (start.data_reduction_factor_serialmx)
        det_specific.SaveScalar("data_reduction_factor_serialmx", start.data_reduction_factor_serialmx.value());

    if (!start.gain_file_names.empty())
        det_specific.SaveVector("gain_file_names", start.gain_file_names);

    if (!start.pixel_mask.empty()) {
        SaveCBORImage("/entry/instrument/detector/detectorSpecific/pixel_mask", start.pixel_mask[0]);
        hdf5_file->HardLink("/entry/instrument/detector/detectorSpecific/pixel_mask",
                            "/entry/instrument/detector/pixel_mask");
    }
}

void NXmx::Detector(const EndMessage &end) {
    SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/nimages", end.max_image_number);
    if (end.images_collected_count)
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/nimages_collected", end.images_collected_count.value());
    if (end.images_sent_to_write_count)
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/nimages_written", end.images_sent_to_write_count.value());
    if (end.efficiency)
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/data_collection_efficiency", end.efficiency.value());
    if (end.max_receiver_delay)
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/max_receiver_delay", end.max_receiver_delay.value());
}

void NXmx::DetectorModule(const std::string &name, const std::vector<int32_t> &origin, const std::vector<int32_t> &size,
                          const std::vector<double> &fast_axis, const std::vector<double> &slow_axis,
                          const std::string &nx_axis, double pixel_size_mm) {
    HDF5Group module_group(*hdf5_file, "/entry/instrument/detector/" + name);

    module_group.NXClass("NXdetector_module");

    module_group.SaveVector("data_origin", origin);
    module_group.SaveVector("data_size", size);

    SaveScalar(module_group, "fast_pixel_direction", pixel_size_mm)->
            Transformation("m", "/entry/instrument/detector/transformations/" + nx_axis,
                           "", "", "translation", fast_axis,
                           {0,0,0}, "");

    SaveScalar(module_group, "slow_pixel_direction", pixel_size_mm)->
            Transformation("m", "/entry/instrument/detector/transformations/" + nx_axis,
                           "", "", "translation", slow_axis,
                           {0,0,0}, "");

    SaveScalar(module_group, "module_offset", 0)->
            Transformation("m", "/entry/instrument/detector/transformations/" + nx_axis,
                           "", "", "translation", {0,0,0});
}

void NXmx::Facility(const StartMessage &start) {
    HDF5Group(*hdf5_file, "/entry/source").NXClass("NXsource");
    SaveScalar(*hdf5_file, "/entry/source/name", start.source_name);

    if (!start.source_type.empty())
        SaveScalar(*hdf5_file, "/entry/source/type", start.source_type);

    HDF5Group(*hdf5_file, "/entry/instrument").NXClass("NXinstrument");
    SaveScalar(*hdf5_file, "/entry/instrument/name", start.instrument_name);
}

void NXmx::Beam(const StartMessage &start) {
    HDF5Group group(*hdf5_file, "/entry/instrument/beam");
    group.NXClass("NXbeam");
    SaveScalar(group, "incident_wavelength", start.incident_wavelength)->Units("angstrom");
    if (start.total_flux)
        SaveScalar(group, "total_flux", start.total_flux.value())->Units("Hz");
}

void NXmx::Metrology(const StartMessage &start) {
    HDF5Group transformations(*hdf5_file, "/entry/instrument/detector/transformations");
    transformations.NXClass("NXtransformations");

    std::vector<double> vector{start.beam_center_x * start.pixel_size_x,
                               start.beam_center_y * start.pixel_size_y,
                               start.detector_distance};

    double vector_length = sqrt(vector[0] * vector[0] + vector[1] * vector[1] + vector[2] * vector[2]);
    std::vector<double> vector_norm{vector[0] / vector_length, vector[1]/vector_length, vector[2]/vector_length};

    SaveScalar(transformations, "translation", vector_length)->
            Transformation("m", ".", "detector", "detector_arm", "translation", vector_norm);

    // https://manual.nexusformat.org/classes/base_classes/NXdetector_module.html?highlight=nxdetector_module
    // The order of indices (i, j or i, j, k) is slow to fast.
    // though EIGER has is the other way round
    // Confusing....
    std::vector<int32_t> origin = {0, 0};
    std::vector<int32_t> size = {static_cast<int32_t>(start.image_size_y),
                                 static_cast<int32_t>(start.image_size_x)};

    DetectorModule("module", origin, size, {-1,0,0}, {0,-1,0}, "translation", start.pixel_size_x);
}

void NXmx::Sample(const StartMessage &start, const EndMessage &end) {
    HDF5Group group(*hdf5_file, "/entry/sample");
    group.NXClass("NXsample");
    if (!start.sample_name.empty())
        group.SaveScalar("name", start.sample_name);

    if (start.space_group_number > 0)
        group.SaveScalar("space_group", start.space_group_number);

    if (start.unit_cell) {
        std::vector<float> v = {start.unit_cell->a, start.unit_cell->b, start.unit_cell->c,
                                start.unit_cell->alpha, start.unit_cell->beta, start.unit_cell->gamma};
        group.SaveVector("unit_cell", v);
    }

    if ((end.max_image_number > 0) && start.goniometer) {
        group.SaveScalar("depends_on", "/entry/sample/transformations/" + start.goniometer->name);
        HDF5Group transformations(group, "transformations");
        transformations.NXClass("NXtransformations");

        std::vector<double> angle_container(end.max_image_number);
        for (int32_t i = 0; i < end.max_image_number; i++)
            angle_container[i] = start.goniometer->start + i * start.goniometer->increment;
        std::vector<double> axis = {start.rotation_axis[0], start.rotation_axis[1], start.rotation_axis[2]};

        SaveVector(transformations, start.goniometer->name, angle_container)->
                Transformation("deg", ".", "", "", "rotation", axis, {0,0,0}, "");
    } else
        group.SaveScalar("depends_on", ".");
}

void NXmx::Attenuator(const StartMessage &start) {
    if (start.attenuator_transmission) {
        HDF5Group group(*hdf5_file, "/entry/instrument/attenuator");
        group.NXClass("NXattenuator");
        SaveScalar(group, "attenuator_transmission", start.attenuator_transmission.value());
    }
}

void NXmx::WriteCalibration(const CompressedImage &image) {
    if (!calibration_group_created) {
        calibration_group_created = true;
        HDF5Group(*hdf5_file, "/entry/instrument/detector/calibration").NXClass("NXcollection");
    }
    SaveCBORImage("/entry/instrument/detector/calibration/" + image.channel, image);
}

void NXmx::SaveCBORImage(const std::string &hdf5_path, const CompressedImage &image) {
    std::vector<hsize_t> dims = {image.ypixel, image.xpixel};

    HDF5DataType data_type(image.pixel_depth_bytes, image.pixel_is_signed);
    HDF5Dcpl dcpl;

    if (image.algorithm != CompressionAlgorithm::NO_COMPRESSION) {
        dcpl.SetCompression(image.algorithm, H5Tget_size(data_type.GetID()), 0);
        dcpl.SetChunking(dims);
    }

    HDF5DataSpace data_space(dims);
    auto dataset = std::make_unique<HDF5DataSet>(*hdf5_file, hdf5_path, data_type, data_space, dcpl);

    if (image.algorithm == CompressionAlgorithm::NO_COMPRESSION)
        dataset->Write(data_type, image.data);
    else
        dataset->WriteDirectChunk(image.data, image.size, {0,0});
}

void NXmx::AzimuthalIntegration(const StartMessage &start, const EndMessage &end) {
    if (!start.az_int_bin_to_q.empty()) {
        HDF5Group rad_int_group(*hdf5_file, "/entry/result/azimIntegration");
        rad_int_group.SaveVector("bin_to_q", start.az_int_bin_to_q);

        for (const auto &[x,y] : end.az_int_result)
            rad_int_group.SaveVector(x, y);
    }
}

void NXmx::ADUHistogram(const EndMessage &end) {
    if (!end.adu_histogram.empty()) {
        HDF5Group adu_histo_group(*hdf5_file, "/entry/result/adu_histogram");
        adu_histo_group.SaveScalar("bin_width", end.adu_histogram_bin_width);
        for (const auto &[x, y]: end.adu_histogram)
            adu_histo_group.SaveVector(x, y);
    }
}

void NXmx::Finalize(const EndMessage &end) {
    if (end.end_date) {
        hdf5_file->Attr("file_time", end.end_date.value());
        hdf5_file->SaveScalar("/entry/end_time", end.end_date.value());
        hdf5_file->SaveScalar("/entry/end_time_estimated", end.end_date.value());
    } else {
        std::string time_now = time_UTC(std::chrono::system_clock::now());
        hdf5_file->Attr("file_time", time_now);
        hdf5_file->SaveScalar("/entry/end_time", time_now);
        hdf5_file->SaveScalar("/entry/end_time_estimated", time_now);
    }

    Detector(end);
    LinkToData(start_message, end);
    Sample(start_message, end);
    AzimuthalIntegration(start_message, end);
    ADUHistogram(end);
}

void NXmx::UserData(const StartMessage &start) {
    if (!start.user_data.empty()
        && start.user_data.contains("hdf5")
        && start.user_data["hdf5"].is_object()) {
        HDF5Group group(*hdf5_file, "/entry/user");
        group.NXClass("NXcollection");

        for (const auto &[x,y]: start.user_data["hdf5"].items()) {
            if (y.is_number())
                group.SaveScalar(x, y.get<double>());
            else if (y.is_string())
                group.SaveScalar(x, y.get<std::string>());
        }
    }
}