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

#pragma once

// The per-term arithmetic of restrained refinement (Refinement.h), written once for the CPU and the GPU:
// the maximum-likelihood amplitude target of one reflection and the geometry restraints, each with its
// derivatives. Plain functions on plain numbers, so the CUDA engine (RefinementGPU.cu) and the CPU path
// (RefinementXray.cpp, RefinementRestraints.cpp) evaluate the same expressions.

#include <cmath>

#if defined(__CUDACC__)
#define REFINE_HD __host__ __device__
#else
#define REFINE_HD
#endif

// What a restraint is, for the evaluators on either side (RefinementRestraints.h builds them).
enum class RestraintKind : int { Bond = 0, Angle, Torsion, Chiral, Plane, Repulsion, BSimilarity };

// One restraint. E = ((value - ideal) / sigma)^2, summed over a set.
struct RestraintRecord {
    int kind = 0;
    int atom[4] = {0, 0, 0, 0};   // Plane: atom[0] the first of its atoms in plane_atoms, atom[1] their number
    float ideal = 0;              // A, degrees, A^3 (chiral) or A (repulsion: the contact distance)
    float sigma = 1;
    int period = 0;               // Torsion
    int slot = 0;                 // the first of its gradient slots, one per atom
    // Repulsion: where the second atom is, as the image x' = rot x + tran of its position (Cartesian);
    // the identity elsewhere.
    float rot[9] = {1, 0, 0, 0, 1, 0, 0, 0, 1};
    float tran[3] = {0, 0, 0};
};

namespace refine_math {

constexpr double PI = 3.14159265358979323846;
constexpr double DEG = PI / 180.0;

// ln I0(x) and I1(x)/I0(x) for x >= 0, from the polynomial approximations of Abramowitz & Stegun (1964),
// Handbook of Mathematical Functions, 9.8.1-9.8.4 (relative error below 2e-7); for large x through the
// exponentially scaled form, so nothing overflows.
template <typename T>
REFINE_HD inline void BesselI0I1(T x, T &log_i0, T &i1_over_i0) {
    if (x < T(3.75)) {
        const T t = (x / T(3.75)) * (x / T(3.75));
        const T i0 = T(1.0) + t * (T(3.5156229) + t * (T(3.0899424) + t * (T(1.2067492) +
                     t * (T(0.2659732) + t * (T(0.0360768) + t * T(0.0045813))))));
        const T i1 = x * (T(0.5) + t * (T(0.87890594) + t * (T(0.51498869) + t * (T(0.15084934) +
                     t * (T(0.02658733) + t * (T(0.00301532) + t * T(0.00032411)))))));
        log_i0 = log(i0);
        i1_over_i0 = i1 / i0;
    } else {
        const T t = T(3.75) / x;
        // sqrt(x) exp(-x) I0(x) and sqrt(x) exp(-x) I1(x)
        const T s0 = T(0.39894228) + t * (T(0.01328592) + t * (T(0.00225319) + t * (T(-0.00157565) +
                     t * (T(0.00916281) + t * (T(-0.02057706) + t * (T(0.02635537) + t * (T(-0.01647633) +
                     t * T(0.00392377))))))));
        const T s1 = T(0.39894228) + t * (T(-0.03988024) + t * (T(-0.00362018) + t * (T(0.00163801) +
                     t * (T(-0.01031555) + t * (T(0.02282967) + t * (T(-0.02895312) + t * (T(0.01787654) +
                     t * T(-0.00420059))))))));
        log_i0 = x + log(s0) - T(0.5) * log(x);
        i1_over_i0 = s1 / s0;
    }
}

// The likelihood of one observed amplitude given the model amplitude Fc - the Rice distribution for an
// acentric reflection, its centric counterpart otherwise - with D the fraction of the model that is
// right and var the variance the rest of the structure and the measurement leave (epsilon beta plus the
// observation's own sigma^2). Returns -ln L without the terms that depend on neither Fc nor the
// parameters, and sets d(-ln L)/dFc.
// Following Pannu & Read (1996) Acta Cryst. A52, 659-668 and Murshudov, Vagin & Dodson (1997) Acta Cryst.
// D53, 240-255
template <typename T>
REFINE_HD inline T AmplitudeNegLogLikelihood(T fo, T fc, T d, T var, bool centric, T &dfc) {
    if (!centric) {
        const T x = T(2.0) * d * fo * fc / var;
        T log_i0, ratio;
        BesselI0I1(x, log_i0, ratio);
        dfc = (T(2.0) * d * d * fc - ratio * T(2.0) * d * fo) / var;
        return log(var) + (fo * fo + d * d * fc * fc) / var - log_i0;
    }
    const T x = d * fo * fc / var;
    // ln cosh x = x + ln(1 + exp(-2x)) - ln 2 for x >= 0
    const T log_cosh = x + log(T(1.0) + exp(T(-2.0) * x)) - T(0.6931471805599453);
    dfc = (d * d * fc - tanh(x) * d * fo) / var;
    return T(0.5) * log(var) + (fo * fo + d * d * fc * fc) / (T(2.0) * var) - log_cosh;
}

// --- geometry restraints: E = ((value - ideal) / sigma)^2 and dE/dx for each atom ---

template <typename T>
struct V3 {
    T x, y, z;
};

template <typename T> REFINE_HD inline V3<T> operator-(V3<T> a, V3<T> b) { return {a.x - b.x, a.y - b.y, a.z - b.z}; }
template <typename T> REFINE_HD inline V3<T> operator+(V3<T> a, V3<T> b) { return {a.x + b.x, a.y + b.y, a.z + b.z}; }
template <typename T> REFINE_HD inline V3<T> operator*(T s, V3<T> a) { return {s * a.x, s * a.y, s * a.z}; }
template <typename T> REFINE_HD inline T Dot(V3<T> a, V3<T> b) { return a.x * b.x + a.y * b.y + a.z * b.z; }
template <typename T> REFINE_HD inline V3<T> Cross(V3<T> a, V3<T> b) {
    return {a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x};
}
template <typename T> REFINE_HD inline T Norm(V3<T> a) { return sqrt(Dot(a, a)); }

// A bond length.
template <typename T>
REFINE_HD inline T BondTerm(V3<T> a, V3<T> b, T ideal, T sigma, V3<T> g[2]) {
    const V3<T> d = a - b;
    const T len = Norm(d);
    const T r = (len - ideal) / sigma;
    const T f = len > T(0) ? T(2.0) * r / (sigma * len) : T(0);
    g[0] = f * d;
    g[1] = (-f) * d;
    return r * r;
}

// A bond angle at b, ideal and sigma in degrees.
template <typename T>
REFINE_HD inline T AngleTerm(V3<T> a, V3<T> b, V3<T> c, T ideal, T sigma, V3<T> g[3]) {
    const V3<T> u = a - b, v = c - b;
    const T lu = Norm(u), lv = Norm(v);
    const V3<T> cr = Cross(u, v);
    const T theta = atan2(Norm(cr), Dot(u, v));
    const T s = T(DEG);
    const T r = (theta - ideal * s) / (sigma * s);
    // dtheta/du = (cos(theta) u/|u| - v/|v|) / (|u| sin(theta)), and the same with u and v swapped
    const T sin_t = fmax(sin(theta), T(1e-6)), cos_t = cos(theta);
    const T f = T(2.0) * r / (sigma * s);
    const V3<T> du = (T(1.0) / (lu * sin_t)) * ((cos_t / lu) * u - (T(1.0) / lv) * v);
    const V3<T> dv = (T(1.0) / (lv * sin_t)) * ((cos_t / lv) * v - (T(1.0) / lu) * u);
    g[0] = f * du;
    g[2] = f * dv;
    g[1] = (-f) * (du + dv);
    return r * r;
}

// The dihedral angle p0-p1-p2-p3 in radians, and its derivatives with respect to the four positions.
// Following Blondel & Karplus (1996) J. Comput. Chem. 17, 1132-1141
template <typename T>
REFINE_HD inline T Dihedral(V3<T> p0, V3<T> p1, V3<T> p2, V3<T> p3, V3<T> d[4]) {
    const V3<T> b1 = p1 - p0, b2 = p2 - p1, b3 = p3 - p2;
    const V3<T> m = Cross(b1, b2), n = Cross(b2, b3);
    const T lb2 = Norm(b2);
    const T mm = fmax(Dot(m, m), T(1e-12)), nn = fmax(Dot(n, n), T(1e-12));
    const T phi = atan2(lb2 * Dot(b1, n), Dot(m, n));
    d[0] = (-lb2 / mm) * m;
    d[3] = (lb2 / nn) * n;
    const T bb = lb2 * lb2;
    const T p = Dot(b1, b2) / bb, q = Dot(b3, b2) / bb;
    d[1] = (-(T(1.0) + p)) * d[0] + q * d[3];
    d[2] = (-T(1.0)) * (d[0] + d[1] + d[3]);
    return phi;
}

// A torsion restrained to ideal (degrees) modulo 360/period; period 0 is read as 1.
template <typename T>
REFINE_HD inline T TorsionTerm(V3<T> p0, V3<T> p1, V3<T> p2, V3<T> p3, T ideal, T sigma, int period, V3<T> g[4]) {
    V3<T> d[4];
    const T phi = Dihedral(p0, p1, p2, p3, d) / T(DEG);
    const T step = T(360.0) / T(period > 0 ? period : 1);
    T delta = phi - ideal;
    delta -= step * floor(delta / step + T(0.5));
    const T r = delta / sigma;
    const T f = T(2.0) * r / (sigma * T(DEG));
    for (int i = 0; i < 4; i++)
        g[i] = f * d[i];
    return r * r;
}

// A chiral volume (c the centre): (a1 - c) . ((a2 - c) x (a3 - c)).
template <typename T>
REFINE_HD inline T ChiralTerm(V3<T> c, V3<T> a1, V3<T> a2, V3<T> a3, T ideal, T sigma, V3<T> g[4]) {
    const V3<T> x1 = a1 - c, x2 = a2 - c, x3 = a3 - c;
    const T vol = Dot(x1, Cross(x2, x3));
    const T r = (vol - ideal) / sigma;
    const T f = T(2.0) * r / sigma;
    g[1] = f * Cross(x2, x3);
    g[2] = f * Cross(x3, x1);
    g[3] = f * Cross(x1, x2);
    g[0] = (-T(1.0)) * (g[1] + g[2] + g[3]);
    return r * r;
}

// A repulsion below the contact distance d0 (none above it).
template <typename T>
REFINE_HD inline T RepulsionTerm(V3<T> a, V3<T> b, T d0, T sigma, V3<T> g[2]) {
    const V3<T> d = a - b;
    const T len = Norm(d);
    if (len >= d0 || len <= T(0)) {
        g[0] = g[1] = V3<T>{0, 0, 0};
        return T(0);
    }
    const T r = (d0 - len) / sigma;
    const T f = T(-2.0) * r / (sigma * len);
    g[0] = f * d;
    g[1] = (-f) * d;
    return r * r;
}

// The unit eigenvector of the smallest eigenvalue of a symmetric 3x3 matrix (a00 a11 a22 a01 a02 a12), by
// cyclic Jacobi rotations.
template <typename T>
REFINE_HD inline V3<T> SmallestEigenvector(const T m_in[6]) {
    T a[3][3] = {{m_in[0], m_in[3], m_in[4]}, {m_in[3], m_in[1], m_in[5]}, {m_in[4], m_in[5], m_in[2]}};
    T v[3][3] = {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}};
    for (int sweep = 0; sweep < 12; sweep++)
        for (int p = 0; p < 2; p++)
            for (int q = p + 1; q < 3; q++) {
                if (fabs(a[p][q]) < T(1e-30))
                    continue;
                const T theta = (a[q][q] - a[p][p]) / (T(2.0) * a[p][q]);
                const T t = (theta >= 0 ? T(1.0) : T(-1.0)) / (fabs(theta) + sqrt(theta * theta + T(1.0)));
                const T c = T(1.0) / sqrt(t * t + T(1.0)), s = t * c;
                for (int k = 0; k < 3; k++) {
                    const T akp = a[k][p], akq = a[k][q];
                    a[k][p] = c * akp - s * akq;
                    a[k][q] = s * akp + c * akq;
                }
                for (int k = 0; k < 3; k++) {
                    const T apk = a[p][k], aqk = a[q][k];
                    a[p][k] = c * apk - s * aqk;
                    a[q][k] = s * apk + c * aqk;
                }
                for (int k = 0; k < 3; k++) {
                    const T vkp = v[k][p], vkq = v[k][q];
                    v[k][p] = c * vkp - s * vkq;
                    v[k][q] = s * vkp + c * vkq;
                }
            }
    int i = 0;
    if (a[1][1] < a[i][i])
        i = 1;
    if (a[2][2] < a[i][i])
        i = 2;
    return {v[0][i], v[1][i], v[2][i]};
}

// A plane: the distances of n atoms from their least-squares plane. The plane is optimal, so the
// derivative of each atom's distance carries no term from the plane moving: dE/dx_i = 2 (n.(x_i - c)) n / sigma^2.
template <typename T>
REFINE_HD inline T PlaneTerm(const V3<T> *x, int n, T sigma, V3<T> *g) {
    V3<T> c{0, 0, 0};
    for (int i = 0; i < n; i++)
        c = c + x[i];
    c = (T(1.0) / T(n)) * c;
    T m[6] = {0, 0, 0, 0, 0, 0};
    for (int i = 0; i < n; i++) {
        const V3<T> d = x[i] - c;
        m[0] += d.x * d.x;
        m[1] += d.y * d.y;
        m[2] += d.z * d.z;
        m[3] += d.x * d.y;
        m[4] += d.x * d.z;
        m[5] += d.y * d.z;
    }
    const V3<T> nv = SmallestEigenvector(m);
    T e = 0;
    for (int i = 0; i < n; i++) {
        const T r = Dot(nv, x[i] - c) / sigma;
        e += r * r;
        g[i] = (T(2.0) * r / sigma) * nv;
    }
    return e;
}

}  // namespace refine_math