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

#pragma once

#include "../bragg_prediction/BraggPrediction.h"
#include "../common/DiffractionExperiment.h"
#include "../scale_merge/HKLKey.h"

// =============================================================================
// PixelRefine — reference-driven profile-fit integration + scaling for stills
// =============================================================================
//
// PixelRefine is the still-image integrator: given a reference set of merged
// intensities I_ref (the current best hypothesis for each reflection's full
// intensity), it integrates one image and returns already-scaled intensities. It
// is an *intensity-wise* operation - the detector geometry is taken as fixed (it
// was refined upstream by XtalOptimizer in IndexAndRefine::RefineGeometryIfNeeded);
// PixelRefine does not touch orientation, cell or detector parameters.
//
// The objective is the factored per-reflection likelihood of FACTORED_MODEL.md,
// Terms 1 + 2:
//
//   Term 2 (shape) — for each resolution shell, the tangential profile width R1 is
//     *measured* from the intensity-weighted second moment of the strong spots:
//     R1 = sqrt(2*<eps_t^2>). A second moment is normalised by the total intensity,
//     so it is decoupled from the per-image scale - which is why measuring R1 is
//     stable where *fitting* it (degenerate with G) is not.
//
//   Term 1 (intensity / scaling) — one residual per reflection: the profile-fit
//     amplitude J (using the Term-2 R1) should equal the scaled reference
//         J_model = G * exp(-B/4d^2) * partiality * pol * I_ref,
//     weighted by the model-expected (Fisher) sigma_J. Only the per-image scale G
//     and Debye-Waller B are optimised. Integration and scaling become one objective;
//     the many empty shoebox pixels enter only through J (with ~zero profile weight)
//     instead of dominating a per-pixel loss.
//
// I_ref is NOT refined here - it is a fixed hypothesis for the pass. The intended
// outer loop is EM-like: run PixelRefine on every image against the current I_ref,
// re-merge to a new I_ref, repeat.
//
// Forward model per pixel (raw detector counts, no per-pixel solid-angle/Lorentz
// weighting - same units as the classical integrator):
//   signal = G * I_ref * B_term * P_radial * P_tangential * pol ,  + I_bkg
//   B_term       = exp(-B |q|^2 / 4)                          (Debye-Waller)
//   P_radial     = exp(-eps_r^2 / R0_eff^2)                   (still partiality, <= 1)
//   P_tangential = exp(-eps_t^2 / R1^2) / (pi R1^2)           (area-normalized profile)
// where eps_r / eps_t are the radial / tangential deviations of the pixel from the
// predicted node, and pol is the per-reflection polarization correction.
//
// X-ray bandwidth (optional): a finite bandwidth thickens the Ewald shell radially,
// adding a fixed, resolution-dependent term to R0 that grows like 1/d^2 (the
// pink-beam/DMM signature): R0_eff^2 = R0^2 + (b*lambda)^2/(2 d^4). b = 0 (the
// default) is a monochromatic no-op; set it for DMM-type data, leave it for Si.
// =============================================================================

struct PixelRefineData {
    // --- model state (input as initial guess, output as refined result) ---
    DiffractionGeometry geom;     // fixed (refined upstream by XtalOptimizer)
    CrystalLattice latt;          // fixed
    char centering = 'P';

    double B_factor = 0.0;      // Debye-Waller B (A^2), refined
    double scale_factor = 1.0;  // overall per-image scale G, refined
    double R[2] = {0.005, 0.005}; // R[0] = radial (partiality) width; R[1] = fallback
                                  // tangential profile width before Term 2 measures it (A^-1)
    bool   refine_B = true;       // refine the per-image B-factor along with G

    // Term 2 measures the physical tangential width R1, but the *integration* profile must
    // be generous (XDS-style: integrate over ~6 sigma) or a tight template centred on the
    // prediction sits off the ~0.4 px centroid-floor scatter and underestimates the
    // intensity (validated on the jet R-free: 0.34 with the raw width -> 0.26 at x6). The
    // measured R1 is multiplied by this before use.
    double r1_multiplier = 6.0;

    // Relative X-ray bandwidth (sigma of dlambda/lambda), e.g. ~0.004 for a 1% FWHM
    // DMM, ~1e-4 for Si(111). Adds a resolution-dependent radial broadening to R[0].
    // 0 = monochromatic (the term switches off entirely).
    double bandwidth = 0.0;

    // Per-image scale G is regularized towards 1 with weight sqrt(n_refl/scale_reg_sigma)
    // (mirrors ScaleOnTheFly). Without this the unconstrained G wanders on weakly
    // measured images and 1/G scrambles the cross-image merge. <= 0 disables.
    double scale_reg_sigma = 2.0;

    // Radial Ewald-sphere acceptance band for prediction (A^-1): a reflection is given
    // a shoebox when ||S|-1/lambda| <= this. Widened from the on-sphere default towards
    // the integrator's profile radius so slightly-misaligned high-resolution reflections
    // are still integrated (multiplicity), while the partiality downweights their tails.
    double ewald_dist_cutoff = 0.0020;

    double max_time_s = 5.0;
    int    shoebox_radius = 3; // half-size of the per-reflection signal box
    // Half-size of the local-background sampling box. Background is the MEAN of the ring
    // shoebox_radius < |dx|,|dy| <= bkg_outer_radius (excluding spot cores), like
    // BraggIntegrate2D. Must be > shoebox_radius.
    int    bkg_outer_radius = 6;

    // --- output ---
    std::vector<Reflection> reflections; // profile-fitted, scaled integration result
    bool   solved = false;
    double final_cost = NAN;
    size_t residual_count = 0;
    double cc = NAN;        // per-image CC of scaled intensities vs reference
    int64_t cc_n = 0;       // number of reflections in the CC
};

class PixelRefine {
    const size_t xpixel, ypixel;
    const DiffractionExperiment &experiment;

    const HKLKeyGenerator hkl_key_generator;
    std::map<HKLKey, double> reference_data;

    // Fills the fixed geometry (beam, distance, detector tilt) and the three
    // real-space lattice column vectors from the current model state, for the
    // per-pixel geometry evaluation.
    void BuildParameterBlocks(const PixelRefineData &data,
                              double beam[2], double &dist_mm,
                              double detector_rot[2],
                              double latt_vec0[3], double latt_vec1[3], double latt_vec2[3]) const;
public:
    PixelRefine(const DiffractionExperiment &experiment,
                const std::vector<MergedReflection> &reference);

    // The BraggPrediction is supplied per call (it is mutated): this keeps a single
    // PixelRefine instance usable from several threads, each passing its own prediction
    // buffer. Only `data` is written; PixelRefine state is const. The image is in raw
    // detector counts (masked/saturated pixels carry the type sentinel); background is
    // estimated locally per shoebox from the image itself.
    template<class T>
    void Run(const T *image,
             BraggPrediction &prediction,
             PixelRefineData &data);
};