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

#include "RefinementXray.h"

#include <algorithm>
#include <cmath>
#include <map>
#include <mutex>

#include <fftw3.h>

#include "gemmi/dencalc.hpp"
#include "gemmi/it92.hpp"
#include "gemmi/math.hpp"

#include "../../common/Logger.h"
#include "../../common/ParallelFor.h"
#include "ModelFFT.h"
#include "ModelGrid.h"
#include "RefinementMath.h"

namespace {

using Table = gemmi::IT92<float>;

// The free reflections a likelihood bin is fitted on, at least.
constexpr size_t FREE_PER_BIN = 100;
constexpr size_t MAX_BINS = 30;

// The density calculator the refinement grids atoms with: DensityCalculator's settings for the grid of
// RigidBodyZoneGrid(), and the given blur.
gemmi::DensityCalculator<Table, float> Calculator(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                                  double d_min, double blur) {
    gemmi::DensityCalculator<Table, float> dc;
    dc.d_min = d_min;
    dc.rate = 1.5;   // RigidBodyRefine.cpp's GRID_RATE, which RigidBodyZoneGrid sizes the grid with
    dc.grid.unit_cell = cell;
    dc.grid.spacegroup = &sg;
    dc.blur = blur;
    return dc;
}

// -ln L summed over rows `idx` at D and beta.
double BinNegLogLikelihood(const RefinementReflections &refl, const std::vector<double> &fc,
                           const std::vector<size_t> &idx, double d, double beta) {
    double sum = 0, dfc = 0;
    for (size_t i : idx) {
        const double var = refl.epsilon[i] * beta + double(refl.sigma[i]) * refl.sigma[i];
        sum += refine_math::AmplitudeNegLogLikelihood<double>(refl.fo[i], fc[i], d, var, refl.centric[i], dfc);
    }
    return sum;
}

// The minimum of f over [lo, hi] by golden-section search.
template <typename F>
double GoldenMinimum(F f, double lo, double hi, int iterations) {
    const double g = 0.5 * (std::sqrt(5.0) - 1.0);
    double a = lo, b = hi;
    double c = b - g * (b - a), d = a + g * (b - a);
    double fc = f(c), fd = f(d);
    for (int i = 0; i < iterations; i++) {
        if (fc < fd) {
            b = d;
            d = c;
            fd = fc;
            c = b - g * (b - a);
            fc = f(c);
        } else {
            a = c;
            c = d;
            fc = fd;
            d = a + g * (b - a);
            fd = f(d);
        }
    }
    return 0.5 * (a + b);
}

class RefinementEngineCPU : public RefinementEngine {
public:
    RefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
                            const RefinementReflections &refl, size_t nthreads)
        : cell_(cell), sg_(sg), d_min_(d_min), refl_(refl), nthreads_(nthreads),
          grid_(RigidBodyZoneGrid(cell, sg, d_min)), composition_(grid_, d_min, refl.hkl),
          orbit_leaders_(OrbitLeaders(grid_, nthreads)) {}

    void SetBlur(double blur) override { blur_ = blur; }
    void SetRowModel(const XrayRowModel &rows) override { rows_ = rows; }
    bool OnGPU() const override { return false; }

    void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
        std::vector<std::complex<double>> fd;
        Compose(model, fd);
        f.assign(fd.begin(), fd.end());
    }

    void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
        gemmi::Grid<float> mask = grid_;
        PutMaskOnGrid(mask, model, orbit_leaders_, nthreads_);
        const gemmi::FPhiGrid<float> fm = MapToFPhi(mask);
        f.clear();
        for (const gemmi::Miller &h : refl_.hkl)
            f.push_back(fm.get_value_by_hkl(h));
    }

    double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) override {
        std::vector<std::complex<double>> f;
        Compose(model, f);
        const size_t n = f.size();
        // The target and, per reflection, dT/dF of the atoms' structure factor (dT/dRe F + i dT/dIm F).
        std::vector<std::complex<double>> dt_df(n, 0.0);
        double sum = 0;
        for (size_t m = 0; m < n; m++) {
            if (refl_.free[m])
                continue;
            const double k = rows_.k_total[m];
            const std::complex<double> fm = k * (f[m] + std::complex<double>(rows_.fmask[m]));
            const double fc = std::abs(fm);
            double dfc = 0;
            sum += refine_math::AmplitudeNegLogLikelihood<double>(refl_.fo[m], fc, rows_.d[m], rows_.var[m],
                                                                 refl_.centric[m], dfc);
            if (fc > 0)
                dt_df[m] = k * dfc * fm / fc;
        }
        if (grad == nullptr)
            return sum;

        // The FFT method's map, dT/drho(x) = (V/N) Re sum_m sum_ops conj(dT/dF_m) scale_m exp(2 pi i h.t)
        // exp(2 pi i (hR).x), as the inverse real transform of a half-l coefficient grid: a term A at index k
        // goes in as A/2 at k, or as conj(A)/2 at -k where k is in the half not stored; on the planes l = 0 and
        // l = nw/2, which hold both, as both.
        const int nu = grid_.nu, nv = grid_.nv, nw = grid_.nw, nwh = nw / 2 + 1;
        const size_t ncomplex = static_cast<size_t>(nu) * nv * nwh;
        std::complex<float> *in = reinterpret_cast<std::complex<float> *>(fftwf_alloc_complex(ncomplex));
        float *out = fftwf_alloc_real(static_cast<size_t>(nu) * nv * nw);
        std::fill(in, in + ncomplex, std::complex<float>(0, 0));
        const std::vector<SymmetryComposition::Term> &terms = composition_.Terms();
        const size_t ops = composition_.Ops();
        auto add = [&](int ku, int kv, int kw, std::complex<double> value) {
            in[gemmi::modulo(ku, nu) + static_cast<size_t>(nu) * (gemmi::modulo(kv, nv) + static_cast<size_t>(nv) * kw)] +=
                std::complex<float>(value);
        };
        for (size_t m = 0; m < n; m++) {
            if (dt_df[m] == 0.0)
                continue;
            const double scale = composition_.Centring() * std::exp(blur_ * 0.25 * composition_.InvD2()[m]);
            const std::complex<double> c = std::conj(dt_df[m]) * scale;
            for (size_t o = 0; o < ops; o++) {
                const SymmetryComposition::Term &t = terms[m * ops + o];
                const std::complex<double> a = 0.5 * c * t.phase;
                const gemmi::Miller &k = t.k;
                if (k[2] > 0 && 2 * k[2] != nw)
                    add(k[0], k[1], k[2], a);
                else if (k[2] < 0 && -2 * k[2] != nw)
                    add(-k[0], -k[1], -k[2], std::conj(a));
                else {
                    const int l = k[2] < 0 ? -k[2] : k[2];
                    add(k[0], k[1], l, a);
                    add(-k[0], -k[1], l, std::conj(a));
                }
            }
        }
        MapFromFPhiExecute(nu, nv, nw, in, out);
        gemmi::Grid<float> map = grid_;
        const float norm = static_cast<float>(cell_.volume / static_cast<double>(map.point_count()));
        for (size_t i = 0; i < map.data.size(); i++)
            map.data[i] = out[i] * norm;
        fftwf_free(in);
        fftwf_free(out);
        AtomGradientsFromMap(map, model, blur_, nthreads_, *grad);
        return sum;
    }

private:
    void Compose(const gemmi::Model &model, std::vector<std::complex<double>> &f) const {
        gemmi::DensityCalculator<Table, float> dc = Calculator(cell_, sg_, d_min_, blur_);
        PutModelDensityOnGrid(dc, model, {}, nthreads_);   // one copy: the symmetry is composed below
        composition_.Compose(MapToFPhi(dc.grid), blur_, f, nullptr, nthreads_);
    }

    gemmi::UnitCell cell_;
    const gemmi::SpaceGroup &sg_;
    double d_min_;
    const RefinementReflections &refl_;
    size_t nthreads_;
    gemmi::Grid<float> grid_;
    SymmetryComposition composition_;
    std::vector<size_t> orbit_leaders_;
    double blur_ = 0;
    XrayRowModel rows_;
};

}  // namespace

RefinementReflections PrepareRefinementReflections(const std::vector<MergedReflection> &merged,
                                                   const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                                   double d_min) {
    const gemmi::GroupOps gops = sg.operations();
    const gemmi::ReciprocalAsu asu(&sg);
    // ASU, last one wins where a reflection is listed twice (both hands of an anomalous merge)
    std::map<gemmi::Miller, const MergedReflection *> by_hkl;
    for (const MergedReflection &r : merged) {
        if (!std::isfinite(r.F) || r.F < 0)
            continue;
        gemmi::Miller h{{r.h, r.k, r.l}};
        if (!asu.is_in(h))
            h = asu.to_asu(h, gops).first;
        by_hkl[h] = &r;
    }
    std::vector<gemmi::Miller> hkl;
    for (const auto &[h, r] : by_hkl)
        hkl.push_back(h);
    const SymmetryComposition composition(RigidBodyZoneGrid(cell, sg, d_min), d_min, hkl);
    RefinementReflections out;
    for (size_t i = 0; i < hkl.size(); i++) {
        if (composition.Row()[i] < 0)
            continue;
        const MergedReflection &r = *by_hkl[hkl[i]];
        out.hkl.push_back(hkl[i]);
        out.fo.push_back(r.F);
        out.sigma.push_back(std::isfinite(r.sigmaF) ? r.sigmaF : 0.0f);
        out.free.push_back(r.rfree_flag);
        out.centric.push_back(gops.is_reflection_centric(hkl[i]));
        out.epsilon.push_back(static_cast<float>(gops.epsilon_factor(hkl[i])));
        out.inv_d2.push_back(cell.calculate_1_d2(hkl[i]));
    }
    return out;
}

BulkSolventMask::BulkSolventMask(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
                                 size_t nthreads)
    : nthreads_(nthreads),
      spacing_(std::max(std::min(d_min / 3.0, MASK_SPACING_MAX), std::cbrt(cell.volume / MASK_POINTS_MAX))) {
    grid_.unit_cell = cell;
    grid_.spacegroup = &sg;
    grid_.set_size_from_spacing(spacing_, gemmi::GridSizeRounding::Up);
    orbit_leaders_ = OrbitLeaders(grid_, nthreads);
}

void BulkSolventMask::Fmask(const gemmi::Model &model, const std::vector<gemmi::Miller> &hkl,
                            std::vector<std::complex<float>> &f) const {
    gemmi::Grid<float> mask = grid_;
    PutMaskOnGrid(mask, model, orbit_leaders_, nthreads_);
    const gemmi::FPhiGrid<float> fm = MapToFPhi(mask);
    f.clear();
    for (const gemmi::Miller &h : hkl)
        f.push_back(fm.get_value_by_hkl(h));
}

void EstimateLikelihoodParameters(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel,
                                  XrayRowModel &rows, size_t nthreads) {
    const size_t n = refl.Size();
    std::vector<double> fc(n);
    for (size_t i = 0; i < n; i++)
        fc[i] = std::abs(fmodel[i]);
    std::vector<size_t> free;
    for (size_t i = 0; i < n; i++)
        if (refl.free[i] && fc[i] > 0)
            free.push_back(i);
    // Too few free reflections to bin: the working set, which says the model is better than it is, but
    // still gives the likelihood a scale.
    if (free.size() < FREE_PER_BIN)
        for (size_t i = 0; i < n; i++)
            if (!refl.free[i] && fc[i] > 0)
                free.push_back(i);
    std::sort(free.begin(), free.end(), [&](size_t a, size_t b) { return refl.inv_d2[a] < refl.inv_d2[b]; });
    const size_t bins = std::clamp<size_t>(free.size() / FREE_PER_BIN, 1, MAX_BINS);

    std::vector<double> centre(bins), bin_d(bins), bin_beta(bins);
    // each bin is fitted on its own
    ParallelFor(static_cast<int>(bins), nthreads, [&](int b) {
        const std::vector<size_t> idx(free.begin() + b * free.size() / bins, free.begin() + (b + 1) * free.size() / bins);
        double s2 = 0, sfo = 0, sfc2 = 0, sfofc = 0;
        for (size_t i : idx) {
            s2 += refl.inv_d2[i];
            sfo += double(refl.fo[i]) * refl.fo[i] / refl.epsilon[i];
            sfc2 += fc[i] * fc[i] / refl.epsilon[i];
            sfofc += refl.fo[i] * fc[i] / refl.epsilon[i];
        }
        centre[b] = s2 / static_cast<double>(idx.size());
        double d = sfc2 > 0 ? std::max(sfofc / sfc2, 0.05) : 1.0;
        double beta = std::max((sfo - d * d * sfc2) / static_cast<double>(idx.size()), 1e-3 * sfo / idx.size());
        // Alternating one-dimensional minimisations of -ln L over ln D and ln beta.
        for (int round = 0; round < 4; round++) {
            const double ld = GoldenMinimum([&](double x) { return BinNegLogLikelihood(refl, fc, idx, std::exp(x), beta); },
                                            std::log(d) - 2.0, std::log(std::min(d * 7.4, 3.0)), 40);
            d = std::exp(ld);
            const double lb = GoldenMinimum([&](double x) { return BinNegLogLikelihood(refl, fc, idx, d, std::exp(x)); },
                                            std::log(beta) - 3.0, std::log(beta) + 3.0, 40);
            beta = std::exp(lb);
        }
        bin_d[b] = d;
        bin_beta[b] = beta;
    });
    rows.d.resize(n);
    rows.var.resize(n);
    for (size_t i = 0; i < n; i++) {
        // linear in 1/d^2 between the bins' centres, flat beyond the outermost
        const double s = refl.inv_d2[i];
        double d = bin_d[0], beta = bin_beta[0];
        if (s >= centre[bins - 1]) {
            d = bin_d[bins - 1];
            beta = bin_beta[bins - 1];
        } else if (s > centre[0]) {
            const size_t b = std::upper_bound(centre.begin(), centre.end(), s) - centre.begin();
            const double t = (s - centre[b - 1]) / (centre[b] - centre[b - 1]);
            d = bin_d[b - 1] + t * (bin_d[b] - bin_d[b - 1]);
            beta = bin_beta[b - 1] + t * (bin_beta[b] - bin_beta[b - 1]);
        }
        rows.d[i] = static_cast<float>(d);
        rows.var[i] = static_cast<float>(refl.epsilon[i] * beta + double(refl.sigma[i]) * refl.sigma[i]);
    }
}

RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector<std::complex<float>> &fmodel) {
    double nw = 0, dw = 0, nf = 0, df = 0;
    RFactors r;
    for (size_t i = 0; i < refl.Size(); i++) {
        const double diff = std::fabs(refl.fo[i] - std::abs(fmodel[i]));
        if (refl.free[i]) {
            nf += diff;
            df += refl.fo[i];
            r.n_free++;
        } else {
            nw += diff;
            dw += refl.fo[i];
            r.n_work++;
        }
    }
    r.r_work = dw > 0 ? nw / dw : 1;
    r.r_free = df > 0 ? nf / df : 1;
    return r;
}

std::vector<std::array<double, 4>> ModelXyzB(const gemmi::Model &model) {
    std::vector<std::array<double, 4>> out;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms)
                out.push_back({a.pos.x, a.pos.y, a.pos.z, a.b_iso});
    return out;
}

namespace {
std::mutex stage_mutex;
std::map<std::string, std::pair<double, int>> stage_times;
}  // namespace

RefinementStage::RefinementStage(const char *name)
    : name_(name), start_(std::chrono::steady_clock::now()), nvtx_(NvtxRangeStart(name)) {}

RefinementStage::~RefinementStage() {
    NvtxRangeEnd(nvtx_);
    const double seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start_).count();
    std::lock_guard lock(stage_mutex);
    std::pair<double, int> &t = stage_times[name_];
    t.first += seconds;
    t.second++;
}

void RefinementStage::Log(Logger &logger) {
    std::lock_guard lock(stage_mutex);
    for (const auto &[name, t] : stage_times)
        logger.Debug("TIMING refine.{} {:.3f} {}", name, t.first, t.second);
    stage_times.clear();
}

double RefinementBlur(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, double b_floor) {
    // set_refmac_compatible_blur(): u_to_b / 1.1 * spacing^2 - b_min, not below 0, the spacing the one asked for
    const gemmi::DensityCalculator<Table, float> dc = Calculator(cell, sg, d_min, 0);
    const double spacing = dc.requested_grid_spacing();
    return std::max(0.0, gemmi::u_to_b() / 1.1 * spacing * spacing - b_floor);
}

std::unique_ptr<RefinementEngine> MakeRefinementEngineCPU(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                                          double d_min, const RefinementReflections &refl,
                                                          size_t nthreads) {
    return std::make_unique<RefinementEngineCPU>(cell, sg, d_min, refl, nthreads);
}

void AtomGradientsFromMap(const gemmi::Grid<float> &map, const gemmi::Model &model, double blur, size_t nthreads,
                          std::vector<std::array<double, 4>> &grad) {
    const gemmi::DensityCalculator<Table, float> dc = Calculator(map.unit_cell, *map.spacegroup, 1.0, blur);
    std::vector<const gemmi::Atom *> atoms;
    for (const gemmi::Chain &ch : model.chains)
        for (const gemmi::Residue &r : ch.residues)
            for (const gemmi::Atom &a : r.atoms)
                atoms.push_back(&a);
    grad.assign(atoms.size(), {0, 0, 0, 0});
    const int nu = map.nu, nv = map.nv, nw = map.nw;
    // As ModelGrid.cpp walks a density box: the orthogonalization over the grid size in single precision.
    const gemmi::UpperTriangularMat33 &om = map.orth_n;
    const float m00 = static_cast<float>(om.a11), m01 = static_cast<float>(om.a12), m02 = static_cast<float>(om.a13),
                m11 = static_cast<float>(om.a22), m12 = static_cast<float>(om.a23), m22 = static_cast<float>(om.a33);
    constexpr double FOUR_PI2 = 4 * refine_math::PI * refine_math::PI;
    ParallelChunks(static_cast<int>(atoms.size()), nthreads, [&](int lo, int hi) {
        for (int i = lo; i < hi; i++) {
            const gemmi::Atom &atom = *atoms[i];
            const float b = static_cast<float>(atom.b_iso + blur);
            const auto precal = Table::get(atom.element, atom.charge, atom.serial).precalculate_density_iso(b, 0.0f);
            const double radius = dc.estimate_radius(precal, b);
            const float max_r2 = static_cast<float>(radius * radius);
            const gemmi::Fractional fpos = map.unit_cell.fractionalize(atom.pos);
            const double n3[3] = {fpos.x * nu, fpos.y * nv, fpos.z * nw};
            float nctr[3];
            int c[3], d[3];
            for (int k = 0; k < 3; k++) {
                nctr[k] = static_cast<float>(n3[k]);
                c[k] = gemmi::iround(n3[k]);
            }
            d[0] = std::min(static_cast<int>(std::ceil(radius / map.spacing[0])), nu - 1);
            d[1] = std::min(static_cast<int>(std::ceil(radius / map.spacing[1])), nv - 1);
            d[2] = std::min(static_cast<int>(std::ceil(radius / map.spacing[2])), nw - 1);
            double gx = 0, gy = 0, gz = 0, gb = 0;
            for (int w = c[2] - d[2]; w <= c[2] + d[2]; w++) {
                const float fz = nctr[2] - static_cast<float>(w);
                const float z = m22 * fz;
                const int w_ = gemmi::modulo(w, nw);
                for (int v = c[1] - d[1]; v <= c[1] + d[1]; v++) {
                    const float fy = nctr[1] - static_cast<float>(v);
                    const float y = m11 * fy + m12 * fz;
                    const float yz = y * y + z * z;
                    if (yz > max_r2)
                        continue;
                    const float x_yz = m01 * fy + m02 * fz;
                    const size_t row = static_cast<size_t>(nu) * (gemmi::modulo(v, nv) + static_cast<size_t>(nv) * w_);
                    for (int u = c[0] - d[0]; u <= c[0] + d[0]; u++) {
                        const float x = m00 * (nctr[0] - static_cast<float>(u)) + x_yz;
                        const float r2 = yz + x * x;
                        if (r2 > max_r2)
                            continue;
                        const double value = map.data[row + gemmi::modulo(u, nu)];
                        double dr2 = 0, db = 0;
                        for (size_t q = 0; q < std::size(precal.a); q++) {
                            const double e = precal.a[q] * gemmi::unsafe_expapprox(std::max(precal.b[q] * r2, -88.f));
                            dr2 += e * precal.b[q];
                            db += e * precal.b[q] * (1.5 + precal.b[q] * r2);
                        }
                        gx += value * dr2 * x;
                        gy += value * dr2 * y;
                        gz += value * dr2 * z;
                        gb += value * db;
                    }
                }
            }
            // rho = occ sum_q a_q exp(b_q r^2): d rho / d(centre) = occ 2 (sum a_q b_q e_q) (centre - point), and
            // with b_q = -4 pi^2 / (B_q + B + blur), d rho / dB = occ sum a_q e_q b_q (1.5 + b_q r^2) / (4 pi^2).
            grad[i] = {2.0 * atom.occ * gx, 2.0 * atom.occ * gy, 2.0 * atom.occ * gz, atom.occ * gb / FOUR_PI2};
        }
    });
}