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

#include "JFJochCrystFELReader.h"

#include <algorithm>
#include <cmath>
#include <filesystem>
#include <numeric>

#include <zlib.h>

#include "../common/JFJochException.h"
#include "../common/Logger.h"
#include "../writer/HDF5Objects.h"

namespace {

bool IsPath(const std::string &s) {
    return !s.empty() && s[0] == '/';
}

std::string ReplaceEvent(std::string path, const std::string &group) {
    const auto pos = path.find('%');
    if (pos != std::string::npos)
        path.replace(pos, 1, group);
    return path;
}

}  // namespace

JFJochCrystFELReader::JFJochCrystFELReader() = default;

JFJochCrystFELReader::~JFJochCrystFELReader() {
    Close();
}

// One number for the whole run from a value the geometry takes from the image files: a scalar, a
// per-frame vector, or one scalar per shot group.
double JFJochCrystFELReader::MeanOverFirstFile(const std::string &path) {
    auto &file = *open_files_[0];
    std::vector<double> values;
    if (path.find('%') != std::string::npos) {
        for (const auto &f: frames_)
            if (f.file == 0)
                values.push_back(HDF5DataSet(file, ReplaceEvent(path, f.group)).ReadScalar<double>());
    } else if (file.GetDimension(path).empty())
        values.push_back(HDF5DataSet(file, path).ReadScalar<double>());
    else
        values = file.ReadVector<double>(path);
    if (values.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No values in " + path);
    return std::accumulate(values.begin(), values.end(), 0.0) / static_cast<double>(values.size());
}

void JFJochCrystFELReader::ReadFiles(const std::string &path, const std::string &geometry_file) {
    Close();
    geom_ = crystfel::Read(geometry_file);

    if (std::filesystem::is_directory(path)) {
        for (const auto &e: std::filesystem::directory_iterator(path))
            if (e.path().extension() == ".h5" || e.path().extension() == ".cxi")
                files_.push_back(e.path().string());
        std::sort(files_.begin(), files_.end());
    } else
        files_.push_back(path);
    if (files_.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "No .h5 or .cxi files in " + path);

    std::unique_lock ul(hdf5_mutex);

    // Every frame of every file, in file order.
    const auto pct = geom_.data.find('%');
    for (size_t f = 0; f < files_.size(); f++) {
        open_files_.emplace_back(std::make_unique<HDF5ReadOnlyFile>(files_[f]));
        auto &file = *open_files_.back();
        if (pct != std::string::npos) {
            const std::string parent = geom_.data.substr(0, pct);
            for (const auto &group: file.FindLeafs(parent.empty() ? "/" : parent)) {
                const auto dataset = ReplaceEvent(geom_.data, group);
                if (file.Exists(dataset))
                    frames_.push_back({f, dataset, -1, group});
            }
        } else {
            const auto dim = file.GetDimension(geom_.data);
            if (dim.size() == 3) {
                for (hsize_t i = 0; i < dim[0]; i++)
                    frames_.push_back({f, geom_.data, static_cast<int64_t>(i), ""});
            } else if (dim.size() == 2)
                frames_.push_back({f, geom_.data, -1, ""});
        }
    }
    if (frames_.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No frames at " + geom_.data + " in " + path);

    const auto dim = open_files_[frames_[0].file]->GetDimension(frames_[0].dataset);
    frame_height_ = dim[dim.size() - 2];
    frame_width_ = dim[dim.size() - 1];

    const auto &p0 = geom_.panels[0];
    const double clen_m = IsPath(p0.clen) ? MeanOverFirstFile(p0.clen) * 1e-3 : std::stod(p0.clen);
    const double distance_mm = (clen_m + p0.coffset) * 1000.0;
    if (geom_.photon_energy.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              geometry_file + " states no photon_energy");
    const double photon_energy_eV = IsPath(geom_.photon_energy) ? MeanOverFirstFile(geom_.photon_energy)
                                                                : std::stod(geom_.photon_energy);
    for (const auto &p: geom_.panels)
        if (p.res != p0.res || p.clen != p0.clen || p.coffset != p0.coffset)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  geometry_file + ": panels at different distances or with different "
                                                  "pixel sizes cannot be put on one flat detector");

    std::vector<uint32_t> file_mask;
    if (!geom_.mask.empty()) {
        if (!geom_.mask_file.empty()) {
            HDF5ReadOnlyFile mask_file(geom_.mask_file);
            file_mask = mask_file.ReadVector<uint32_t>(geom_.mask, {0, 0}, {frame_height_, frame_width_});
        } else
            file_mask = open_files_[0]->ReadVector<uint32_t>(geom_.mask, {0, 0}, {frame_height_, frame_width_});
    }
    CheckDirectChunks();
    ul.unlock();

    // CrystFEL's laboratory frame and the flat image's (column, row, beam) are both right-handed with
    // the beam along z, so the point (X, Y) lands at column beam_x + X, row beam_y + Y - and CrystFEL's
    // horizontal x stays the image's horizontal, where the polarisation correction expects it. A pixel
    // centre is at its index, so beam_x is chosen to put the first panel's pixel centres on whole columns.
    auto position = [](const crystfel::Panel &p, double fs, double ss) {
        return std::pair{p.corner_x + fs * p.fs_x + ss * p.ss_x, p.corner_y + fs * p.fs_y + ss * p.ss_y};
    };
    double max_x = -1e9, min_x = 1e9, max_y = -1e9, min_y = 1e9;
    for (const auto &p: geom_.panels)
        for (double fs: {0.0, double(p.max_fs - p.min_fs + 1)})
            for (double ss: {0.0, double(p.max_ss - p.min_ss + 1)}) {
                const auto [x, y] = position(p, fs, ss);
                max_x = std::max(max_x, x);
                min_x = std::min(min_x, x);
                max_y = std::max(max_y, y);
                min_y = std::min(min_y, y);
            }
    const auto [x0, y0] = position(p0, 0.5, 0.5);
    const double beam_x = std::ceil(x0 - min_x) + 1 - x0;
    const double beam_y = std::ceil(y0 - min_y) + 1 - y0;
    width_ = static_cast<size_t>(std::ceil(beam_x + max_x)) + 2;
    height_ = static_cast<size_t>(std::ceil(beam_y + max_y)) + 2;

    // Everything starts as a gap; the panels' pixels are put on top, less the bad regions.
    std::vector<uint32_t> mask(width_ * height_, 1u << PixelMask::ModuleGapPixelBit);
    destination_.assign(frame_width_ * frame_height_, -1);
    adu_per_photon_.assign(frame_width_ * frame_height_, 1.0f);
    counts_in_photons_ = true;
    std::optional<double> saturation;

    for (const auto &p: geom_.panels) {
        if (p.max_fs >= static_cast<int64_t>(frame_width_) || p.max_ss >= static_cast<int64_t>(frame_height_))
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Panel " + p.name + " lies outside the stored frame");
        double gain = 1.0;
        if (p.adu_per_photon.has_value())
            gain = p.adu_per_photon.value();
        else if (p.adu_per_eV.has_value())
            gain = p.adu_per_eV.value() * photon_energy_eV;
        else
            counts_in_photons_ = false;
        if (p.max_adu.has_value())
            saturation = std::min(saturation.value_or(INT64_MAX), p.max_adu.value() / gain);
        if (p.no_index)
            continue;

        for (int64_t ss = p.min_ss; ss <= p.max_ss; ss++) {
            for (int64_t fs = p.min_fs; fs <= p.max_fs; fs++) {
                const auto [x, y] = position(p, fs - p.min_fs + 0.5, ss - p.min_ss + 0.5);
                bool bad = false;
                for (const auto &b: geom_.bad) {
                    if (b.in_lab)
                        bad |= x >= b.min_x && x <= b.max_x && y >= b.min_y && y <= b.max_y;
                    else if (b.panel.empty() || b.panel == p.name)
                        bad |= fs >= b.min_fs && fs <= b.max_fs && ss >= b.min_ss && ss <= b.max_ss;
                }
                if (bad)
                    continue;

                const size_t src = ss * frame_width_ + fs;
                const auto col = std::lround(beam_x + x);
                const auto row = std::lround(beam_y + y);
                const size_t dst = static_cast<size_t>(row) * width_ + static_cast<size_t>(col);
                destination_[src] = static_cast<int64_t>(dst);
                adu_per_photon_[src] = static_cast<float>(gain);

                const uint32_t m = src < file_mask.size() ? file_mask[src] : 0;
                const bool masked = (geom_.mask_bad != 0 && (m & geom_.mask_bad) != 0)
                                    || (geom_.mask_good != 0 && (m & geom_.mask_good) != geom_.mask_good);
                mask[dst] = masked ? (1u << PixelMask::ErrorPixelBit) : 0;
            }
        }
    }

    dataset_ = std::make_shared<JFJochReaderDataset>();
    dataset_->experiment = default_experiment;
    DetectorSetup detector = DetDECTRIS(static_cast<int64_t>(width_), static_cast<int64_t>(height_),
                                        "CrystFEL geometry", {});
    detector.PixelSize_um(static_cast<float>(1e6 / p0.res));
    // Nothing a parallax or sensor-efficiency correction needs is stated in a geometry file.
    detector.SensorThickness_um(0);
    if (saturation.has_value())
        detector.SaturationLimit(static_cast<int64_t>(std::floor(saturation.value())));
    detector.BitDepthImage(32);
    detector.MinFrameTime(std::chrono::microseconds(0));
    detector.MinCountTime(std::chrono::microseconds(0));
    detector.ReadOutTime(std::chrono::nanoseconds(0));
    dataset_->experiment.Detector(detector);
    dataset_->experiment.BeamX_pxl(static_cast<float>(beam_x));
    dataset_->experiment.BeamY_pxl(static_cast<float>(beam_y));
    dataset_->experiment.DetectorDistance_mm(static_cast<float>(distance_mm));
    dataset_->experiment.IncidentEnergy_keV(static_cast<float>(photon_energy_eV / 1000.0));
    dataset_->experiment.ImagesPerTrigger(static_cast<int64_t>(frames_.size()));
    dataset_->pixel_mask = std::make_shared<const PixelMask>(mask);

    Logger("CrystFELReader").Info("{} frames in {} files, {} panels; {} x {} image, beam at ({:.2f}, {:.2f}) px, "
                                  "{:.3f} mm, {:.1f} eV, {}{}", frames_.size(), files_.size(), geom_.panels.size(),
                                  width_, height_, beam_x, beam_y, distance_mm, photon_energy_eV,
                                  counts_in_photons_ ? "converted to photons" : "in detector units (no gain stated)",
                                  direct_chunks_ ? ", chunks decoded in parallel" : "");
    SetStartMessage(dataset_);
}

void JFJochCrystFELReader::CheckDirectChunks() {
    const auto &frame = frames_[0];
    HDF5DataSet dataset(*open_files_[frame.file], frame.dataset);

    const hid_t dcpl = H5Dget_create_plist(dataset.GetID());
    const int rank = frame.index >= 0 ? 3 : 2;
    std::vector<hsize_t> chunk(rank, 0);
    bool ok = H5Pget_layout(dcpl) == H5D_CHUNKED && H5Pget_chunk(dcpl, rank, chunk.data()) == rank
              && chunk[rank - 1] == frame_width_ && chunk[rank - 2] == frame_height_ && (rank == 2 || chunk[0] == 1);
    bool deflate = false;
    shuffled_ = false;
    for (int i = 0; ok && i < H5Pget_nfilters(dcpl); i++) {
        unsigned int flags;
        size_t nelements = 0;
        const H5Z_filter_t filter = H5Pget_filter2(dcpl, i, &flags, &nelements, nullptr, 0, nullptr, nullptr);
        if (filter == H5Z_FILTER_DEFLATE)
            deflate = true;
        else if (filter == H5Z_FILTER_SHUFFLE)
            shuffled_ = true;
        else
            ok = false;
    }
    H5Pclose(dcpl);

    const hid_t type = H5Dget_type(dataset.GetID());
    stored_size_ = H5Tget_size(type);
    stored_float_ = H5Tget_class(type) == H5T_FLOAT;
    stored_signed_ = H5Tget_sign(type) == H5T_SGN_2;
    ok &= (H5Tget_class(type) == H5T_INTEGER && (stored_size_ == 2 || stored_size_ == 4))
          || (stored_float_ && stored_size_ == 4);
    H5Tclose(type);

    direct_chunks_ = ok && deflate;
}

bool JFJochCrystFELReader::CountsInPhotons() const {
    return counts_in_photons_;
}

uint64_t JFJochCrystFELReader::GetNumberOfImages() const {
    return frames_.size();
}

void JFJochCrystFELReader::Close() {
    std::unique_lock ul(hdf5_mutex);
    open_files_.clear();
    files_.clear();
    frames_.clear();
    dataset_.reset();
}

template <class Buffer>
void JFJochCrystFELReader::DecodeInto(int64_t image_number, Buffer &buffer) const {
    const auto &frame = frames_.at(image_number);
    std::vector<float> stored(frame_width_ * frame_height_);
    if (direct_chunks_) {
        std::vector<uint8_t> compressed;
        {
            std::unique_lock ul(hdf5_mutex);
            HDF5DataSet dataset(*open_files_[frame.file], frame.dataset);
            if (frame.index >= 0)
                dataset.ReadDirectChunk(compressed, {static_cast<hsize_t>(frame.index), 0, 0});
            else
                dataset.ReadDirectChunk(compressed, {0, 0});
        }
        const size_t n = stored.size();
        std::vector<uint8_t> raw(n * stored_size_);
        uLongf raw_size = raw.size();
        if (uncompress(raw.data(), &raw_size, compressed.data(), compressed.size()) != Z_OK || raw_size != raw.size())
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Cannot inflate " + frame.dataset + " of " + files_[frame.file]);
        if (shuffled_) {
            std::vector<uint8_t> tmp(raw.size());
            for (size_t i = 0; i < n; i++)
                for (size_t b = 0; b < stored_size_; b++)
                    tmp[i * stored_size_ + b] = raw[b * n + i];
            raw.swap(tmp);
        }
        for (size_t i = 0; i < n; i++) {
            const uint8_t *v = raw.data() + i * stored_size_;
            if (stored_float_)
                stored[i] = *reinterpret_cast<const float *>(v);
            else if (stored_size_ == 2)
                stored[i] = stored_signed_ ? *reinterpret_cast<const int16_t *>(v) : *reinterpret_cast<const uint16_t *>(v);
            else
                stored[i] = stored_signed_ ? *reinterpret_cast<const int32_t *>(v) : *reinterpret_cast<const uint32_t *>(v);
        }
    } else {
        std::unique_lock ul(hdf5_mutex);
        HDF5DataSet dataset(*open_files_[frame.file], frame.dataset);
        if (frame.index >= 0)
            dataset.ReadVector(stored, {static_cast<hsize_t>(frame.index), 0, 0}, {1, frame_height_, frame_width_});
        else
            dataset.ReadVector(stored, {0, 0}, {frame_height_, frame_width_});
    }

    buffer.assign(width_ * height_ * sizeof(int32_t), 0);
    auto out = reinterpret_cast<int32_t *>(buffer.data());
    for (size_t i = 0; i < stored.size(); i++)
        if (destination_[i] >= 0)
            out[destination_[i]] = static_cast<int32_t>(std::lround(stored[i] / adu_per_photon_[i]));
}

bool JFJochCrystFELReader::LoadImage_i(std::shared_ptr<JFJochReaderDataset> &dataset,
                                       DataMessage &message,
                                       std::vector<uint8_t> &buffer,
                                       int64_t image_number,
                                       bool update_dataset) {
    (void) update_dataset;
    if (!dataset || image_number < 0 || static_cast<size_t>(image_number) >= frames_.size())
        return false;
    DecodeInto(image_number, buffer);
    message.image = CompressedImage(buffer.data(), buffer.size(), width_, height_,
                                    CompressedImageMode::Int32, CompressionAlgorithm::NO_COMPRESSION);
    message.number = image_number;
    return true;
}

bool JFJochCrystFELReader::ReadRawImage(int64_t image_number, JFJochReaderRawImage &image) {
    if (image_number < 0 || static_cast<size_t>(image_number) >= frames_.size())
        return false;
    DecodeInto(image_number, image.image_buffer);
    image.image = CompressedImage(image.image_buffer.data(), image.image_buffer.size(), width_, height_,
                                  CompressedImageMode::Int32, CompressionAlgorithm::NO_COMPRESSION);
    return true;
}

std::vector<SpotToSave> JFJochCrystFELReader::ReadSpots(int64_t) const {
    return {};
}