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

#include "SweepLayout.h"

#include <algorithm>
#include <cmath>
#include <cctype>
#include <filesystem>
#include <optional>

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

namespace {

// The difference between two recorded angles, brought into (-180, 180]. A sweep that runs past 360
// starts over at 0 in some writers' headers and keeps counting in others', and this reads both the
// same way. Safe because no gap seen in a deposited series comes close to half a turn - the widest
// in the corpus here is 0.7 degrees - so a folded step is the step and not an aliased one.
double Fold(double d) {
    while (d <= -180.0) d += 360.0;
    while (d > 180.0) d -= 360.0;
    return d;
}

std::string Name(const std::string &path) {
    return std::filesystem::path(path).filename().string();
}

// The frame number a file name ends in. Every one-file-per-image format here numbers its frames that
// way, and the numbering is a second, independent statement of how long the series should be: a
// directory holding fewer files than its own numbering spans is a directory that was not unpacked
// whole, which is worth saying out loud - one staged sweep here was short by 690 frames for months,
// and the only sign of it was a resolution nobody could explain.
std::optional<int64_t> TrailingNumber(const std::string &path) {
    const std::string name = Name(path);
    size_t end = name.find_last_of('.');
    if (end == std::string::npos)
        end = name.size();
    // ".cbf.gz" and the like: step back over as many trailing extensions as there are.
    while (end > 0 && !std::isdigit(static_cast<unsigned char>(name[end - 1]))) {
        const size_t dot = name.find_last_of('.', end - 1);
        if (dot == std::string::npos)
            return {};
        end = dot;
    }
    size_t begin = end;
    while (begin > 0 && std::isdigit(static_cast<unsigned char>(name[begin - 1])))
        begin--;
    if (begin == end)
        return {};
    return std::stoll(name.substr(begin, end - begin));
}

// The first few frames of a list, by name, for a message a user has to act on. All of them would be
// hundreds of lines on the series this exists for.
std::string NameSome(const std::vector<std::string> &paths) {
    const size_t show = std::min<size_t>(paths.size(), 5);
    std::string out;
    for (size_t i = 0; i < show; i++)
        out += (i ? ", " : "") + Name(paths[i]);
    if (paths.size() > show)
        out += fmt::format(" and {} more", paths.size() - show);
    return out;
}

// Whether two readings of the same instrument setting are the same reading. Relative, because what
// counts as the same distance depends on the distance; the bounds are far wider than a read-back
// jitters and far narrower than a real move.
bool Same(double a, double b, double rel_tol, double abs_tol) {
    return std::abs(a - b) <= std::max(abs_tol, rel_tol * std::max(std::abs(a), std::abs(b)));
}

}  // namespace

namespace sweep {

Layout Place(const std::vector<Frame> &frames, const std::string &logger_name) {
    if (frames.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No images to place on a sweep");
    Logger logger(logger_name);

    // ---- one sweep, or several sets of images that happen to share a name?
    //
    // Everything below takes the geometry from the first frame and applies it to all of them, so a
    // series whose headers disagree about where the detector was is not a sweep at all and must not
    // be silently averaged into one. A folder of screening shots is the case that matters: it fails
    // deep inside indexing, as a lattice nobody can explain, when it should fail here by name.
    const Frame &f0 = frames[0];
    std::vector<std::string> bad_distance, bad_beam, bad_wavelength, bad_increment;
    for (const auto &f : frames) {
        if (!Same(f.distance_m, f0.distance_m, 0.005, 1e-6))
            bad_distance.push_back(f.path);
        if (std::abs(f.beam_x_px - f0.beam_x_px) > 2.0 || std::abs(f.beam_y_px - f0.beam_y_px) > 2.0)
            bad_beam.push_back(f.path);
        if (!Same(f.wavelength_A, f0.wavelength_A, 0.001, 1e-9))
            bad_wavelength.push_back(f.path);
        if (!Same(f.increment_deg, f0.increment_deg, 0.01, 1e-6))
            bad_increment.push_back(f.path);
    }
    const auto refuse = [&](const char *what, const std::vector<std::string> &who, double first) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              fmt::format("The images named here are not one sweep: {} differ(s) from "
                                          "{} ({:g}) in {}. Process the sweeps separately.",
                                          what, Name(f0.path), first, NameSome(who)));
    };
    if (!bad_distance.empty()) refuse("the detector distance", bad_distance, f0.distance_m);
    if (!bad_beam.empty()) refuse("the beam centre", bad_beam, f0.beam_x_px);
    if (!bad_wavelength.empty()) refuse("the wavelength", bad_wavelength, f0.wavelength_A);
    if (!bad_increment.empty()) refuse("the oscillation width", bad_increment, f0.increment_deg);

    // ---- the angles, unwrapped so a sweep that passes 360 keeps counting
    std::vector<double> angle(frames.size());
    angle[0] = f0.angle_deg;
    for (size_t i = 1; i < frames.size(); i++)
        angle[i] = angle[i - 1] + Fold(frames[i].angle_deg - angle[i - 1]);

    // ---- a first guess at the rotation step
    //
    // The smallest move between two frames that ARE adjacent in the series, signed with the way it
    // went. Taking the FIRST pair instead - which is what this code used to do - reads a gap as the
    // step and compresses the whole sweep by however much is missing. The header's own
    // Angle_increment is not used for this: it is the oscillation WIDTH, which a series with
    // overlapping or spaced wedges does not step by.
    //
    // A difference smaller than half the oscillation width is not a step but jitter in the recorded
    // angle: no instrument slices finer than it exposes, so wedges overlapping twofold would be a
    // read-back wobble, and taking one as the step would spread the sweep over millions of slots.
    //
    // It is only a guess, and deliberately the crudest one: a single recorded difference, and the
    // one biased furthest low by the read-back noise. It is used below to count STEPS between
    // neighbours, never to place a frame outright.
    const double too_fine = 0.5 * std::abs(f0.increment_deg);
    double guess = 0;
    for (size_t i = 1; i < frames.size(); i++) {
        const double d = Fold(angle[i] - angle[i - 1]);
        if (std::abs(d) > std::max(too_fine, 1e-6) && (guess == 0 || std::abs(d) < std::abs(guess)))
            guess = d;
    }

    Layout out;
    // A series that never turns: a grid scan, a set of stills, or a single image. There is no sweep
    // to place anything on, so the files are the slots and the header's nominal increment stands.
    if (guess == 0) {
        out.files.reserve(frames.size());
        for (const auto &f : frames)
            out.files.push_back(f.path);
        out.start_deg = f0.angle_deg;
        out.increment_deg = f0.increment_deg;
        out.present = frames.size();
        return out;
    }

    // ---- how many steps apart each pair of NEIGHBOURS is
    //
    // Counted from one local difference at a time, so nothing accumulates: a recorded angle carries
    // a few parts in 100000 of read-back noise (these headers are written from 32-bit floats), which
    // is nowhere near half a step for one pair however long the sweep is. Measuring the frame's
    // position against the FIRST frame instead multiplies that noise by the frame number - on a 2700
    // frame sweep the guess above was 1.3e-4 low and by frame 1924 the drift had reached a quarter
    // of a step, so a perfectly regular series was refused as scattered.
    //
    // This is also what carries a sweep past a full turn. Each fold is one step forward, so the
    // count keeps climbing through 360 and a frame taken on the second revolution lands beyond the
    // first, not on top of it. A series that genuinely jumps BACK - two sweeps of one crystal
    // concatenated - folds to a large negative count in a single pair and collides below.
    std::vector<int64_t> slot(frames.size());
    slot[0] = 0;
    for (size_t i = 1; i < frames.size(); i++)
        slot[i] = slot[i - 1] + std::llround(Fold(angle[i] - angle[i - 1]) / guess);

    // ---- the step the whole series agrees on
    //
    // A straight line through (steps, angle): every frame's own header has a say, so the read-back
    // noise averages out instead of one unlucky pair setting the scale for thousands of frames.
    double mean_slot = 0, mean_angle = 0;
    for (size_t i = 0; i < frames.size(); i++) {
        mean_slot += static_cast<double>(slot[i]);
        mean_angle += angle[i];
    }
    mean_slot /= static_cast<double>(frames.size());
    mean_angle /= static_cast<double>(frames.size());
    double cov = 0, var = 0;
    for (size_t i = 0; i < frames.size(); i++) {
        const double ds = static_cast<double>(slot[i]) - mean_slot;
        cov += ds * (angle[i] - mean_angle);
        var += ds * ds;
    }
    // Two frames a step apart give var > 0; var == 0 only if every frame landed on one step, which
    // the guess above has already ruled out.
    const double step = var > 0 ? cov / var : guess;
    const double start = mean_angle - step * mean_slot;

    // ---- every frame on that step, or this is not a rotation series
    std::vector<std::string> off_grid;
    for (size_t i = 0; i < frames.size(); i++)
        if (std::abs(angle[i] - (start + step * static_cast<double>(slot[i]))) > 0.25 * std::abs(step))
            off_grid.push_back(frames[i].path);
    if (!off_grid.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              fmt::format("The images named here do not lie on one rotation series: "
                                          "{} starts at {:.4f} deg, the series steps by {:.4f} deg, "
                                          "and {} sit(s) off that step. Screening images taken at "
                                          "scattered angles are not a sweep.",
                                          Name(f0.path), angle[0], step, NameSome(off_grid)));

    // ---- the slots
    //
    // Numbered from the frame that comes FIRST on the spindle, which is not always the first file:
    // slot 0 is where the goniometer's start angle is, and the files were only ever sorted by name.
    const int64_t first = *std::min_element(slot.begin(), slot.end());
    const int64_t last = *std::max_element(slot.begin(), slot.end());
    out.files.assign(static_cast<size_t>(last - first) + 1, std::string());
    std::vector<std::string> duplicates;
    for (size_t i = 0; i < frames.size(); i++) {
        std::string &at = out.files[static_cast<size_t>(slot[i] - first)];
        if (!at.empty())
            duplicates.push_back(frames[i].path);
        at = frames[i].path;
    }
    // A sweep of more than one full turn reaches the same recorded angle again and is NOT this: it
    // gets there by stepping forward the whole way, so its second revolution occupies fresh slots.
    // What collides here is a series that jumped back - two sweeps of one crystal concatenated.
    if (!duplicates.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              fmt::format("The images named here go back to a point of the rotation "
                                          "an earlier image of the series already took: {}. Two "
                                          "sweeps of the same crystal have to be processed "
                                          "separately.",
                                          NameSome(duplicates)));

    out.start_deg = start + step * static_cast<double>(first);
    out.increment_deg = step;
    out.present = frames.size();

    if (out.present < out.files.size())
        logger.Warning("{} of the {} images the sweep spans are present; the {} missing ones are "
                       "left as gaps, so every image keeps the spindle angle its own header states "
                       "({:.2f} to {:.2f} deg). The merge will be that much less complete.",
                       out.present, out.files.size(), out.files.size() - out.present,
                       out.start_deg,
                       out.start_deg + step * static_cast<double>(out.files.size() - 1));

    // The numbering says the same thing a second way, and says it about the ends of the series too,
    // which the angles cannot: a sweep missing its first and last frames still spans only the angles
    // that are there. Where the two disagree with the file count, the directory is short.
    int64_t lo = 0, hi = 0;
    bool numbered = true;
    for (size_t i = 0; i < frames.size() && numbered; i++) {
        const auto n = TrailingNumber(frames[i].path);
        if (!n.has_value())
            numbered = false;
        else if (i == 0)
            lo = hi = *n;
        else
            lo = std::min(lo, *n), hi = std::max(hi, *n);
    }
    if (numbered && hi - lo + 1 > static_cast<int64_t>(out.present))
        logger.Warning("The file numbering runs {}..{}, which is {} frames, but the directory holds "
                       "{}: {} are not there. If this series should be complete, it was not unpacked "
                       "or copied whole - check the source.",
                       lo, hi, hi - lo + 1, out.present, hi - lo + 1 - static_cast<int64_t>(out.present));
    return out;
}

}  // namespace sweep