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

#pragma once

#include <array>
#include <cstddef>
#include <vector>

#include <cuda_runtime.h>

#include "../indexing/CUDAMemHelpers.h"

// The rigid body's scale fit on the GPU: what RigidBodyTarget::Residuals does with a gemmi
// Scaling<float> (use_solvent, k_sol and b_sol fixed) - fit_isotropic_b_approximately() followed by
// gemmi's Levenberg-Marquardt, and FitModelScale's k_sol/b_sol grid. The sums over the reflections run
// on the device, accumulated in double in a fixed order; the Levenberg-Marquardt control (the damping,
// the 7x7 solve, the stop rules) is gemmi's own, ported line for line and run on the host.

struct ModelScaleParams {
    double k_overall = 1.0;
    double b_star[6] = {0, 0, 0, 0, 0, 0};   // gemmi SMat33 order u11 u22 u33 u12 u13 u23
};

struct ModelSolventFit {
    double k_sol = 0.35, b_sol = 46.0;       // gemmi's Scaling defaults when the fit does not run
    double r = 1.0;                          // ModelScaleReport::r_work_fit
    int n_grid = 0;
    ModelScaleParams scale;                  // the winner's k_overall and b_star
};

// One fit's parameters for one launch, and which sums to take over the points.
struct ModelScaleFitState {
    double k_overall;
    double b_star[6];
    double k_sol, b_sol;
    int mode;
    int column;   // the fit's row of |Fcalc + solvent| on the device
    int weighted; // the robust fit's weights (row `column`) multiply the WSSR and J^T J, J^T r; 0: weight 1
    double c;     // the robust fit's scale of the residuals (ModelScaling.cpp, RobustReweight)
};

class ModelScaleGPU {
public:
    static size_t DeviceBytes(size_t max_points);
    ModelScaleGPU(cudaStream_t stream, size_t max_points);

    // Per zone: the Scaling points in gemmi prepare_points() order, adp_symmetry_constraints(sg) rows
    // and the cell's fractionalization matrix (UnitCell::frac.mat, row-major). n <= max_points.
    void SetPoints(const std::vector<std::array<int, 3>> &hkl,
                   const std::vector<double> &stol2,
                   const std::vector<float> &fobs,
                   const std::vector<float> &sigma,
                   const std::vector<std::array<double, 6>> &constraints,
                   const double frac[9]);

    // fit_isotropic_b_approximately() + fit_parameters() at a fixed solvent pair, from k_overall = 1 and
    // b_star = 0. d_fcmol, d_fmask: one float2 per point on the device.
    ModelScaleParams Fit(const float2 *d_fcmol, const float2 *d_fmask, double k_sol, double b_sol);

    // FitModelScale (ModelScaling.cpp) with the default box, b_sol up to b_hi: the same grid, the same
    // winner. Throws ModelScaleGPUTooFewReflections where FitModelScale would chain the grid points
    // (five or fewer reflections for the isotropic fit).
    // robust: the fine pass by robust least squares, as FitModelScale with ModelScaleBox::robust - the
    // reweighting on the device, the median of each fit's first residuals included.
    // k_start, b_start: ModelScaleBox's - the fine pass around them, no coarse pass (negative: as FitModelScale).
    ModelSolventFit FitSolvent(const float2 *d_fcmol, const float2 *d_fmask, double b_hi = 80.0,
                               bool robust = false, double k_start = -1, double b_start = -1);
    // What a robust FitSolvent takes on the card beyond DeviceBytes(), reserved at its first call.
    static size_t RobustBytes(size_t max_points);

private:
    struct LevMarRun;

    // Queues the copy of `fits` to the device. host_fits_ is written only after the stream has been
    // synchronized, which every Reduce() ends with.
    void UploadFits(const std::vector<ModelScaleFitState> &fits);
    // The sums of the first `nfits` uploaded fits over all points, one row of SLOTS doubles per fit, on
    // the host.
    const double *Reduce(int nfits);
    // The stream's queued work done, polled for: a fit's launches are short and the host waits for every one
    // of them, so sleeping until the device raises the wake-up (set_gpu_blocking_sync) would cost more than
    // the work itself.
    void Wait();
    const double *Sums(const std::vector<ModelScaleFitState> &fits);
    void FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std::vector<ModelScaleFitState> &fits);
    // gemmi's Levenberg-Marquardt for every fit in `fits` from the parameters it holds.
    void LevMarBatch(std::vector<ModelScaleFitState> &fits);
    // The robust rounds of a batch FitBatch() has fitted: reweight, fit again, ROBUST_ROUNDS times.
    void RobustBatch(std::vector<ModelScaleFitState> &fits);

    cudaStream_t stream_;
    size_t max_points_;
    int n_ = 0;
    int n_strong_ = 0;                        // points fit_isotropic_b_approximately() fits on
    int n_params_ = 1;                        // k_overall + one per constraint row
    double constraints_[6][6] = {};
    double frac_[9] = {};

    CudaDevicePtr<int> hkl_;                  // 3 per point
    CudaDevicePtr<double> stol2_;
    CudaDevicePtr<float> fobs_;
    CudaDevicePtr<unsigned char> strong_;
    CudaDevicePtr<float> f_abs_;              // per fit of a batch, per point
    CudaDevicePtr<ModelScaleFitState> fits_;
    CudaDevicePtr<double> partial_;           // per fit, per block, per slot
    CudaDevicePtr<double> sums_;              // per fit, per slot
    CudaHostPtr<ModelScaleFitState> host_fits_;
    CudaHostPtr<double> host_sums_;
    CudaDevicePtr<float> weight_, residual_;  // the robust fit's, per fit of a batch, per point
    CudaDevicePtr<double> median_;            // the robust fit's c, per fit of a batch
    CudaHostPtr<double> host_median_;
};