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

#include "ModelScaleGPU.h"

#include <algorithm>
#include <cmath>
#include <stdexcept>
#include <string>
#include <utility>

#include "../../common/JFJochException.h"
#include "RigidBodyGPUEngine.h"   // ModelScaleGPUTooFewReflections

namespace {

void cuda_err(cudaError_t val) {
    if (val != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}

// A fixed launch shape, so every point is summed by the same thread in the same order whatever the
// card or the stream: the result is the same bit for bit run to run.
constexpr int THREADS = 128;
constexpr int BLOCKS = 48;
constexpr int SLOTS = 36;      // 1 + 28 + 7: the WSSR, the lower half of J^T J and J^T r for 7 parameters
constexpr int MAX_FITS = 88;   // FitModelScale's coarse pass, the larger of its two

// What a launch sums over the points for one fit.
constexpr int MODE_ISOTROPIC = 0;  // fit_isotropic_b_approximately(): sx, sy, sxx, sxy, n
constexpr int MODE_WSSR = 1;       // compute_wssr()
constexpr int MODE_MATRICES = 2;   // compute_lm_matrices(): WSSR, J^T J, J^T r
constexpr int MODE_R = 3;          // FitModelScale's R: sum |fobs - value|, sum fobs

struct Constraints {
    float row[6][6];
};

// |Fcalc + k_sol exp(-b_sol s^2) Fmask| of every point for each fit, in gemmi's arithmetic: the solvent
// scale cast to float and the sum and its modulus in float, as a complex<float> is. The solvent pair is
// fixed for the whole of a fit, so this is done once per fit rather than at every step.
__global__ void solvent_amplitudes(const double *__restrict__ stol2, const float2 *__restrict__ fcmol,
                                   const float2 *__restrict__ fmask, int n,
                                   const ModelScaleFitState *__restrict__ fits, float *__restrict__ f_abs) {
    const ModelScaleFitState f = fits[blockIdx.y];
    for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) {
        const float solvent = static_cast<float>(f.k_sol * exp(-f.b_sol * stol2[i]));
        const float2 fc = fcmol[i], fm = fmask[i];
        f_abs[static_cast<size_t>(f.column) * n + i] = hypotf(fc.x + solvent * fm.x, fc.y + solvent * fm.y);
    }
}

// The sums of gemmi Scaling<float> for one fit (blockIdx.y) over the points.
template <int NA>
__global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict__ stol2,
                           const float *__restrict__ fobs, const unsigned char *__restrict__ strong,
                           const float *__restrict__ f_abs_all, int n,
                           const ModelScaleFitState *__restrict__ fits, Constraints c, double *__restrict__ partial) {
    constexpr int NS = 1 + NA * (NA + 1) / 2 + NA;   // WSSR, J^T J lower half, J^T r
    const ModelScaleFitState f = fits[blockIdx.y];
    const float *f_abs_fit = f_abs_all + static_cast<size_t>(f.column) * n;
    float uf[6];
    for (int k = 0; k < 6; k++)
        uf[k] = static_cast<float>(f.b_star[k]);
    double acc[NS];
    for (int s = 0; s < NS; s++)
        acc[s] = 0;

    for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) {
        const float f_abs = f_abs_fit[i];
        const double fo = fobs[i];

        if (f.mode == MODE_ISOTROPIC) {
            if (!strong[i])
                continue;
            const double x = stol2[i];
            const double y = logf(static_cast<float>(fo / f_abs));
            acc[0] += x;
            acc[1] += y;
            acc[2] += x * x;
            acc[3] += x * y;
            acc[4] += 1;
            continue;
        }

        if (f.mode == MODE_R) {
            const double hx = hkl[3 * i], hy = hkl[3 * i + 1], hz = hkl[3 * i + 2];
            const double *u = f.b_star;
            const double r_u_r = hx * hx * u[0] + hy * hy * u[1] + hz * hz * u[2]
                                 + 2 * (hx * hy * u[3] + hx * hz * u[4] + hy * hz * u[5]);
            // compute_value(): |F| times the scale cast to float, in float
            const float value = f_abs * static_cast<float>(f.k_overall * exp(-0.25 * r_u_r));
            acc[0] += fabs(fo - value);
            acc[1] += fo;
            continue;
        }

        // Below, the anisotropic factor is taken in float rather than in double as gemmi takes it: double
        // arithmetic is most of the cost of a step on a card with little double throughput. In
        // compute_value() gemmi casts the scale to float anyway, so the WSSR moves by the last bit of some
        // points; the derivatives only set the direction of the next step, which the WSSR then accepts or
        // not.
        const float fx = hkl[3 * i], fy = hkl[3 * i + 1], fz = hkl[3 * i + 2];
        const float r_u_r_f = fx * fx * uf[0] + fy * fy * uf[1] + fz * fz * uf[2]
                              + 2 * (fx * fy * uf[3] + fx * fz * uf[4] + fy * fz * uf[5]);
        const float k_aniso = expf(-0.25f * r_u_r_f);
        if (f.mode == MODE_WSSR) {
            const float value = f_abs * static_cast<float>(f.k_overall * k_aniso);
            const double dy = fo - value;
            acc[0] += dy * dy;
            continue;
        }

        // compute_value_and_derivatives() with k_sol and b_sol fixed, per point in float, summed in double
        const float fe = f_abs * k_aniso;
        const float y = static_cast<float>(f.k_overall) * fe;
        const float du[6] = {-0.25f * y * (fx * fx), -0.25f * y * (fy * fy), -0.25f * y * (fz * fz),
                             -0.5f * y * (fx * fy), -0.5f * y * (fx * fz), -0.5f * y * (fy * fz)};
        double dy_da[NA];
        dy_da[0] = fe;
        for (int j = 1; j < NA; j++) {
            const float *r = c.row[j - 1];
            dy_da[j] = r[0] * du[0] + r[1] * du[1] + r[2] * du[2] + r[3] * du[3] + r[4] * du[4] + r[5] * du[5];
        }
        const double dy = fo - y;
        acc[0] += dy * dy;
        int s = 1;
        for (int j = 0; j < NA; j++)
            for (int k = 0; k <= j; k++)
                acc[s++] += dy_da[j] * dy_da[k];
        for (int j = 0; j < NA; j++)
            acc[s++] += dy * dy_da[j];
    }

    // The block's sum of each slot the mode fills: a fixed shuffle tree within each warp, then the warps
    // in order.
    const int used = f.mode == MODE_MATRICES ? NS : f.mode == MODE_ISOTROPIC ? 5 : f.mode == MODE_R ? 2 : 1;
    __shared__ double warp_sum[THREADS / 32][NS];
    const int lane = threadIdx.x % 32, warp = threadIdx.x / 32;
    for (int s = 0; s < used; s++) {
        double v = acc[s];
        for (int offset = 16; offset > 0; offset /= 2)
            v += __shfl_down_sync(0xffffffffu, v, offset);
        if (lane == 0)
            warp_sum[warp][s] = v;
    }
    __syncthreads();
    if (threadIdx.x < SLOTS) {
        double t = 0;
        if (threadIdx.x < used)
            for (int w = 0; w < THREADS / 32; w++)
                t += warp_sum[w][threadIdx.x];
        partial[(static_cast<size_t>(blockIdx.y) * BLOCKS + blockIdx.x) * SLOTS + threadIdx.x] = t;
    }
}

// Each fit's block partials summed in block order.
__global__ void sum_blocks(const double *__restrict__ partial, double *__restrict__ sums) {
    const int s = threadIdx.x;
    double t = 0;
    for (int b = 0; b < BLOCKS; b++)
        t += partial[(static_cast<size_t>(blockIdx.x) * BLOCKS + b) * SLOTS + s];
    sums[blockIdx.x * SLOTS + s] = t;
}

// gemmi's jordan_solve (levmar.hpp): Gauss-Jordan elimination with partial pivoting, x returned in b.
void JordanSolve(double *a, double *b, int n) {
    for (int i = 0; i < n; i++) {
        int maxnr = -1;
        double amax = 0;
        for (int j = i; j < n; j++) {
            const double aji = std::fabs(a[n * j + i]);
            if (aji > amax) {
                maxnr = j;
                amax = aji;
            }
        }
        if (maxnr == -1) {
            for (int j = i; j < n; j++)
                if (a[n * i + j] != 0. || b[i] != 0.)
                    throw std::runtime_error("Trying to reverse singular matrix. Column " + std::to_string(i) +
                                             " is zeroed.");
            continue;
        }
        if (maxnr != i) {
            for (int j = i; j < n; j++)
                std::swap(a[n * maxnr + j], a[n * i + j]);
            std::swap(b[i], b[maxnr]);
        }
        const double c = 1.0 / a[i * n + i];
        for (int j = i; j < n; j++)
            a[i * n + j] *= c;
        b[i] *= c;
        for (int k = 0; k < n; k++)
            if (k != i) {
                const double d = a[k * n + i];
                for (int j = i; j < n; j++)
                    a[k * n + j] -= a[i * n + j] * d;
                b[k] -= b[i] * d;
            }
    }
}

}  // namespace

// gemmi's LevMar::fit (levmar.hpp), unrolled into the requests it makes of the target, so that many
// fits can share one launch per step. Each call of compute_lm_matrices() or compute_wssr() there is one
// request here; everything between two of them is the same code in the same order.
struct ModelScaleGPU::LevMarRun {
    // gemmi's defaults
    const int eval_limit = 100;
    const double lambda_limit = 1e+15;
    const double stop_rel_change = 1e-5;
    const double lambda_up_factor = 10;
    const double lambda_down_factor = 0.1;

    double lambda = 0.001;
    std::vector<double> initial_a, best_a, trial;
    std::vector<double> alpha, beta;
    double initial_wssr = NAN, wssr = NAN;
    int eval_count = 0;
    int small_change_counter = 0;
    bool done = false;

    int mode = MODE_MATRICES;   // the next request, at `at`
    std::vector<double> at;

    explicit LevMarRun(const std::vector<double> &a) : initial_a(a), best_a(a), at(a) {
        alpha.resize(a.size() * a.size());
        beta.resize(a.size());
    }

    void Take(const double *sums) {
        const size_t na = beta.size();
        if (mode == MODE_MATRICES) {
            int s = 1;
            for (size_t j = 0; j < na; j++)
                for (size_t k = 0; k <= j; k++)
                    alpha[na * j + k] = alpha[na * k + j] = sums[s++];
            for (size_t j = 0; j < na; j++)
                beta[j] = sums[s++];
            if (eval_count == 0) {
                initial_wssr = wssr = sums[0];
                eval_count = 1;
            } else {
                ++eval_count;
                lambda *= lambda_down_factor;
            }
            Propose();
            return;
        }
        const double new_wssr = sums[0];
        ++eval_count;
        if (new_wssr < wssr) {
            const double rel_change = (wssr - new_wssr) / wssr;
            wssr = new_wssr;
            best_a = trial;
            if (wssr == 0) {
                done = true;
                return;
            }
            if (rel_change < stop_rel_change) {
                if (++small_change_counter >= 2) {
                    done = true;
                    return;
                }
            } else {
                small_change_counter = 0;
            }
            mode = MODE_MATRICES;
            at = best_a;
        } else {
            if (lambda > lambda_limit) {
                done = true;
                return;
            }
            lambda *= lambda_up_factor;
            Propose();
        }
    }

    // The damped normal equations J^T J + lambda diag(J^T J), solved for the next trial.
    void Propose() {
        if (eval_limit > 0 && eval_count >= eval_limit) {
            done = true;
            return;
        }
        const size_t na = beta.size();
        std::vector<double> temp_alpha = alpha;
        for (size_t j = 0; j < na; j++)
            temp_alpha[na * j + j] *= (1.0 + lambda);
        trial = beta;
        JordanSolve(temp_alpha.data(), trial.data(), static_cast<int>(na));
        for (size_t i = 0; i < na; i++)
            trial[i] += best_a[i];
        mode = MODE_WSSR;
        at = trial;
    }

    const std::vector<double> &Result() const {
        return wssr < initial_wssr ? best_a : initial_a;
    }
};

size_t ModelScaleGPU::DeviceBytes(size_t max_points) {
    return max_points * (3 * sizeof(int) + sizeof(double) + sizeof(float) + sizeof(unsigned char))
           + MAX_FITS * max_points * sizeof(float)
           + MAX_FITS * (sizeof(ModelScaleFitState) + (BLOCKS + 1) * SLOTS * sizeof(double));
}

ModelScaleGPU::ModelScaleGPU(cudaStream_t stream, size_t max_points)
    : stream_(stream), max_points_(max_points),
      hkl_(3 * max_points), stol2_(max_points), fobs_(max_points), strong_(max_points),
      f_abs_(MAX_FITS * max_points), fits_(MAX_FITS), partial_(static_cast<size_t>(MAX_FITS) * BLOCKS * SLOTS),
      sums_(static_cast<size_t>(MAX_FITS) * SLOTS),
      host_fits_(MAX_FITS), host_sums_(static_cast<size_t>(MAX_FITS) * SLOTS) {}

void ModelScaleGPU::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]) {
    if (hkl.size() > max_points_ || stol2.size() != hkl.size() || fobs.size() != hkl.size() ||
        sigma.size() != hkl.size() || constraints.empty() || constraints.size() > 6)
        throw std::invalid_argument("ModelScaleGPU::SetPoints: inconsistent points");
    n_ = static_cast<int>(hkl.size());
    n_params_ = 1 + static_cast<int>(constraints.size());
    for (size_t j = 0; j < constraints.size(); j++)
        for (int k = 0; k < 6; k++)
            constraints_[j][k] = constraints[j][k];
    std::copy(frac, frac + 9, frac_);

    // fit_isotropic_b_approximately()'s filter: it skips weak reflections
    std::vector<unsigned char> strong(n_);
    n_strong_ = 0;
    for (int i = 0; i < n_; i++) {
        strong[i] = !(fobs[i] < 1 || fobs[i] < sigma[i]);
        n_strong_ += strong[i];
    }
    if (n_ == 0)
        return;
    cuda_err(cudaMemcpyAsync(hkl_, hkl.data(), n_ * 3 * sizeof(int), cudaMemcpyHostToDevice, stream_));
    cuda_err(cudaMemcpyAsync(stol2_, stol2.data(), n_ * sizeof(double), cudaMemcpyHostToDevice, stream_));
    cuda_err(cudaMemcpyAsync(fobs_, fobs.data(), n_ * sizeof(float), cudaMemcpyHostToDevice, stream_));
    cuda_err(cudaMemcpyAsync(strong_, strong.data(), n_, cudaMemcpyHostToDevice, stream_));
    cuda_err(cudaStreamSynchronize(stream_));   // the host vectors go out of scope
}

void ModelScaleGPU::UploadFits(const std::vector<ModelScaleFitState> &fits) {
    if (fits.size() > static_cast<size_t>(MAX_FITS))
        throw std::logic_error("ModelScaleGPU: too many fits in one launch");
    std::copy(fits.begin(), fits.end(), host_fits_.get());
    cuda_err(cudaMemcpyAsync(fits_, host_fits_, fits.size() * sizeof(ModelScaleFitState), cudaMemcpyHostToDevice,
                             stream_));
}

const double *ModelScaleGPU::Sums(const std::vector<ModelScaleFitState> &fits) {
    UploadFits(fits);
    return Reduce(static_cast<int>(fits.size()));
}

const double *ModelScaleGPU::Reduce(int nfits) {
    Constraints c{};
    for (int j = 0; j < 6; j++)
        for (int k = 0; k < 6; k++)
            c.row[j][k] = static_cast<float>(constraints_[j][k]);
    const dim3 grid(BLOCKS, nfits);
    switch (n_params_) {
        case 2: scale_sums<2><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
        case 3: scale_sums<3><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
        case 4: scale_sums<4><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
        case 5: scale_sums<5><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
        case 6: scale_sums<6><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
        case 7: scale_sums<7><<<grid, THREADS, 0, stream_>>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break;
        default: throw std::logic_error("ModelScaleGPU: unexpected number of parameters");
    }
    cuda_err(cudaGetLastError());
    sum_blocks<<<nfits, SLOTS, 0, stream_>>>(partial_, sums_);
    cuda_err(cudaGetLastError());
    cuda_err(cudaMemcpyAsync(host_sums_, sums_, nfits * SLOTS * sizeof(double), cudaMemcpyDeviceToHost, stream_));
    cuda_err(cudaStreamSynchronize(stream_));
    return host_sums_;
}

// fit_isotropic_b_approximately() then fit_parameters() for every fit in `fits`, which start from the
// k_overall and b_star they hold, and are left holding the answer.
void ModelScaleGPU::FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std::vector<ModelScaleFitState> &fits) {
    const size_t nfits = fits.size();
    for (size_t i = 0; i < nfits; i++) {
        fits[i].column = static_cast<int>(i);
        fits[i].mode = MODE_ISOTROPIC;
    }
    UploadFits(fits);
    solvent_amplitudes<<<dim3(BLOCKS, static_cast<unsigned>(nfits)), THREADS, 0, stream_>>>(stol2_, d_fcmol, d_fmask,
                                                                                           n_, fits_, f_abs_);
    cuda_err(cudaGetLastError());
    const double *iso = Reduce(static_cast<int>(nfits));
    for (size_t i = 0; i < nfits; i++) {
        const double *s = iso + i * SLOTS;
        const double sx = s[0], sy = s[1], sxx = s[2], sxy = s[3], n = s[4];
        if (n <= 5)
            continue;
        const double slope = (n * sxy - sx * sy) / (n * sxx - sx * sx);
        const double intercept = (sy - slope * sx) / n;
        const double b_iso = -slope;
        fits[i].k_overall = std::exp(intercept);
        // set_b_overall({b_iso, b_iso, b_iso, 0, 0, 0}): b_star = F B F^T (SMat33::transformed_by)
        const double *m = frac_;
        auto elem = [&](int r, int q) {
            return m[3 * r] * (m[3 * q] * b_iso) + m[3 * r + 1] * (m[3 * q + 1] * b_iso) + m[3 * r + 2] * (m[3 * q + 2] * b_iso);
        };
        const double b[6] = {elem(0, 0), elem(1, 1), elem(2, 2), elem(0, 1), elem(0, 2), elem(1, 2)};
        std::copy(b, b + 6, fits[i].b_star);
    }

    // get_parameters() and set_parameters() with k_sol and b_sol fixed
    const int na = n_params_;
    auto get_parameters = [&](const ModelScaleFitState &f) {
        std::vector<double> a(na);
        a[0] = f.k_overall;
        for (int j = 1; j < na; j++) {
            const double *r = constraints_[j - 1];
            const double *u = f.b_star;
            a[j] = r[0] * u[0] + r[1] * u[1] + r[2] * u[2] + r[3] * u[3] + r[4] * u[4] + r[5] * u[5];
        }
        return a;
    };
    auto set_parameters = [&](ModelScaleFitState &f, const std::vector<double> &a) {
        f.k_overall = a[0];
        std::fill(f.b_star, f.b_star + 6, 0.0);
        for (int j = 1; j < na; j++)
            for (int k = 0; k < 6; k++)
                f.b_star[k] += constraints_[j - 1][k] * a[j];
    };

    std::vector<LevMarRun> runs;
    for (const auto &f : fits)
        runs.emplace_back(get_parameters(f));
    while (true) {
        std::vector<size_t> live;
        std::vector<ModelScaleFitState> request;
        for (size_t i = 0; i < nfits; i++)
            if (!runs[i].done) {
                ModelScaleFitState f = fits[i];
                set_parameters(f, runs[i].at);
                f.mode = runs[i].mode;
                live.push_back(i);
                request.push_back(f);
            }
        if (live.empty())
            break;
        const double *sums = Sums(request);
        for (size_t j = 0; j < live.size(); j++)
            runs[live[j]].Take(sums + j * SLOTS);
    }
    for (size_t i = 0; i < nfits; i++)
        set_parameters(fits[i], runs[i].Result());
}

ModelScaleParams ModelScaleGPU::Fit(const float2 *d_fcmol, const float2 *d_fmask, double k_sol, double b_sol) {
    ModelScaleParams out;
    if (n_ == 0)
        return out;
    std::vector<ModelScaleFitState> fits(1, ModelScaleFitState{1.0, {0, 0, 0, 0, 0, 0}, k_sol, b_sol, 0, 0});
    FitBatch(d_fcmol, d_fmask, fits);
    out.k_overall = fits[0].k_overall;
    std::copy(fits[0].b_star, fits[0].b_star + 6, out.b_star);
    return out;
}

// FitModelScale's grid (ModelScaling.cpp, default ModelScaleBox): a coarse pass over the box and a
// finer one around its winner, every grid point its own fit, the winner the lowest finite R in grid
// order. The fits of a pass run together, one launch per Levenberg-Marquardt step.
ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d_fmask) {
    ModelSolventFit out;
    if (n_ < 20)
        return out;
    if (n_strong_ <= 5)
        throw ModelScaleGPUTooFewReflections();

    const double k_lo = 0.10, k_hi = 0.60, b_lo = 10.0, b_hi = 80.0;   // ModelScaleBox{}
    double best_r = -1;
    auto try_points = [&](std::vector<ModelScaleFitState> &fits) {
        FitBatch(d_fcmol, d_fmask, fits);
        for (auto &f : fits)
            f.mode = MODE_R;
        const double *r = Sums(fits);
        for (size_t i = 0; i < fits.size(); i++) {
            ++out.n_grid;
            const double den = r[i * SLOTS + 1];
            const double ri = den > 0 ? r[i * SLOTS] / den : 1.0;
            if (std::isfinite(ri) && (best_r < 0 || ri < best_r)) {
                best_r = ri;
                out.k_sol = fits[i].k_sol;
                out.b_sol = fits[i].b_sol;
                out.scale.k_overall = fits[i].k_overall;
                std::copy(fits[i].b_star, fits[i].b_star + 6, out.scale.b_star);
            }
        }
    };
    auto point = [](double ks, double bs) {
        return ModelScaleFitState{1.0, {0, 0, 0, 0, 0, 0}, ks, bs, 0, 0};
    };

    std::vector<ModelScaleFitState> coarse;
    for (double ks = k_lo; ks <= k_hi + 1e-9; ks += 0.05)
        for (double bs = b_lo; bs <= b_hi + 1e-9; bs += 10.0)
            coarse.push_back(point(ks, bs));
    try_points(coarse);
    const double k0 = out.k_sol, b0 = out.b_sol;
    const double k_hi2 = std::min(k_hi, k0 + 0.05);
    const double b_hi2 = std::min(b_hi, b0 + 10.0);
    std::vector<ModelScaleFitState> fine;
    for (double ks = std::max(k_lo, k0 - 0.05); ks <= k_hi2 + 1e-9; ks += 0.025)
        for (double bs = std::max(b_lo, b0 - 10.0); bs <= b_hi2 + 1e-9; bs += 5.0)
            fine.push_back(point(ks, bs));
    try_points(fine);
    out.r = best_r;
    return out;
}