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

#include "JFJochHDF5Reader.h"
#include "../common/PixelMask.h"

std::vector<hsize_t> GetDimension(HDF5Object& object, const std::string& path) {
    const auto dim = object.GetDimension(path);
    if (dim.size() != 3)
        throw JFJochException(JFJochExceptionCategory::HDF5, "Wrong dimension of /entry/data/data");
    return dim;
}

void JFJochHDF5Reader::ReadFile(const std::string& filename) {
    std::unique_lock ul(master_file_mutex);
    try {
        auto dataset = std::make_shared<JFJochReaderDataset>();
        master_file = std::make_unique<HDF5ReadOnlyFile>(filename);

        dataset->geom.BeamX_pxl(master_file->GetFloat("/entry/instrument/detector/beam_center_x"));
        dataset->geom.BeamY_pxl(master_file->GetFloat("/entry/instrument/detector/beam_center_y"));
        dataset->geom.DetectorDistance_mm(master_file->GetFloat("/entry/instrument/detector/distance") * 1000.0);
        dataset->geom.PixelSize_mm(master_file->GetFloat("/entry/instrument/detector/x_pixel_size") * 1000.0);
        dataset->geom.Wavelength_A(master_file->GetFloat("/entry/instrument/beam/incident_wavelength"));

        dataset->saturation_value = master_file->GetInt("/entry/instrument/detector/saturation_value");
        dataset->error_value = master_file->GetOptInt("/entry/instrument/detector/error_value");

        dataset->jfjoch_release = master_file->GetString("/entry/instrument/detector/jfjoch_release");
        dataset->bit_depth_image = master_file->GetInt("/entry/instrument/detector/bit_depth_image");

        dataset->instrument_name = master_file->GetString("/entry/instrument/name");
        dataset->source_name = master_file->GetString("/entry/source/name");

        dataset->frame_time = master_file->GetFloat("/entry/instrument/detector/frame_time");
        dataset->count_time = master_file->GetFloat("/entry/instrument/detector/count_time");
        dataset->detector_name = master_file->GetString("/entry/instrument/detector/description");

        if (master_file->Exists("/entry/roi"))
            dataset->roi = master_file->FindLeafs("/entry/roi");
        for (const auto &s: dataset->roi) {
            dataset->roi_max.emplace_back(master_file->ReadVector<int64_t>("/entry/roi/" + s + "/max"));
            dataset->roi_sum.emplace_back(master_file->ReadVector<int64_t>("/entry/roi/" + s + "/sum"));
            dataset->roi_sum_sq.emplace_back(master_file->ReadVector<int64_t>("/entry/roi/" + s + "/sum_sq"));
            dataset->roi_npixel.emplace_back(master_file->ReadVector<int64_t>("/entry/roi/" + s + "/npixel"));
        }
        if (master_file->Exists("/entry/instrument/attenuator"))
            dataset->attenuator_transmission = master_file->GetOptFloat("/entry/instrument/attenuator/attenuator_transmission");
        dataset->total_flux = master_file->GetOptFloat("/entry/instrument/beam/total_flux");

        if (master_file->Exists("/entry/image/max_value"))
            dataset->max_value = master_file->ReadVector<int64_t>("/entry/image/max_value");

        if (master_file->Exists("/entry/az_int") && master_file->Exists("/entry/az_int/bin_to_q")) {
            HDF5DataSet bin_to_q_dataset(*master_file, "/entry/az_int/bin_to_q");
            bin_to_q_dataset.ReadVector(dataset->az_int_bin_to_q);
        }

        if (master_file->Exists("/entry/data/data")) {
            legacy_format = false;
            auto dim = GetDimension(*master_file, "/entry/data/data");
            number_of_images = dim[0];
            dataset->image_size_y = dim[1];
            dataset->image_size_x = dim[2];

            images_per_file = number_of_images;

            dataset->efficiency = master_file->ReadVector<float>(
                    "/entry/instrument/detector/detectorSpecific/data_collection_efficiency_image");

            dataset->spot_count = master_file->ReadOptVector<uint32_t>("/entry/MX/nPeaks");
            dataset->indexing_result = master_file->ReadOptVector<uint8_t>("/entry/MX/imageIndexed");
            dataset->bkg_estimate = master_file->ReadOptVector<float>("/entry/MX/bkgEstimate");
        } else if (master_file->Exists("/entry/data/data_000001")) {
            legacy_format = true;

            dataset->image_size_x = master_file->GetInt("/entry/instrument/detector/detectorSpecific/x_pixels_in_detector");
            dataset->image_size_y = master_file->GetInt("/entry/instrument/detector/detectorSpecific/y_pixels_in_detector");

            //size_t expected_images = master_file->GetInt("/entry/instrument/detector/detectorSpecific/nimages");

            images_per_file = GetDimension(*master_file, "/entry/data/data_000001")[0];
            number_of_images = images_per_file;
            if (number_of_images > 0) {
                uint32_t file_id = 1;
                while (true) {
                    char buff[32];
                    snprintf(buff, 32, "/entry/data/data_%06d", file_id + 1);
                    if (!master_file->Exists(buff))
                        break;
                    number_of_images += GetDimension(*master_file, buff)[0];
                    file_id++;
                }
            }
        } else {
            dataset->image_size_x = master_file->GetInt("/entry/instrument/detector/detectorSpecific/x_pixels_in_detector");
            dataset->image_size_y = master_file->GetInt("/entry/instrument/detector/detectorSpecific/y_pixels_in_detector");

            number_of_images = 0;
        }

        if (dataset->image_size_x * dataset->image_size_y > 0)
            dataset->pixel_mask = master_file->ReadOptVector<uint32_t>(
                    "/entry/instrument/detector/pixel_mask",
                    {0,0},
                    {dataset->image_size_y, dataset->image_size_x}
            );
        if (dataset->pixel_mask.empty())
            dataset->pixel_mask = std::vector<uint32_t>(dataset->image_size_x * dataset->image_size_y);
        dataset->number_of_images = number_of_images;
        SetStartMessage(dataset);
    } catch (const std::exception& e) {
        master_file = {};
        number_of_images = 0;
        SetStartMessage({});
        throw;
    }
}

void JFJochHDF5Reader::SetStartMessage(std::shared_ptr<JFJochReaderDataset> val) {
    std::unique_lock ul(start_message_mutex);
    start_message = val;
}

std::shared_ptr<JFJochReaderDataset> JFJochHDF5Reader::GetStartMessage() const {
    std::unique_lock ul(start_message_mutex);
    if (!start_message)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "File not loaded");
    return start_message;
}

uint64_t JFJochHDF5Reader::GetNumberOfImages() const {
    std::unique_lock ul(master_file_mutex);
    return number_of_images;
}

void JFJochHDF5Reader::LoadImageDataset(std::vector<int32_t> &output,
                                        const std::string &name,
                                        hsize_t number,
                                        hsize_t width,
                                        hsize_t height) {
    output.resize(width * height);
    std::vector<hsize_t> size = {1, height, width};
    std::vector<hsize_t> start = {static_cast<hsize_t>(number), 0, 0};

    HDF5DataSet dataset(*master_file, name);
    HDF5DataType datatype(dataset);

    if (!datatype.IsSigned() && (datatype.GetElemSize() == 4)) {
        std::vector<uint32_t> output_tmp;
        output_tmp.resize(width * height);
        dataset.ReadVector(output_tmp, start, size);

        for (int i = 0; i < output_tmp.size(); i++) {
            if (output_tmp[i] >= INT32_MAX)
                output[i] = INT32_MAX;
            else
                output[i] = static_cast<int32_t>(output_tmp[i]);
        }
    } else
        dataset.ReadVector(output, start, size);
}

std::shared_ptr<JFJochReaderImage> JFJochHDF5Reader::LoadImageInternal(int64_t image_number) {
    std::unique_lock ul(master_file_mutex);

    if (!master_file)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot load image if file not loaded");

    auto new_image = std::make_shared<JFJochReaderImage>();

    if (image_number >= number_of_images)
        throw JFJochException(JFJochExceptionCategory::HDF5, "Image out of bounds");

    new_image->number = image_number;

    if (legacy_format) {
        const uint32_t file_id = image_number / images_per_file;
        char buff[32];
        snprintf(buff, 32, "/entry/data/data_%06d", file_id + 1);
        const uint32_t image_id = image_number % images_per_file;
        LoadImageDataset(new_image->image,
                         buff,
                         image_id,
                         start_message->image_size_x,
                         start_message->image_size_y);
    } else {
        LoadImageDataset(new_image->image, "/entry/data/data",
                         image_number,
                         start_message->image_size_x,
                         start_message->image_size_y);
    }

    if (!legacy_format
        && (start_message->spot_count.size() > image_number)
        && (start_message->spot_count[image_number] > 0)) {
        auto spot_count = start_message->spot_count[image_number];
        auto spot_x = master_file->ReadVector<float>(
                "/entry/MX/peakXPosRaw",
                {(hsize_t) image_number, 0},
                {1, spot_count}
        );
        auto spot_y = master_file->ReadVector<float>(
                "/entry/MX/peakYPosRaw",
                {(hsize_t) image_number, 0},
                {1, spot_count}
        );
        auto spot_intensity = master_file->ReadVector<float>(
                "/entry/MX/peakTotalIntensity",
                {(hsize_t) image_number, 0},
                {1, spot_count}
        );
        auto spot_indexed = master_file->ReadVector<uint8_t>(
                "/entry/MX/peakIndexed",
                {(hsize_t) image_number, 0},
                {1, spot_count}
        );
        for (int i = 0; i < spot_count; i++) {
            new_image->spots.emplace_back(SpotToSave{
                    .x = spot_x.at(i),
                    .y = spot_y.at(i),
                    .intensity = spot_intensity.at(i),
                    .indexed = (spot_indexed.at(i) != 0)
            });
        }
    }

    if (!legacy_format && !start_message->az_int_bin_to_q.empty()) {
        new_image->az_int = master_file->ReadOptVector<float>(
                "/entry/azint/image",
                {(hsize_t) image_number, 0},
                {1, start_message->az_int_bin_to_q.size()}
        );
    }

    size_t good_pixel = 0;

    for (int i = 0; i < new_image->image.size(); i++) {
        int32_t value = new_image->image[i];

        uint32_t mask_val = 0;
        if (!start_message->pixel_mask.empty())
            mask_val = start_message->pixel_mask.at(i);

        if ((mask_val & (
                (1<<PixelMask::ModuleGapPixelBit)
                | (1<<PixelMask::ChipGapPixelBit)
                | (1<<PixelMask::ModuleEdgePixelBit))) != 0) {
            new_image->image[i] = GAP_PXL_VALUE;
        } else if (new_image->image[i] <= start_message->error_value || (mask_val != 0)) {
            new_image->image[i] = ERROR_PXL_VALUE;
            new_image->error_pixel.emplace(i);
        } else if (new_image->image[i] >= start_message->saturation_value) {
            new_image->image[i] = SATURATED_PXL_VALUE;
            new_image->saturated_pixel.emplace(i);
        } else {
            good_pixel++;
            new_image->valid_pixel.emplace(value, i);
        }
    }

    new_image->dataset = start_message;
    return new_image;
}

void JFJochHDF5Reader::LoadImage(int64_t image_number) {
    auto new_image = LoadImageInternal(image_number);
    {
        std::unique_lock ul(current_image_mutex);
        current_image = new_image;
    }
}

std::shared_ptr<JFJochReaderImage> JFJochHDF5Reader::CopyImage() {
    std::unique_lock ul(current_image_mutex);
    return current_image;
}

void JFJochHDF5Reader::Close() {
    std::unique_lock ul(master_file_mutex);
    master_file = {};
    number_of_images = 0;
    SetStartMessage({});
}