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

#include "JFJochXTCReader.h"

#include <algorithm>
#include <cmath>
#include <map>
#include <numeric>

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

namespace {

// The pixel size of a Rayonix image, from the model name - an MX340-HS is 340 mm across - and the
// image's width. Only where the width times the binning is a whole readout of 3840, 5760 or 7680
// pixels, i.e. no region of interest was read out.
std::optional<double> RayonixPixel_um(const std::string &model, uint32_t binning, uint32_t width) {
    if (!model.starts_with("MX") || binning == 0)
        return std::nullopt;
    const auto digits = model.substr(2, model.find_first_not_of("0123456789", 2) - 2);
    if (digits.empty())
        return std::nullopt;
    const uint64_t native = static_cast<uint64_t>(width) * binning;
    if (native != 3840 && native != 5760 && native != 7680)
        return std::nullopt;
    return std::stod(digits) * 1000.0 / width;
}

}  // namespace

bool JFJochXTCReader::CanRead(const std::string &path) {
    return xtc::CanRead(path);
}

void JFJochXTCReader::ReadFiles(const std::vector<std::string> &streams) {
    Close();
    Logger logger("XTCReader");
    std::vector<xtc::Frame> all;
    std::vector<double> energies;
    std::string model;
    uint32_t binning = 0;
    size_t damaged = 0;
    for (const auto &stream: streams) {
        auto index = xtc::IndexFile(stream);
        all.insert(all.end(), index.frames.begin(), index.frames.end());
        energies.insert(energies.end(), index.photon_energy_eV.begin(), index.photon_energy_eV.end());
        if (!index.rayonix_model.empty()) {
            model = index.rayonix_model;
            binning = index.rayonix_binning;
        }
        damaged += index.damaged;
        if (index.truncated)
            logger.Warning("{} ends inside its last event, which is left out", stream);
    }

    // The diffraction detector is the camera with the largest image. Another camera of the same size
    // would make that choice a guess, so it is refused; an event without it is left out.
    std::map<uint64_t, uint64_t> area;
    for (const auto &f: all)
        area[f.source] = static_cast<uint64_t>(f.width) * f.height;
    if (area.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No camera images found in the XTC streams given");
    const auto largest = std::max_element(area.begin(), area.end(),
                                          [](const auto &a, const auto &b) { return a.second < b.second; });
    if (std::count_if(area.begin(), area.end(), [&](const auto &a) { return a.second == largest->second; }) > 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "The XTC streams hold images of the same size from more than one camera; which is the "
                                     "diffraction detector is not clear");
    for (const auto &f: all)
        if (f.source == largest->first)
            frames_.push_back(f);
    for (const auto &f: frames_)
        if (f.width != frames_.front().width || f.height != frames_.front().height)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  f.file + ": the images are not all the same size");
    if (damaged > 0)
        logger.Warning("{} damaged, compressed or unknown-version images left out", damaged);

    const auto &f0 = frames_.front();
    if (!energies.empty())
        photon_energy_eV_ = std::accumulate(energies.begin(), energies.end(), 0.0) / static_cast<double>(energies.size());
    pixel_um_ = RayonixPixel_um(model, binning, f0.width);

    dataset_ = std::make_shared<JFJochReaderDataset>();
    dataset_->experiment = default_experiment;
    DetectorSetup detector = DetDECTRIS(f0.width, f0.height, model.empty() ? "XTC camera" : model, {});
    // A CCD: the phosphor converts at the surface, so there is no depth to correct parallax for.
    detector.SensorThickness_um(0);
    detector.SaturationLimit(65535);
    detector.BitDepthImage(32);
    detector.MinFrameTime(std::chrono::microseconds(0));
    detector.MinCountTime(std::chrono::microseconds(0));
    detector.ReadOutTime(std::chrono::nanoseconds(0));
    if (pixel_um_)
        detector.PixelSize_um(static_cast<float>(*pixel_um_));

    dataset_->experiment.Detector(detector);
    dataset_->experiment.BeamX_pxl(0.5f * static_cast<float>(f0.width));
    dataset_->experiment.BeamY_pxl(0.5f * static_cast<float>(f0.height));
    if (photon_energy_eV_)
        dataset_->experiment.IncidentEnergy_keV(static_cast<float>(*photon_energy_eV_ / 1000.0));
    logger.Info("{} images in XTC ({}); {}; {}; one flat panel centred on the beam", frames_.size(),
                model.empty() ? "camera" : model,
                photon_energy_eV_ ? fmt::format("photon energy {:.1f} eV, the mean of {} beam records",
                                                *photon_energy_eV_, energies.size())
                                  : std::string("no beam records"),
                pixel_um_ ? fmt::format("{:.2f} um pixels ({}x{} binning)", *pixel_um_, binning, binning)
                          : std::string("pixel size not stated"));

    dataset_->error_value = -1;
    dataset_->experiment.ImagesPerTrigger(static_cast<int64_t>(frames_.size()));
    dataset_->pixel_mask = std::make_shared<const PixelMask>(f0.width, f0.height);
    SetStartMessage(dataset_);
}

void JFJochXTCReader::ApplyGain(float adu_per_photon) {
    adu_per_photon_ = adu_per_photon;
    counts_in_photons_ = true;
    dataset_->experiment.Detector().SaturationLimit(static_cast<int64_t>(std::floor(65535.0f / adu_per_photon)));
    SetStartMessage(dataset_);
}

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

void JFJochXTCReader::Close() {
    frames_.clear();
    dataset_.reset();
    photon_energy_eV_.reset();
    pixel_um_.reset();
    adu_per_photon_ = 1.0f;
    counts_in_photons_ = false;
}

template <class Buffer>
CompressedImage JFJochXTCReader::DecodeInto(int64_t image_number, Buffer &buffer,
                                            std::vector<uint8_t> &scratch) const {
    if (image_number < 0 || static_cast<size_t>(image_number) >= frames_.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Image number out of range");
    const auto &f = frames_[image_number];
    buffer.resize(static_cast<size_t>(f.width) * f.height * sizeof(int32_t));
    auto out = reinterpret_cast<int32_t *>(buffer.data());
    xtc::ReadFrame(f, out, scratch);
    if (adu_per_photon_ != 1.0f)
        for (size_t i = 0; i < static_cast<size_t>(f.width) * f.height; i++)
            out[i] = static_cast<int32_t>(std::lround(static_cast<float>(out[i]) / adu_per_photon_));
    return CompressedImage(buffer.data(), buffer.size(), f.width, f.height,
                           CompressedImageMode::Int32, CompressionAlgorithm::NO_COMPRESSION);
}

bool JFJochXTCReader::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;
    std::vector<uint8_t> scratch;
    message.image = DecodeInto(image_number, buffer, scratch);
    message.number = image_number;
    return true;
}

bool JFJochXTCReader::ReadRawImage(int64_t image_number, JFJochReaderRawImage &image) {
    if (image_number < 0 || static_cast<size_t>(image_number) >= frames_.size())
        return false;
    image.image = DecodeInto(image_number, image.image_buffer, image.read_buffer);
    return true;
}

std::vector<SpotToSave> JFJochXTCReader::ReadSpots(int64_t) const {
    return {};   // raw XTC stores no analysis results
}