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

#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"

#ifdef JFJOCH_USE_CUDA

#include <algorithm>
#include <cmath>
#include <complex>
#include <cstring>
#include <random>
#include <vector>

#include "gemmi/asumask.hpp"
#include "gemmi/scaling.hpp"
#include "gemmi/symmetry.hpp"
#include "gemmi/unitcell.hpp"

#include "../rugnux/ModelScaling.h"
#include "../rugnux/ModelScaleGPU.h"

namespace {

constexpr double PI_ = 3.14159265358979323846;

// Scaling points in the reciprocal asymmetric unit of `sg` to d_min, sorted as prepare_points() leaves
// them. Fcalc has random phases and a Wilson fall-off, the mask term is strong at low resolution and
// roughly opposite in phase, and |Fobs| comes from a known overall scale, anisotropic B and solvent pair
// with 5% noise - so the fit has a right answer and a realistic shape of residual.
gemmi::Scaling<float> MakeScaling(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) {
    gemmi::Scaling<float> scaling(cell, &sg);
    scaling.use_solvent = true;
    const gemmi::GroupOps gops = sg.operations();
    const gemmi::ReciprocalAsu asu(&sg);
    const gemmi::Miller lim = cell.get_hkl_limits(d_min);
    std::vector<gemmi::Miller> hkl;
    for (int h = -lim[0]; h <= lim[0]; h++)
        for (int k = -lim[1]; k <= lim[1]; k++)
            for (int l = -lim[2]; l <= lim[2]; l++) {
                const gemmi::Miller m{{h, k, l}};
                if ((h == 0 && k == 0 && l == 0) || cell.calculate_d(m) < d_min || !asu.is_in(m) ||
                    gops.is_systematically_absent(m))
                    continue;
                hkl.push_back(m);
            }
    std::sort(hkl.begin(), hkl.end());

    // B_cart of 25/35/30 A^2 with an off-diagonal term, as B*
    gemmi::SMat33<double> b_cart{25, 35, 30, 4, -3, 2};
    const gemmi::SMat33<double> b_star = b_cart.transformed_by(cell.frac.mat);
    std::mt19937 rng(20260928);
    std::uniform_real_distribution<double> uniform(0.0, 1.0);
    std::normal_distribution<double> gauss(0.0, 1.0);
    const double k_overall = 3.0, k_sol = 0.38, b_sol = 52.0;
    for (const gemmi::Miller &m : hkl) {
        const double stol2 = cell.calculate_stol_sq(m);
        const double phase = 2 * PI_ * uniform(rng);
        const double amp = 200.0 * std::exp(-10.0 * stol2) * std::sqrt(-std::log(1.0 - 0.999 * uniform(rng)));
        const std::complex<double> fc = std::polar(amp, phase);
        const std::complex<double> fm = std::polar(1500.0 * std::exp(-20.0 * stol2) * (0.5 + uniform(rng)),
                                                   phase + PI_ + 0.5 * gauss(rng));
        const std::complex<float> fcf(fc), fmf(fm);
        const std::complex<double> total = std::complex<double>(fcf) + k_sol * std::exp(-b_sol * stol2) * std::complex<double>(fmf);
        const double fobs = k_overall * std::exp(-0.25 * b_star.r_u_r(m)) * std::abs(total) * (1.0 + 0.05 * gauss(rng));
        gemmi::Scaling<float>::Point p{};
        p.hkl = m;
        p.stol2 = stol2;
        p.fcmol = fcf;
        p.fmask = fmf;
        p.fobs = static_cast<float>(std::fabs(fobs));
        p.sigma = static_cast<float>(0.05 * std::fabs(fobs) + 0.5);
        scaling.points.push_back(p);
    }
    return scaling;
}

struct GpuPoints {
    CudaStream stream;
    ModelScaleGPU scale;
    CudaDevicePtr<float2> fcmol, fmask;

    explicit GpuPoints(const gemmi::Scaling<float> &s)
        : scale(stream, s.points.size()), fcmol(s.points.size()), fmask(s.points.size()) {
        std::vector<std::array<int, 3>> hkl;
        std::vector<double> stol2;
        std::vector<float> fobs, sigma;
        std::vector<float2> fc, fm;
        for (const auto &p : s.points) {
            hkl.push_back(p.hkl);
            stol2.push_back(p.stol2);
            fobs.push_back(p.fobs);
            sigma.push_back(p.sigma);
            fc.push_back(make_float2(p.fcmol.real(), p.fcmol.imag()));
            fm.push_back(make_float2(p.fmask.real(), p.fmask.imag()));
        }
        std::vector<std::array<double, 6>> constraints(s.constraint_matrix.begin(), s.constraint_matrix.end());
        double frac[9];
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                frac[3 * i + j] = s.cell.frac.mat[i][j];
        scale.SetPoints(hkl, stol2, fobs, sigma, constraints, frac);
        cudaMemcpy(fcmol, fc.data(), fc.size() * sizeof(float2), cudaMemcpyHostToDevice);
        cudaMemcpy(fmask, fm.data(), fm.size() * sizeof(float2), cudaMemcpyHostToDevice);
    }
};

double MaxAbs(const double b[6]) {
    double m = 0;
    for (int i = 0; i < 6; i++)
        m = std::max(m, std::fabs(b[i]));
    return m;
}

void CheckSameScale(const ModelScaleParams &gpu, const gemmi::Scaling<float> &cpu) {
    const double b_cpu[6] = {cpu.b_star.u11, cpu.b_star.u22, cpu.b_star.u33,
                             cpu.b_star.u12, cpu.b_star.u13, cpu.b_star.u23};
    CHECK(std::fabs(gpu.k_overall - cpu.k_overall) <= 1e-4 * std::fabs(cpu.k_overall));
    for (int i = 0; i < 6; i++)
        CHECK(std::fabs(gpu.b_star[i] - b_cpu[i]) <= 1e-4 * MaxAbs(b_cpu) + 1e-12);
}

struct Case {
    const char *name;
    double a, b, c, alpha, beta, gamma;
    const char *hm;
};

const Case CASES[] = {
    {"triclinic", 41, 47, 53, 82, 97, 104, "P 1"},
    {"monoclinic", 72, 44, 51, 90, 112, 90, "C 1 2 1"},
    {"tetragonal", 64, 64, 81, 90, 90, 90, "P 4"},
    {"hexagonal", 58, 58, 96, 90, 90, 120, "P 6"},
    {"cubic", 92, 92, 92, 90, 90, 90, "P 2 3"},
};

}  // namespace

TEST_CASE("ModelScaleGPU_FitMatchesGemmi", "[ModelValidation][gpu]") {
    if (get_gpu_count() == 0)
        return;
    for (const Case &c : CASES) {
        INFO(c.name);
        gemmi::UnitCell cell(c.a, c.b, c.c, c.alpha, c.beta, c.gamma);
        const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(c.hm);
        gemmi::Scaling<float> cpu = MakeScaling(cell, sg, 2.8);
        REQUIRE(cpu.points.size() > 2000);
        GpuPoints gpu(cpu);

        for (const double k_sol : {0.25, 0.4})
            for (const double b_sol : {30.0, 60.0}) {
                gemmi::Scaling<float> ref = cpu;
                ref.k_sol = k_sol;
                ref.b_sol = b_sol;
                ref.fix_k_sol = true;
                ref.fix_b_sol = true;
                ref.fit_isotropic_b_approximately();
                ref.fit_parameters();
                const ModelScaleParams p = gpu.scale.Fit(gpu.fcmol, gpu.fmask, k_sol, b_sol);
                CheckSameScale(p, ref);

                gemmi::Scaling<float> with_gpu = ref;
                with_gpu.k_overall = p.k_overall;
                with_gpu.b_star = {p.b_star[0], p.b_star[1], p.b_star[2], p.b_star[3], p.b_star[4], p.b_star[5]};
                CHECK(std::fabs(with_gpu.calculate_r_factor() - ref.calculate_r_factor()) < 1e-6);
            }
    }
}

TEST_CASE("ModelScaleGPU_SolventGridMatchesFitModelScale", "[ModelValidation][gpu]") {
    if (get_gpu_count() == 0)
        return;
    for (const Case &c : CASES) {
        INFO(c.name);
        gemmi::UnitCell cell(c.a, c.b, c.c, c.alpha, c.beta, c.gamma);
        const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(c.hm);
        gemmi::Scaling<float> cpu = MakeScaling(cell, sg, 2.8);
        GpuPoints gpu(cpu);

        const ModelScaleReport report = FitModelScale(cpu, {}, 8);
        const ModelSolventFit fit = gpu.scale.FitSolvent(gpu.fcmol, gpu.fmask);
        CHECK(fit.n_grid == report.n_grid);
        CHECK(fit.k_sol == cpu.k_sol);   // the same grid point, so the same double
        CHECK(fit.b_sol == cpu.b_sol);
        CHECK(std::fabs(fit.r - report.r_work_fit) < 1e-6);
        CheckSameScale(fit.scale, cpu);
    }
}

TEST_CASE("ModelScaleGPU_Deterministic", "[ModelValidation][gpu]") {
    if (get_gpu_count() == 0)
        return;
    gemmi::UnitCell cell(72, 44, 51, 90, 112, 90);
    gemmi::Scaling<float> cpu = MakeScaling(cell, *gemmi::find_spacegroup_by_name("C 1 2 1"), 2.5);
    GpuPoints gpu(cpu);
    const ModelSolventFit first = gpu.scale.FitSolvent(gpu.fcmol, gpu.fmask);
    const ModelScaleParams first_fit = gpu.scale.Fit(gpu.fcmol, gpu.fmask, first.k_sol, first.b_sol);
    for (int repeat = 0; repeat < 3; repeat++) {
        const ModelSolventFit again = gpu.scale.FitSolvent(gpu.fcmol, gpu.fmask);
        CHECK(again.k_sol == first.k_sol);
        CHECK(again.b_sol == first.b_sol);
        CHECK(std::memcmp(&again.r, &first.r, sizeof(double)) == 0);
        CHECK(std::memcmp(&again.scale, &first.scale, sizeof(ModelScaleParams)) == 0);
        const ModelScaleParams fit = gpu.scale.Fit(gpu.fcmol, gpu.fmask, first.k_sol, first.b_sol);
        CHECK(std::memcmp(&fit, &first_fit, sizeof(ModelScaleParams)) == 0);
    }
}

#endif