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

#pragma once

#include <array>
#include <cstdint>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "../../common/DiffractionGeometry.h"
#include "../../common/CrystalLattice.h"
#include "../../common/GoniometerAxis.h"
#include "../../common/UnitCell.h"
#include "../../common/Logger.h"
#include "../IntegrationOutcome.h"
#include "gemmi/symmetry.hpp"

// Post-integration geometry refinement for rotation data. Unlike the at-indexing XtalOptimizer, this runs
// AFTER integration/merge, where each reflection has an OBSERVED rocking centroid phi_obs (the intensity-
// weighted mean goniometer angle over the frames it spans) and an observed spot position. One JOINT fit
// refines the crystal (orientation, cell, rotation axis) and the detector (distance, beam centre) together
// against two residuals:
//   * an Ewald excitation residual evaluated at phi_obs (distance-independent) -> pins the absolute cell
//     scale that the positional residual leaves degenerate with the distance. Because phi_obs is the real
//     rocking angle (not a frame centre) it is unbiased.
//   * the positional detector<->reciprocal residual at each partial's observed spot -> pins the distance.
// Where there are more reflections than the fit's caps, the strongest by I/sigma are the ones kept;
// inside the fit every reflection carries the same weight.

struct PostRefineResult {
    bool ok = false;
    DiffractionGeometry geom;                 // refined (distance / beam left as configured)
    UnitCell cell{};                          // refined unit cell
    int events_used = 0;
    int obs_used = 0;
    double distance_before_mm = 0.0, distance_after_mm = 0.0;
    double beam_x_before_px = 0.0, beam_x_after_px = 0.0;   // refined beam centre (GEOM mode)
    double beam_y_before_px = 0.0, beam_y_after_px = 0.0;
    // The joint fit is one decision - the crystal and the detector are refined together and committed
    // together - so these two are always equal. Both are kept because the report names them separately.
    bool cell_refined = false;                // the joint fit passed cross-validation (crystal half)
    bool detector_refined = false;            // the joint fit passed cross-validation (detector half)
    // The fit moved the distance or a cell length by more than one step of its trust region (see
    // STEP_BOUND in PostRefine.cpp), so it walked rather than settling beside the geometry it started
    // from. Such a move is not refused for its size - it is RATIFIED by re-integrating and re-indexing
    // at it, which a second lattice does not survive and a wrong header does (see RunAllPasses).
    bool large_move = false;
    // What was committed is the fit with the distance HELD at the header - everything else as free as
    // ever - because freeing it did not lower the held-out excitation residual by more than that
    // residual's noise (see the joint fit in PostRefine.cpp). The report names it.
    bool distance_held = false;
    // The held-out residual at the geometry the fit STARTED from - what the integration this fit was
    // handed actually realises there, before the fit moves anything - and its standard error over the
    // held-out values. Two passes at two geometries are compared on this, not on the fit's own
    // prediction of where it is going (see the geometry walk in RunAllPasses). NaN where the fit had
    // too few observations to measure it.
    double held_out_before = std::numeric_limits<double>::quiet_NaN();
    double held_out_before_se = std::numeric_limits<double>::quiet_NaN();
    // The cell the fit started from: the one this pass's own indexing refined. Beside `cell` (what the
    // fit committed), and against the next pass's cell_before, it says whether re-indexing at the
    // committed geometry ratifies the committed cell.
    UnitCell cell_before{};
    // Which test refused the fit and what it wanted, where it was refused; empty where it committed.
    // Nothing else says it - the geometry that comes out of a refusal is the nominal one whatever
    // refused it, so a silent refusal cannot be told from a fit that never wanted to move.
    std::string refused_reason;
    // GONIOMETER ROTATION SCALE: the factor by which the stage actually turned relative to the angle
    // stored in the file (which is the COMMANDED value, hence a stage calibration error is invisible in
    // the header). Fitted after the joint fit as a single free parameter, with the crystal and the axis
    // direction held at their committed values. Always the fitted value; 1.0 = header and stage agree.
    double rotation_scale = 1.0;
    // Whether the run ACTED on it: set by the caller where the scale, walked to its fixed point, was
    // adopted on the evidence of the whole sweep (rugnux WalkRotationScale) - never by the fit itself.
    bool rotation_scale_suspect = false;
};

struct PostRefineSettings {
    gemmi::CrystalSystem crystal_system = gemmi::CrystalSystem::Triclinic;
    bool refine_geometry = false;     // XtalOptimizer-equivalent: the crystal (from phi_obs and the observed
                                      // spot positions) and the detector distance + beam centre, in one
                                      // cross-validated fit. The only supported refinement mode.
    double excitation_weight = 1.0;   // weight of the phi/excitation residual vs the positional one
    // The geometry handed in still has the file's distance, so "the header distance is right" is a
    // hypothesis the fit can hold and test (PostRefineResult::distance_held). Once a run has moved off
    // the header - a walk, see RunAllPasses - that hypothesis is already refuted and the distance is
    // simply free.
    bool distance_at_header = true;
    int min_events = 50;
    int num_threads = 1;
    // An independent measurement of the beam centre in pixels, where the run has one (rugnux's pre-scan
    // fit to the isotropy of the scattered background, which runs on every rotation run). The fit bounds
    // how far it may move the beam from whichever of this and the nominal centre is NEARER; see
    // BEAM_BOUND_PXL in PostRefine.cpp for why the nominal centre alone is not enough to bound it.
    std::optional<std::array<double, 2>> measured_beam_px;
};

// One integrated partial, flattened across all images. Kept as narrow as the sort and the event
// split allow: on a large cell this array is gigabytes, and the scatter and every level of the
// per-bucket sort move all of it. The goniometer angle is not stored - it is a function of the
// image number alone, and is rebuilt from it where it is needed.
struct PostRefinePartial {
    // Sixteen bits are plenty for a Miller index - |h| <= a / d_min, in the hundreds even on the
    // longest axis at atomic resolution - and here four bytes of every partial are four bytes of an
    // array that is gigabytes on a large cell.
    int16_t h, k, l;
    float img;
    float I, sigma;
    float obs_x, obs_y;  // observed spot centroid (pixels); NAN if the box sum found no centroid
};

// The gathered form the fit runs on: every usable partial, bucketed by h (bstart = bucket starts,
// order = the buckets largest first) and sorted by (h, k, l, img) inside each bucket. Gathered
// apart from the fit itself so a caller that is finished with the per-image reflection vectors can
// hand them back before the fit's own allocations arrive - on a fine-sliced long axis the vectors
// are several times this array.
struct PostRefineObservations {
    std::unique_ptr<PostRefinePartial[]> pts;
    size_t n_pts = 0;
    std::vector<int32_t> bstart;   // size H+1: where each h bucket starts in pts
    std::vector<int> order;        // size H: bucket indices, largest bucket first
};

// Gather and sort the usable partials out of the per-image outcomes. release_reflections: each
// image's reflection vector is freed the moment it has been read - only for a caller that will
// never touch them again (the rotation geometry pre-pass, whose reflections are never written).
// An allocation failure returns an empty gather (n_pts = 0), which the fit refuses - the same
// silent refusal the fit gives every other failure.
PostRefineObservations GatherPostRefineObservations(std::vector<IntegrationOutcome> &outcomes,
                                                    size_t nthreads, bool release_reflections,
                                                    Logger &logger);

// nominal_geom / reference_latt: the current detector geometry and the phi=0 reference lattice (orientation
// + cell) from rotation indexing. obs: the gathered per-image integrated reflections (observed_x/y, I,
// sigma, image_number), from GatherPostRefineObservations. axis: the goniometer. Returns ok=false
// (geometry untouched) on failure.
PostRefineResult PostRefineRotationGeometry(PostRefineObservations obs,
                                            const GoniometerAxis &axis,
                                            const DiffractionGeometry &nominal_geom,
                                            const CrystalLattice &reference_latt,
                                            const PostRefineSettings &settings,
                                            Logger &logger);

// Whether the held-out residual `after` realises at its geometry is below the one `before` realises at
// its own by more than the standard error of the difference - the cross-validation noise of the two
// measurements. False where either was not measured.
bool HeldOutResidualFell(const PostRefineResult &before, const PostRefineResult &after);

// Whether re-indexing at the geometry `fit` committed pushed the committed cell back: on the cell length
// the fit moved most, the cell `next` (the fit of the pass run at that geometry) started from lies on the
// side the fit moved away from. Returns +-(axis + 1), signed as the fit's move, where it did; 0 where the
// re-indexing followed the fit, or the fit committed nothing.
int ReindexPushesCellBack(const PostRefineResult &fit, const PostRefineResult &next);