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

#include "JFJochXTCReader.h"

#include <cmath>
#include <filesystem>
#include <fstream>

#include <nlohmann/json.hpp>

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

namespace {

using Vec = std::array<double, 3>;

double Dot(const Vec &a, const Vec &b) { return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]; }

Vec Cross(const Vec &a, const Vec &b) {
    return {a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], a[0] * b[1] - a[1] * b[0]};
}

Vec ReadVec(const nlohmann::json &j) { return {j.at(0).get<double>(), j.at(1).get<double>(), j.at(2).get<double>()}; }

// A DIALS laboratory vector in the internal frame. DIALS uses the imgCIF laboratory frame (z from
// the sample to the source, y up), which differs from the internal one by a half turn about x.
Coord DialsToInternal(const Vec &v) {
    return {static_cast<float>(v[0]), static_cast<float>(-v[1]), static_cast<float>(-v[2])};
}

// The DIALS experiment list beside the streams, or in the directory above them.
std::string FindExpt(const std::string &stream) {
    const auto dir = std::filesystem::absolute(stream).parent_path();
    for (const auto &d: {dir, dir.parent_path()}) {
        std::error_code ec;
        for (const auto &e: std::filesystem::directory_iterator(d, ec))
            if (e.is_regular_file() && e.path().extension() == ".expt")
                return e.path().string();
    }
    return {};
}

}  // namespace

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

void JFJochXTCReader::ReadGeometry(const std::string &expt) {
    std::ifstream in(expt);
    const auto j = nlohmann::json::parse(in);
    const auto &beam = j.at("beam").at(0);
    const auto &panels = j.at("detector").at(0).at("panels");
    if (panels.size() != 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              expt + " describes a detector of " + std::to_string(panels.size())
                              + " panels; only a single-panel detector is read from XTC");
    const auto &panel = panels.at(0);

    const double wavelength_A = beam.at("wavelength").get<double>();
    // DIALS states the direction from the sample to the source; the beam travels the other way.
    const Vec direction = ReadVec(beam.at("direction"));
    const Vec b = {-direction[0], -direction[1], -direction[2]};
    const Vec fast = ReadVec(panel.at("fast_axis")), slow = ReadVec(panel.at("slow_axis"));
    const Vec origin = ReadVec(panel.at("origin"));   // mm, corner of the first pixel
    const double pixel_mm = panel.at("pixel_size").at(0).get<double>();

    // Where the beam meets the detector plane, in pixels along fast and slow, and how far along the
    // detector normal the plane is.
    const Vec normal = Cross(fast, slow);
    const double t = Dot(origin, normal) / Dot(b, normal);
    const Vec hit = {t * b[0] - origin[0], t * b[1] - origin[1], t * b[2] - origin[2]};
    if (std::fabs(Dot(b, normal)) < 0.9999)
        Logger("XTCReader").Warning("{}: the detector is tilted by {:.2f} deg against the beam; it is read "
                                    "as square-on", expt, std::acos(std::fabs(Dot(b, normal))) * 180.0 / M_PI);

    auto &experiment = dataset_->experiment;
    const auto &f0 = frames_.front();
    DetectorSetup detector = DetDECTRIS(f0.width, f0.height, "XTC camera", {});
    detector.PixelSize_um(static_cast<float>(pixel_mm * 1000.0));
    // 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);
    const auto orientation = DetectorOrientation::Match(DialsToInternal(fast), DialsToInternal(slow));
    if (!orientation.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              expt + ": the panel's fast and slow axes are not a square-on layout this reader supports");
    detector.ImageOrientation(orientation.value());
    detector.MinFrameTime(std::chrono::microseconds(0));
    detector.MinCountTime(std::chrono::microseconds(0));
    detector.ReadOutTime(std::chrono::nanoseconds(0));
    experiment.Detector(detector);

    experiment.BeamX_pxl(static_cast<float>(Dot(hit, fast) / pixel_mm));
    experiment.BeamY_pxl(static_cast<float>(Dot(hit, slow) / pixel_mm));
    experiment.DetectorDistance_mm(static_cast<float>(std::fabs(Dot(origin, normal))));
    experiment.IncidentEnergy_keV(WVL_1A_IN_KEV / static_cast<float>(wavelength_A));

    // The polarization degree is 2f - 1 of the fraction f DIALS states; positive where the plane
    // contains the lab x axis (normal along y), as in the SMV reader.
    if (beam.contains("polarization_fraction") && beam.contains("polarization_normal")) {
        const Coord n = DialsToInternal(ReadVec(beam.at("polarization_normal")));
        const float degree = 2.0f * beam.at("polarization_fraction").get<float>() - 1.0f;
        experiment.PolarizationFactor(std::fabs(n.x) > std::fabs(n.y) ? -degree : degree);
    }
}

void JFJochXTCReader::ReadFiles(const std::string &path) {
    frames_.clear();
    for (const auto &stream: xtc::CollectStreams(path)) {
        auto f = xtc::IndexFrames(stream);
        frames_.insert(frames_.end(), f.begin(), f.end());
    }
    if (frames_.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No camera images found in the XTC of " + path);
    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");

    const std::string expt = FindExpt(frames_.front().file);
    if (expt.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "XTC carries no detector geometry: put the DIALS experiment list (.expt) of "
                              "the run beside the XTC files, or in the directory above them");
    Logger("XTCReader").Info("{} images in XTC, geometry from {}", frames_.size(), expt);

    dataset_ = std::make_shared<JFJochReaderDataset>();
    dataset_->experiment = default_experiment;
    ReadGeometry(expt);
    dataset_->error_value = -1;
    dataset_->experiment.ImagesPerTrigger(static_cast<int64_t>(frames_.size()));
    dataset_->pixel_mask = std::make_shared<const PixelMask>(frames_.front().width, frames_.front().height);

    SetStartMessage(dataset_);
}

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

void JFJochXTCReader::Close() {
    frames_.clear();
    dataset_.reset();
}

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));
    xtc::ReadFrame(f, reinterpret_cast<int32_t *>(buffer.data()), scratch);
    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
}