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

#include "JFJochSFRMReader.h"

#include <cmath>

#include "../common/JFJochException.h"
#include "../common/JFJochMath.h"
#include "../common/Logger.h"
#include "ImgCIF.h"
#include "SweepLayout.h"

namespace {


// The instrument a .sfrm header describes, as an imgCIF axis chain, so that one piece of code turns
// either description into the geometry. The chain is the one APEX writes when it converts frames to
// imgCIF (the pixels come out identical but for the baseline, so the image axes are the same), with
// one thing settled from data rather than taken from it: omega turns in the sense of the 2theta arm.
// An omega scan of APEX's conversion states the opposite increment to the frames it came from, and a
// SMART 6000 omega sweep indexes only that way round. Phi and chi are as the conversion composes them
// from a kappa setting. CENTER is placed as the conversion places it, each coordinate counted from the
// far edge of the stored image: on an APEX II phi sweep read both ways it indexes 97.5% of the frames
// against 90.3% and merges at ISa 17.5 against 2.1, where the image's own .cbf conversion gives 17.
//
// The goniometer is stated as Eulerian angles whatever the head is - a kappa instrument's frames carry
// the omega, chi and phi its kappa setting is equivalent to - so the chain is omega, carrying chi about
// the beam, carrying phi. On the 2theta arm the detector sits DISTANC (plus DETTYPE's face-to-grid
// offset) from the sample, tilted by the DETPAR pitch, roll and yaw. The stored rows run along the
// spindle axis and the stored columns along the direction the arm swings.
imgcif::Description Chain(const sfrm::Header &h, const std::string &path) {
    if (h.center.size() < 2)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, path + " states no CENTER");
    if (!(h.pixels_per_cm_512 > 0.0))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + " states no pixel size (DETTYPE)");
    if (h.detpar[0] != 0.0 || h.detpar[1] != 0.0 || h.detpar[2] != 0.0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + " states non-zero DETPAR offsets, whose meaning this reader does not know");
    // DETTYPE gives pixels per centimetre on a 512-pixel frame; a binned or unbinned frame scales it.
    const double pixel_mm = 10.0 / h.pixels_per_cm_512 * 512.0 / static_cast<double>(h.ncols);
    const double distance_mm = (h.distance_cm + h.grid_offset_cm) * 10.0;
    const double half_cols = static_cast<double>(h.ncols) / 2.0 * pixel_mm;
    const double half_rows = static_cast<double>(h.nrows) / 2.0 * pixel_mm;

    imgcif::Description d;
    d.axes["OMEGA"] = {"rotation", "goniometer", ".", {-1, 0, 0}, {}, h.angles[1], 0};
    d.axes["CHI"] = {"rotation", "goniometer", "OMEGA", {0, 0, 1}, {}, h.angles[3], 0};
    d.axes["PHI"] = {"rotation", "goniometer", "CHI", {1, 0, 0}, {}, h.angles[2], 0};
    d.axes["TWOTHETA"] = {"rotation", "detector", ".", {1, 0, 0}, {}, -h.angles[0], 0};
    d.axes["DX"] = {"translation", "detector", "TWOTHETA", {0, 0, -1}, {}, distance_mm, 0};
    d.axes["YAW"] = {"rotation", "detector", "DX", {-1, 0, 0}, {}, h.detpar[5], 0};
    d.axes["PITCH"] = {"rotation", "detector", "YAW", {0, -1, 0}, {}, h.detpar[3], 0};
    d.axes["ROLL"] = {"rotation", "detector", "PITCH", {0, 0, -1}, {}, h.detpar[4], 0};
    d.axes["H"] = {"translation", "detector", "ROLL", {0, 1, 0}, {}, (h.center[0] * pixel_mm - half_cols), 0};
    d.axes["V"] = {"translation", "detector", "H", {1, 0, 0}, {}, (h.center[1] * pixel_mm - half_rows), 0};
    d.axes["ELEMENT_X"] = {"translation", "detector", "V", {0, 1, 0}, {-half_rows, -half_cols, 0}, 0, 0};
    d.axes["ELEMENT_Y"] = {"translation", "detector", "ELEMENT_X", {1, 0, 0}, {}, 0, 0};
    d.fast = {"ELEMENT_X", h.ncols, pixel_mm / 2.0, pixel_mm};
    d.slow = {"ELEMENT_Y", h.nrows, pixel_mm / 2.0, pixel_mm};

    static const char *AXIS_NAMES[] = {"", "TWOTHETA", "OMEGA", "PHI", "CHI"};
    if (h.axis < 1 || h.axis > 4)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + " scans AXIS " + std::to_string(h.axis)
                              + ", which is not a rotation this reader supports");
    d.axes[AXIS_NAMES[h.axis]].increment = h.increment_deg;
    d.wavelength_A = h.wavelength_A;
    // The polarization APEX's conversion states for the same frame: 1 - cos^2 of the monochromator's
    // scattering angle (FILTER2), in the horizontal plane. A tube with mirrors and no monochromator
    // states 0, unpolarized - these are laboratory frames, and the synchrotron default does not apply.
    const double cos_mono = std::cos(h.monochromator_2theta_deg * PI / 180.0);
    d.polarization_ratio = 1.0 - cos_mono * cos_mono;
    d.polarization_norm_deg = 0;
    d.overload = h.saturation;
    d.exposure_s = h.exposure_s;
    d.detector = h.detector;
    return d;
}

}  // namespace

bool JFJochSFRMReader::CanRead(const std::string &path) {
    return sfrm::CanRead(path);
}

void JFJochSFRMReader::ReadFiles(const std::string &path) {
    files_ = sfrm::CollectSweep(path);
    if (files_.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No SFRM images found for " + path);

    header0_ = sfrm::ReadHeader(files_[0]);
    if (!(header0_.wavelength_A > 0.0))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              files_[0] + " states no wavelength in its SFRM header");
    const auto geometry0 = imgcif::Evaluate(Chain(header0_, files_[0]));
    const auto plane0 = imgcif::PlaneOf(geometry0, files_[0]);

    dataset_ = std::make_shared<JFJochReaderDataset>();
    dataset_->experiment = default_experiment;

    DetectorSetup detector = DetDECTRIS(header0_.ncols, header0_.nrows, header0_.detector, {});
    detector.PixelSize_um(plane0.pixel_mm * 1000.0f);
    // A CCD has no sensor depth to correct for: the phosphor converts at the surface.
    detector.SensorThickness_um(0);
    // CCDPARM[4] is the count at which the CCD saturates; the stored container (with its overflow
    // tables) holds far more, so it is the only limit there is.
    if (header0_.saturation > 0.0)
        detector.SaturationLimit(static_cast<int64_t>(header0_.saturation));
    else
        Logger("SFRMReader").Warning("{} states no saturation (CCDPARM), so no pixel will be called "
                                     "saturated.", files_[0]);
    detector.BitDepthImage(32);
    detector.ImageOrientation(plane0.orientation);
    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(plane0.beam_x_pxl);
    dataset_->experiment.BeamY_pxl(plane0.beam_y_pxl);
    dataset_->experiment.DetectorDistance_mm(plane0.distance_mm);
    dataset_->experiment.PoniRot1_rad(plane0.rot1).PoniRot2_rad(plane0.rot2).PoniRot3_rad(plane0.rot3);
    dataset_->experiment.IncidentEnergy_keV(WVL_1A_IN_KEV / static_cast<float>(header0_.wavelength_A));
    dataset_->experiment.PolarizationFactor(imgcif::PolarizationFactor(Chain(header0_, files_[0])));
    if (header0_.exposure_s > 0.0)
        dataset_->experiment.FrameTime(
            std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::duration<double>(header0_.exposure_s)),
            std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::duration<double>(header0_.exposure_s)));

    // Where every image sits on the spindle, from its own header - and the check that every frame is
    // of the same sweep: the same axis turning about the same direction. A run whose fixed angles moved
    // part-way is more than one sweep, and reading it as one would put every frame after the move at
    // the wrong orientation.
    std::vector<sweep::Frame> frames(files_.size());
    sweep::ForEachInOrder(files_.size(), [&](size_t i) {
        const auto h = sfrm::ReadHeader(files_[i]);
        const auto g = imgcif::Evaluate(Chain(h, files_[i]));
        if (g.scan_axis != geometry0.scan_axis
                || (imgcif::ToInternal(g.scan_direction) - imgcif::ToInternal(geometry0.scan_direction)).Length() > 1e-4f)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  files_[i] + " turns " + g.scan_axis + " about a different direction from "
                                  + files_[0] + ": the goniometer was reset between them, so the run is "
                                    "more than one sweep");
        const auto plane = imgcif::PlaneOf(g, files_[i]);
        frames[i] = {files_[i], g.scan_start_deg, g.scan_increment_deg, plane.distance_mm * 1e-3,
                     plane.beam_x_pxl, plane.beam_y_pxl, h.wavelength_A};
    });

    const auto layout = sweep::Place(frames, "SFRMReader");
    files_ = layout.files;
    dataset_->experiment.Goniometer(imgcif::SweepAxis(geometry0.scan_axis, layout.start_deg, layout.increment_deg,
                                                    geometry0.scan_direction));

    dataset_->error_value = -1;
    dataset_->experiment.ImagesPerTrigger(static_cast<int64_t>(files_.size()));

    // Outside the phosphor's active area a frame reads below its baseline offset - the stored value
    // is zero there in every frame - so frame 0's negative counts are the pixels no frame can be
    // trusted at, exactly as a PILATUS marks its gaps.
    std::vector<int32_t> first;
    sfrm::Read(files_[0], first);
    std::vector<uint32_t> mask(first.size(), 0);
    for (size_t i = 0; i < first.size(); i++)
        if (first[i] < 0)
            mask[i] = 1;
    dataset_->pixel_mask = std::make_shared<const PixelMask>(mask);

    SetStartMessage(dataset_);
}

uint64_t JFJochSFRMReader::GetNumberOfImages() const {
    return files_.size();
}

void JFJochSFRMReader::Close() {
    files_.clear();
    dataset_.reset();
}

template <class Buffer>
CompressedImage JFJochSFRMReader::DecodeInto(int64_t image_number, Buffer &buffer,
                                             std::vector<uint8_t> &scratch) const {
    if (!HasImage(image_number))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No image at this point of the sweep");

    const size_t npixel = static_cast<size_t>(header0_.nrows) * static_cast<size_t>(header0_.ncols);
    buffer.resize(npixel * sizeof(int32_t));
    const auto h = sfrm::ReadInto(files_[image_number], reinterpret_cast<int32_t *>(buffer.data()), npixel,
                                  scratch);
    if (h.nrows != header0_.nrows || h.ncols != header0_.ncols)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "SFRM image size differs from the first image of the sweep");

    return CompressedImage(buffer.data(), buffer.size(),
                           static_cast<size_t>(header0_.ncols), static_cast<size_t>(header0_.nrows),
                           CompressedImageMode::Int32, CompressionAlgorithm::NO_COMPRESSION);
}

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

bool JFJochSFRMReader::HasImage(int64_t image_number) const {
    return image_number >= 0 && static_cast<size_t>(image_number) < files_.size()
           && !files_[image_number].empty();
}

bool JFJochSFRMReader::ReadRawImage(int64_t image_number, JFJochReaderRawImage &image) {
    if (!HasImage(image_number))
        return false;
    image.image = DecodeInto(image_number, image.image_buffer, image.read_buffer);
    return true;
}

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